diff --git a/internal/dialect/openai_usage.go b/internal/dialect/openai_usage.go new file mode 100644 index 00000000..71d050e1 --- /dev/null +++ b/internal/dialect/openai_usage.go @@ -0,0 +1,126 @@ +package dialect + +import ( + "encoding/json" + "fmt" + + "gpt-load/internal/usage" +) + +func (d *OpenAI) ExtractUsage(body []byte) (usage.Result, error) { + object, err := decodeJSONObject(body) + if err != nil { + return usage.Result{}, fmt.Errorf("decode OpenAI usage response") + } + patch, found := openAIUsagePatch(object, true) + var accumulator usage.Accumulator + if found { + if err := accumulator.ReplaceSnapshot(patch); err != nil { + return usage.Result{}, fmt.Errorf("normalize OpenAI usage response") + } + } + result, _ := accumulator.Finalize(true) + return result, nil +} + +func (d *OpenAI) NewUsageStreamExtractor() UsageStreamExtractor { + return &openAIUsageStreamExtractor{} +} + +type openAIUsageStreamExtractor struct { + accumulator usage.Accumulator + diagnostics usage.Diagnostics +} + +func (e *openAIUsageStreamExtractor) Observe(payload []byte) error { + object, err := decodeJSONObject(payload) + if err != nil { + e.diagnostics.Add(usage.DiagnosticInvalidPayload) + if mergeErr := e.accumulator.MergePatch(usage.Patch{Diagnostics: e.diagnostics}); mergeErr != nil { + return mergeErr + } + return fmt.Errorf("decode OpenAI usage stream payload") + } + + patch, found := openAIUsagePatch(object, openAIUsageFinal(object)) + if !found { + return nil + } + patch.Diagnostics.Merge(e.diagnostics) + if err := e.accumulator.ReplaceSnapshot(patch); err != nil { + return err + } + return nil +} + +func (e *openAIUsageStreamExtractor) Finalize() (usage.Result, bool) { + return e.accumulator.Finalize(true) +} + +func openAIUsagePatch(root map[string]json.RawMessage, final bool) (usage.Patch, bool) { + usageObject, diagnostics := usageOptionalObject(root, "usage") + if usageObject == nil { + return usage.Patch{}, false + } + + prompt, promptDiagnostics := usageInteger(usageObject, "prompt_tokens", true) + diagnostics.Merge(promptDiagnostics) + completion, completionDiagnostics := usageInteger(usageObject, "completion_tokens", true) + diagnostics.Merge(completionDiagnostics) + total, totalDiagnostics := usageInteger(usageObject, "total_tokens", false) + diagnostics.Merge(totalDiagnostics) + + var cached *int64 + if details, detailDiagnostics := usageOptionalObject(usageObject, "prompt_tokens_details"); details != nil { + diagnostics.Merge(detailDiagnostics) + cached, detailDiagnostics = usageInteger(details, "cached_tokens", false) + diagnostics.Merge(detailDiagnostics) + cacheWrite, cacheWriteDiagnostics := usageInteger(details, "cache_write_tokens", false) + diagnostics.Merge(cacheWriteDiagnostics) + if cacheWrite != nil && *cacheWrite > 0 { + diagnostics.Add(usage.DiagnosticUnsupportedBillableDetail) + } + } else { + diagnostics.Merge(detailDiagnostics) + } + + patch := usage.Patch{Final: final, Diagnostics: diagnostics} + if prompt != nil { + cacheValue := int64(0) + if cached != nil { + cacheValue = *cached + } + uncached, ok := usage.SubtractCached(*prompt, cacheValue) + if !ok { + uncached = 0 + patch.Diagnostics.Add(usage.DiagnosticNegativeValue) + } + patch.UncachedInput = &uncached + patch.CacheRead = &cacheValue + } + if completion != nil { + patch.Output = completion + } + + tokens := usage.Tokens{} + if patch.UncachedInput != nil { + tokens.UncachedInput = *patch.UncachedInput + } + if patch.CacheRead != nil { + tokens.CacheRead = *patch.CacheRead + } + if patch.Output != nil { + tokens.Output = *patch.Output + } + usageNormalizedTotal(total, tokens, &patch.Diagnostics) + return patch, true +} + +func openAIUsageFinal(root map[string]json.RawMessage) bool { + rawChoices, exists := root["choices"] + if !exists { + return false + } + var choices []json.RawMessage + return json.Unmarshal(rawChoices, &choices) == nil && choices != nil && len(choices) == 0 +} diff --git a/internal/dialect/openai_usage_test.go b/internal/dialect/openai_usage_test.go new file mode 100644 index 00000000..5260a363 --- /dev/null +++ b/internal/dialect/openai_usage_test.go @@ -0,0 +1,316 @@ +package dialect + +import ( + "net/http" + "strings" + "testing" + + "gpt-load/internal/usage" +) + +func TestUsageOpenAICanonicalFixtures(t *testing.T) { + extractor := NewOpenAI(http.DefaultClient) + result, err := extractor.ExtractUsage(readUsageFixture(t, "openai", "nonstream.json")) + if err != nil { + t.Fatal(err) + } + want := usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30} + if result.State != usage.StateComplete || result.Tokens != want { + t.Fatalf("non-stream result = %#v", result) + } + + stream := extractor.NewUsageStreamExtractor() + observeUsageJSONL(t, stream, readUsageFixture(t, "openai", "stream.jsonl")) + streamResult, ok := stream.Finalize() + if !ok || streamResult.State != usage.StateComplete || streamResult.Tokens != want { + t.Fatalf("stream result = %#v, %t", streamResult, ok) + } + if _, ok := stream.Finalize(); ok { + t.Fatal("second Finalize() succeeded") + } +} + +func TestUsageOpenAINonStreamOptionalFields(t *testing.T) { + extractor := NewOpenAI(http.DefaultClient) + tests := []struct { + name string + body string + want usage.Tokens + diagnostic usage.DiagnosticCode + }{ + { + name: "cached missing", + body: `{"usage":{"prompt_tokens":100,"completion_tokens":30,"total_tokens":130,"prompt_tokens_details":{}}}`, + want: usage.Tokens{UncachedInput: 100, Output: 30}, + }, + { + name: "cached zero and cache write zero", + body: `{"usage":{"prompt_tokens":100,"completion_tokens":30,"total_tokens":130,"prompt_tokens_details":{"cached_tokens":0,"cache_write_tokens":0}}}`, + want: usage.Tokens{UncachedInput: 100, Output: 30}, + }, + { + name: "positive cache write is unrepresentable", + body: `{"usage":{"prompt_tokens":100,"completion_tokens":30,"total_tokens":130,"prompt_tokens_details":{"cached_tokens":20,"cache_write_tokens":5}}}`, + want: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}, + diagnostic: usage.DiagnosticUnsupportedBillableDetail, + }, + { + name: "negative cache write is clamped", + body: `{"usage":{"prompt_tokens":100,"completion_tokens":30,"total_tokens":130,"prompt_tokens_details":{"cached_tokens":20,"cache_write_tokens":-5}}}`, + want: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}, + diagnostic: usage.DiagnosticNegativeValue, + }, + { + name: "invalid cache write is diagnosed", + body: `{"usage":{"prompt_tokens":100,"completion_tokens":30,"total_tokens":130,"prompt_tokens_details":{"cached_tokens":20,"cache_write_tokens":1.5}}}`, + want: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}, + diagnostic: usage.DiagnosticInvalidNumber, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := extractor.ExtractUsage([]byte(tt.body)) + if err != nil { + t.Fatal(err) + } + if result.State != usage.StateComplete || result.Tokens != tt.want { + t.Fatalf("result = %#v, want complete with %#v", result, tt.want) + } + if tt.diagnostic != "" && !result.Diagnostics.Has(tt.diagnostic) { + t.Fatalf("diagnostics = %#v, want %q", result.Diagnostics, tt.diagnostic) + } + }) + } +} + +func TestUsageOpenAINonStreamRequiredAndStrictNumbers(t *testing.T) { + extractor := NewOpenAI(http.DefaultClient) + tests := []struct { + name string + body string + want usage.Tokens + state usage.State + diagnostic usage.DiagnosticCode + }{ + { + name: "prompt missing", + body: `{"usage":{"completion_tokens":30}}`, + want: usage.Tokens{Output: 30}, + state: usage.StateComplete, + diagnostic: usage.DiagnosticMissingRequiredField, + }, + { + name: "prompt null", + body: `{"usage":{"prompt_tokens":null,"completion_tokens":30}}`, + want: usage.Tokens{Output: 30}, + state: usage.StateComplete, + diagnostic: usage.DiagnosticMissingRequiredField, + }, + { + name: "completion missing", + body: `{"usage":{"prompt_tokens":100}}`, + want: usage.Tokens{UncachedInput: 100}, + state: usage.StateComplete, + diagnostic: usage.DiagnosticMissingRequiredField, + }, + { + name: "completion null", + body: `{"usage":{"prompt_tokens":100,"completion_tokens":null}}`, + want: usage.Tokens{UncachedInput: 100}, + state: usage.StateComplete, + diagnostic: usage.DiagnosticMissingRequiredField, + }, + { + name: "negative prompt", + body: `{"usage":{"prompt_tokens":-100,"completion_tokens":30}}`, + want: usage.Tokens{Output: 30}, + state: usage.StateComplete, + diagnostic: usage.DiagnosticNegativeValue, + }, + { + name: "fractional prompt", + body: `{"usage":{"prompt_tokens":1.5,"completion_tokens":30}}`, + want: usage.Tokens{Output: 30}, + state: usage.StateComplete, + diagnostic: usage.DiagnosticInvalidNumber, + }, + { + name: "exponent prompt", + body: `{"usage":{"prompt_tokens":1e2,"completion_tokens":30}}`, + want: usage.Tokens{Output: 30}, + state: usage.StateComplete, + diagnostic: usage.DiagnosticInvalidNumber, + }, + { + name: "string prompt", + body: `{"usage":{"prompt_tokens":"100","completion_tokens":30}}`, + want: usage.Tokens{Output: 30}, + state: usage.StateComplete, + diagnostic: usage.DiagnosticInvalidNumber, + }, + { + name: "overflow prompt", + body: `{"usage":{"prompt_tokens":9223372036854775808,"completion_tokens":30}}`, + want: usage.Tokens{Output: 30}, + state: usage.StateComplete, + diagnostic: usage.DiagnosticInvalidNumber, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := extractor.ExtractUsage([]byte(tt.body)) + if err != nil { + t.Fatal(err) + } + if result.State != tt.state || result.Tokens != tt.want || !result.Diagnostics.Has(tt.diagnostic) { + t.Fatalf("result = %#v, want state %q tokens %#v and diagnostic %q", result, tt.state, tt.want, tt.diagnostic) + } + }) + } +} + +func TestUsageOpenAIClampsCachedAndChecksTotal(t *testing.T) { + extractor := NewOpenAI(http.DefaultClient) + tests := []struct { + name string + body string + want usage.Tokens + wantDelta int64 + code usage.DiagnosticCode + }{ + { + name: "cached exceeds prompt", + body: `{"usage":{"prompt_tokens":100,"completion_tokens":30,"prompt_tokens_details":{"cached_tokens":120}}}`, + want: usage.Tokens{CacheRead: 120, Output: 30}, + code: usage.DiagnosticNegativeValue, + }, + { + name: "total mismatch", + body: `{"usage":{"prompt_tokens":100,"completion_tokens":30,"total_tokens":140,"prompt_tokens_details":{"cached_tokens":20}}}`, + want: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}, + wantDelta: 10, + code: usage.DiagnosticInconsistentTotal, + }, + { + name: "total shortfall", + body: `{"usage":{"prompt_tokens":100,"completion_tokens":30,"total_tokens":120,"prompt_tokens_details":{"cached_tokens":20}}}`, + want: usage.Tokens{UncachedInput: 80, CacheRead: 20, Output: 30}, + wantDelta: -10, + code: usage.DiagnosticInconsistentTotal, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := extractor.ExtractUsage([]byte(tt.body)) + if err != nil { + t.Fatal(err) + } + if result.Tokens != tt.want || !result.Diagnostics.Has(tt.code) { + t.Fatalf("result = %#v, want tokens %#v and diagnostic %q", result, tt.want, tt.code) + } + if tt.code == usage.DiagnosticInconsistentTotal { + delta, ok := result.Diagnostics.TotalDelta() + if !ok || delta != tt.wantDelta { + t.Fatalf("TotalDelta() = %d, %t, want %d, true", delta, ok, tt.wantDelta) + } + } + }) + } +} + +func TestUsageOpenAIStreamSnapshotsAndFinality(t *testing.T) { + extractor := NewOpenAI(http.DefaultClient) + tests := []struct { + name string + steps []string + state usage.State + want usage.Tokens + }{ + { + name: "latest snapshot wins", + steps: []string{ + `{"choices":[{"delta":{}}],"usage":{"prompt_tokens":100,"completion_tokens":10}}`, + `{"choices":[{"delta":{}}],"usage":{"prompt_tokens":100,"completion_tokens":20}}`, + `{"choices":[ ],"usage":{"prompt_tokens":100,"completion_tokens":30}}`, + }, + state: usage.StateComplete, + want: usage.Tokens{UncachedInput: 100, Output: 30}, + }, + { + name: "usage without final marker is partial", + steps: []string{`{"choices":[{"delta":{}}],"usage":{"prompt_tokens":100,"completion_tokens":30}}`}, + state: usage.StatePartial, + want: usage.Tokens{UncachedInput: 100, Output: 30}, + }, + { + name: "no usage is missing", + steps: []string{`{"choices":[{"delta":{"content":"hello"}}]}`}, + state: usage.StateMissing, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + stream := extractor.NewUsageStreamExtractor() + for _, step := range tt.steps { + if err := stream.Observe([]byte(step)); err != nil { + t.Fatal(err) + } + } + result, ok := stream.Finalize() + if !ok || result.State != tt.state || result.Tokens != tt.want { + t.Fatalf("Finalize() = %#v, %t, want state %q tokens %#v", result, ok, tt.state, tt.want) + } + }) + } +} + +func TestUsageOpenAIStreamMalformedPayloadKeepsDiagnosticAndState(t *testing.T) { + stream := NewOpenAI(http.DefaultClient).NewUsageStreamExtractor() + if err := stream.Observe([]byte(`{"choices":[{"delta":{}}],"usage":{"prompt_tokens":100,"completion_tokens":10}}`)); err != nil { + t.Fatal(err) + } + canary := []byte(`{"usage-secret-CANARY":"do-not-leak",`) + if err := stream.Observe(canary); err == nil { + t.Fatal("Observe(malformed) succeeded") + } else if strings.Contains(err.Error(), "usage-secret-CANARY") || strings.Contains(err.Error(), "do-not-leak") { + t.Fatalf("malformed error leaked payload: %q", err) + } + if err := stream.Observe([]byte(`{"choices":[],"usage":{"prompt_tokens":100,"completion_tokens":30}}`)); err != nil { + t.Fatal(err) + } + result, ok := stream.Finalize() + if !ok || result.State != usage.StateComplete || result.Tokens != (usage.Tokens{UncachedInput: 100, Output: 30}) || + !result.Diagnostics.Has(usage.DiagnosticInvalidPayload) { + t.Fatalf("Finalize() = %#v, %t, want latest complete usage with invalid payload diagnostic", result, ok) + } +} + +func TestUsageOpenAIMalformedBodiesAndObserveOwnership(t *testing.T) { + extractor := NewOpenAI(http.DefaultClient) + for _, body := range [][]byte{[]byte(`{}`), []byte(`{"usage":null}`)} { + result, err := extractor.ExtractUsage(body) + if err != nil || result.State != usage.StateMissing || result.Tokens != (usage.Tokens{}) { + t.Fatalf("ExtractUsage(%q) = %#v, %v, want missing zero result without error", body, result, err) + } + } + for _, body := range [][]byte{[]byte(`[]`), []byte(`{"usage":{}} {}`)} { + if _, err := extractor.ExtractUsage(body); err == nil { + t.Fatalf("ExtractUsage(%q) succeeded", body) + } + } + + stream := extractor.NewUsageStreamExtractor() + payload := []byte(`{"choices":[],"usage":{"prompt_tokens":100,"completion_tokens":30}}`) + if err := stream.Observe(payload); err != nil { + t.Fatal(err) + } + copy(payload, []byte(`this caller mutation must not be retained`)) + result, ok := stream.Finalize() + if !ok || result.State != usage.StateComplete || result.Tokens != (usage.Tokens{UncachedInput: 100, Output: 30}) { + t.Fatalf("Finalize() after caller mutation = %#v, %t", result, ok) + } +} diff --git a/internal/dialect/testdata/usage/openai/nonstream.json b/internal/dialect/testdata/usage/openai/nonstream.json new file mode 100644 index 00000000..fc31a82f --- /dev/null +++ b/internal/dialect/testdata/usage/openai/nonstream.json @@ -0,0 +1,14 @@ +{ + "usage": { + "prompt_tokens": 100, + "completion_tokens": 30, + "total_tokens": 130, + "prompt_tokens_details": { + "cached_tokens": 20, + "cache_write_tokens": 0 + }, + "completion_tokens_details": { + "reasoning_tokens": 12 + } + } +} diff --git a/internal/dialect/testdata/usage/openai/stream.jsonl b/internal/dialect/testdata/usage/openai/stream.jsonl new file mode 100644 index 00000000..5acece89 --- /dev/null +++ b/internal/dialect/testdata/usage/openai/stream.jsonl @@ -0,0 +1,3 @@ +{"choices":[{"delta":{"content":"hello"}}]} +{"choices":[{"delta":{}}],"usage":{"prompt_tokens":100,"completion_tokens":30,"total_tokens":130,"prompt_tokens_details":{"cached_tokens":20,"cache_write_tokens":0},"completion_tokens_details":{"reasoning_tokens":12}}} +{"choices":[],"usage":{"prompt_tokens":100,"completion_tokens":30,"total_tokens":130,"prompt_tokens_details":{"cached_tokens":20,"cache_write_tokens":0},"completion_tokens_details":{"reasoning_tokens":12}}} diff --git a/internal/dialect/usage.go b/internal/dialect/usage.go new file mode 100644 index 00000000..3319b692 --- /dev/null +++ b/internal/dialect/usage.go @@ -0,0 +1,17 @@ +package dialect + +import "gpt-load/internal/usage" + +// UsageExtractor optionally exposes provider-specific usage extraction. +type UsageExtractor interface { + ExtractUsage(body []byte) (usage.Result, error) + NewUsageStreamExtractor() UsageStreamExtractor +} + +// UsageStreamExtractor extracts usage from one provider response stream. +type UsageStreamExtractor interface { + Observe(payload []byte) error + Finalize() (usage.Result, bool) +} + +var _ UsageExtractor = (*OpenAI)(nil) diff --git a/internal/dialect/usage_json.go b/internal/dialect/usage_json.go new file mode 100644 index 00000000..117ab655 --- /dev/null +++ b/internal/dialect/usage_json.go @@ -0,0 +1,61 @@ +package dialect + +import ( + "bytes" + "encoding/json" + "strconv" + + "gpt-load/internal/usage" +) + +func usageInteger(object map[string]json.RawMessage, field string, required bool) (*int64, usage.Diagnostics) { + raw, exists := object[field] + if !exists || bytes.Equal(bytes.TrimSpace(raw), []byte("null")) { + if required { + return nil, usageDiagnostic(usage.DiagnosticMissingRequiredField) + } + return nil, usage.Diagnostics{} + } + + value, err := strconv.ParseInt(string(raw), 10, 64) + if err != nil { + return nil, usageDiagnostic(usage.DiagnosticInvalidNumber) + } + if value < 0 { + zero := int64(0) + return &zero, usageDiagnostic(usage.DiagnosticNegativeValue) + } + return &value, usage.Diagnostics{} +} + +func usageOptionalObject(object map[string]json.RawMessage, field string) (map[string]json.RawMessage, usage.Diagnostics) { + raw, exists := object[field] + if !exists || bytes.Equal(bytes.TrimSpace(raw), []byte("null")) { + return nil, usage.Diagnostics{} + } + nested, err := decodeJSONObject(raw) + if err != nil { + return nil, usageDiagnostic(usage.DiagnosticInvalidNumber) + } + return nested, usage.Diagnostics{} +} + +func usageNormalizedTotal(reported *int64, tokens usage.Tokens, diagnostics *usage.Diagnostics) { + if reported == nil { + return + } + normalized, ok := usage.CheckedTotal(tokens) + if !ok { + diagnostics.Add(usage.DiagnosticInvalidNumber) + return + } + if *reported != normalized { + diagnostics.SetTotalDelta(*reported - normalized) + } +} + +func usageDiagnostic(code usage.DiagnosticCode) usage.Diagnostics { + var diagnostics usage.Diagnostics + diagnostics.Add(code) + return diagnostics +} diff --git a/internal/dialect/usage_test.go b/internal/dialect/usage_test.go new file mode 100644 index 00000000..81983c2d --- /dev/null +++ b/internal/dialect/usage_test.go @@ -0,0 +1,80 @@ +package dialect + +import ( + "bufio" + "bytes" + "context" + "net/http" + "os" + "path/filepath" + "testing" + + "gpt-load/internal/health" + "gpt-load/internal/protocol" + "gpt-load/internal/state" +) + +var ( + _ Dialect = (*OpenAI)(nil) + _ Dialect = (*Anthropic)(nil) + _ Dialect = (*Gemini)(nil) + _ Dialect = (*usageDialectOnly)(nil) +) + +func TestUsageOptionalCapability(t *testing.T) { + dialect := Dialect(&usageDialectOnly{}) + if dialect.Protocol() != protocol.OpenAI { + t.Fatalf("Protocol() = %q, want %q", dialect.Protocol(), protocol.OpenAI) + } + if _, ok := dialect.(UsageExtractor); ok { + t.Fatal("Dialect-only implementation unexpectedly has UsageExtractor capability") + } +} + +func readUsageFixture(t *testing.T, provider, name string) []byte { + t.Helper() + body, err := os.ReadFile(filepath.Join("testdata", "usage", provider, name)) + if err != nil { + t.Fatalf("read %s usage fixture %s: %v", provider, name, err) + } + return body +} + +func observeUsageJSONL(t *testing.T, stream UsageStreamExtractor, body []byte) { + t.Helper() + scanner := bufio.NewScanner(bytes.NewReader(body)) + for line := 1; scanner.Scan(); line++ { + payload := bytes.TrimSpace(scanner.Bytes()) + if len(payload) == 0 { + continue + } + if err := stream.Observe(payload); err != nil { + t.Fatalf("Observe(line %d) error = %v", line, err) + } + } + if err := scanner.Err(); err != nil { + t.Fatalf("scan usage fixture: %v", err) + } +} + +type usageDialectOnly struct{} + +func (d *usageDialectOnly) Protocol() protocol.Protocol { return protocol.OpenAI } + +func (d *usageDialectOnly) ExtractModel(*ParsedRequest) (string, bool, error) { return "", false, nil } + +func (d *usageDialectOnly) BuildUpstreamURL(string, *ParsedRequest) (string, error) { return "", nil } + +func (d *usageDialectOnly) InjectCredential(http.Header, string) {} + +func (d *usageDialectOnly) ListModels(context.Context, string, string, state.HeaderRules) ([]string, error) { + return nil, nil +} + +func (d *usageDialectOnly) Probe(context.Context, string, string, state.HeaderRules, string) error { + return nil +} + +func (d *usageDialectOnly) ClassifyStatus(int, []byte) health.FailureCategory { + return health.FailureCategoryOK +}