okooooooo/gpt-load
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(dialect): 实现 OpenAI 用量抽取

Browse Source
This commit is contained in:
tbphp committed 2026-07-26 09:22:39 +08:00
1 parent 31d0822c54
commit 4c233fa6f9
7 files changed
+617
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

internal/dialect/openai_usage.go
+126
View File
@@ -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
}
internal/dialect/openai_usage_test.go
+316
View File
@@ -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)
}
}
internal/dialect/testdata/usage/openai/nonstream.json
Vendored
+14
View File
@@ -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
}
}
}
internal/dialect/testdata/usage/openai/stream.jsonl
Vendored
+3
View File
@@ -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}}}
internal/dialect/usage.go
+17
View File
@@ -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)
internal/dialect/usage_json.go
+61
View File
@@ -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
}
internal/dialect/usage_test.go
+80
View File
@@ -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
}