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

This commit is contained in:
tbphp committed 2026-07-26 09:46:36 +08:00
1 parent b8d448fb48
commit 376385d0df
5 files changed
+540

No files matched your search

+190
View File
@@ -0,0 +1,190 @@
package dialect
import (
"encoding/json"
"fmt"
"gpt-load/internal/usage"
)
func (d *Anthropic) ExtractUsage(body []byte) (usage.Result, error) {
root, err := decodeJSONObject(body)
if err != nil {
return usage.Result{}, fmt.Errorf("decode Anthropic usage response")
}
usageObject, diagnostics := usageOptionalObject(root, "usage")
var accumulator usage.Accumulator
if usageObject == nil {
if err := accumulator.MergePatch(usage.Patch{Diagnostics: diagnostics}); err != nil {
return usage.Result{}, fmt.Errorf("normalize Anthropic usage response")
}
} else {
patch, _ := anthropicUsagePatch(usageObject, true, true)
patch.Diagnostics.Merge(diagnostics)
if err := accumulator.MergePatch(patch); err != nil {
return usage.Result{}, fmt.Errorf("normalize Anthropic usage response")
}
}
result, _ := accumulator.Finalize(true)
return result, nil
}
func (d *Anthropic) NewUsageStreamExtractor() UsageStreamExtractor {
return &anthropicUsageStreamExtractor{}
}
type anthropicUsageStreamExtractor struct {
accumulator usage.Accumulator
validStart bool
}
func (e *anthropicUsageStreamExtractor) Observe(payload []byte) error {
root, err := decodeJSONObject(payload)
if err != nil {
if mergeErr := e.mergeDiagnostics(usageDiagnostic(usage.DiagnosticInvalidPayload)); mergeErr != nil {
return mergeErr
}
return fmt.Errorf("decode Anthropic usage stream payload")
}
switch anthropicEventType(root) {
case "message_start":
return e.observeStart(root)
case "message_delta":
return e.observeDelta(root)
default:
return nil
}
}
func (e *anthropicUsageStreamExtractor) Finalize() (usage.Result, bool) {
return e.accumulator.Finalize(true)
}
func (e *anthropicUsageStreamExtractor) observeStart(root map[string]json.RawMessage) error {
message, diagnostics := usageOptionalObject(root, "message")
if message == nil {
if !diagnostics.Has(usage.DiagnosticInvalidNumber) {
diagnostics.Add(usage.DiagnosticMissingRequiredField)
}
return e.mergeDiagnostics(diagnostics)
}
usageObject, usageDiagnostics := usageOptionalObject(message, "usage")
diagnostics.Merge(usageDiagnostics)
if usageObject == nil {
if !diagnostics.Has(usage.DiagnosticInvalidNumber) {
diagnostics.Add(usage.DiagnosticMissingRequiredField)
}
return e.mergeDiagnostics(diagnostics)
}
patch, validInput := anthropicUsagePatch(usageObject, false, false)
patch.Diagnostics.Merge(diagnostics)
if err := e.accumulator.MergePatch(patch); err != nil {
return err
}
if validInput {
e.validStart = true
}
return nil
}
func (e *anthropicUsageStreamExtractor) observeDelta(root map[string]json.RawMessage) error {
usageObject, diagnostics := usageOptionalObject(root, "usage")
if usageObject == nil {
if !diagnostics.Has(usage.DiagnosticInvalidNumber) {
diagnostics.Add(usage.DiagnosticMissingRequiredField)
}
return e.mergeDiagnostics(diagnostics)
}
output, outputDiagnostics := usageInteger(usageObject, "output_tokens", true)
diagnostics.Merge(outputDiagnostics)
validOutput := output != nil && !outputDiagnostics.Has(usage.DiagnosticInvalidNumber) && !outputDiagnostics.Has(usage.DiagnosticNegativeValue)
patch := usage.Patch{Output: output, Diagnostics: diagnostics}
if validOutput {
if e.validStart {
patch.Final = true
} else {
patch.Diagnostics.Add(usage.DiagnosticInvalidEventSequence)
}
}
return e.accumulator.MergePatch(patch)
}
func (e *anthropicUsageStreamExtractor) mergeDiagnostics(diagnostics usage.Diagnostics) error {
return e.accumulator.MergePatch(usage.Patch{Diagnostics: diagnostics})
}
func anthropicUsagePatch(usageObject map[string]json.RawMessage, includeOutput, final bool) (usage.Patch, bool) {
input, diagnostics := usageInteger(usageObject, "input_tokens", true)
validInput := input != nil && !diagnostics.Has(usage.DiagnosticInvalidNumber) && !diagnostics.Has(usage.DiagnosticNegativeValue)
patch := usage.Patch{Final: final, Diagnostics: diagnostics}
if input != nil {
patch.UncachedInput = input
}
cacheRead, cacheReadDiagnostics := usageInteger(usageObject, "cache_read_input_tokens", false)
patch.Diagnostics.Merge(cacheReadDiagnostics)
patch.CacheRead = usageValueOrZero(cacheRead)
aggregate, aggregateDiagnostics := usageInteger(usageObject, "cache_creation_input_tokens", false)
patch.Diagnostics.Merge(aggregateDiagnostics)
aggregateValid := aggregate != nil && !aggregateDiagnostics.Has(usage.DiagnosticInvalidNumber) && !aggregateDiagnostics.Has(usage.DiagnosticNegativeValue)
detail, detailDiagnostics := usageOptionalObject(usageObject, "cache_creation")
patch.Diagnostics.Merge(detailDiagnostics)
if detail == nil {
patch.CacheWrite5M = usageValueOrZero(aggregate)
patch.CacheWrite1H = usageValueOrZero(nil)
if aggregateValid {
patch.Diagnostics.Add(usage.DiagnosticCacheWriteDefaulted5M)
}
} else {
write5M, write5MDiagnostics := usageInteger(detail, "ephemeral_5m_input_tokens", false)
patch.Diagnostics.Merge(write5MDiagnostics)
write1H, write1HDiagnostics := usageInteger(detail, "ephemeral_1h_input_tokens", false)
patch.Diagnostics.Merge(write1HDiagnostics)
patch.CacheWrite5M = usageValueOrZero(write5M)
patch.CacheWrite1H = usageValueOrZero(write1H)
if aggregateValid {
detailTotal, ok := usage.CheckedAdd(*patch.CacheWrite5M, *patch.CacheWrite1H)
if !ok {
patch.Diagnostics.Add(usage.DiagnosticInvalidNumber)
} else if *aggregate != detailTotal {
patch.Diagnostics.SetTotalDelta(*aggregate - detailTotal)
}
}
}
if includeOutput {
output, outputDiagnostics := usageInteger(usageObject, "output_tokens", true)
patch.Diagnostics.Merge(outputDiagnostics)
if output != nil {
patch.Output = output
}
}
return patch, validInput
}
func anthropicEventType(root map[string]json.RawMessage) string {
raw, exists := root["type"]
if !exists {
return ""
}
var eventType string
if err := json.Unmarshal(raw, &eventType); err != nil {
return ""
}
return eventType
}
func usageValueOrZero(value *int64) *int64 {
zero := int64(0)
if value != nil {
zero = *value
}
return &zero
}
+332
View File
@@ -0,0 +1,332 @@
package dialect
import (
"math"
"net/http"
"strings"
"testing"
"gpt-load/internal/usage"
)
func TestUsageAnthropicCanonicalFixtures(t *testing.T) {
extractor, ok := any(NewAnthropic(http.DefaultClient)).(UsageExtractor)
if !ok {
t.Fatal("Anthropic does not expose UsageExtractor capability")
}
want := usage.Tokens{
UncachedInput: 80,
CacheRead: 20,
CacheWrite5M: 5,
CacheWrite1H: 7,
Output: 30,
}
result, err := extractor.ExtractUsage(readUsageFixture(t, "anthropic", "nonstream.json"))
if err != nil {
t.Fatal(err)
}
if result.State != usage.StateComplete || result.Tokens != want {
t.Fatalf("non-stream result = %#v, want complete with %#v", result, want)
}
stream := extractor.NewUsageStreamExtractor()
observeUsageJSONL(t, stream, readUsageFixture(t, "anthropic", "stream.jsonl"))
streamResult, finalized := stream.Finalize()
if !finalized || streamResult.State != usage.StateComplete || streamResult.Tokens != want {
t.Fatalf("stream result = %#v, %t, want complete with %#v", streamResult, finalized, want)
}
if _, finalized := stream.Finalize(); finalized {
t.Fatal("second Finalize() succeeded")
}
}
func TestUsageAnthropicCacheCreationMapping(t *testing.T) {
extractor := NewAnthropic(http.DefaultClient)
tests := []struct {
name string
body string
want usage.Tokens
diagnostics []usage.DiagnosticCode
delta *int64
}{
{
name: "detail wins when aggregate matches",
body: `{"usage":{"input_tokens":80,"output_tokens":30,"cache_read_input_tokens":20,"cache_creation":{"ephemeral_5m_input_tokens":5,"ephemeral_1h_input_tokens":7},"cache_creation_input_tokens":12}}`,
want: usage.Tokens{UncachedInput: 80, CacheRead: 20, CacheWrite5M: 5, CacheWrite1H: 7, Output: 30},
},
{
name: "aggregate falls back to five minutes",
body: `{"usage":{"input_tokens":80,"output_tokens":30,"cache_creation_input_tokens":12}}`,
want: usage.Tokens{UncachedInput: 80, CacheWrite5M: 12, Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticCacheWriteDefaulted5M},
},
{
name: "aggregate surplus records positive delta",
body: `{"usage":{"input_tokens":80,"output_tokens":30,"cache_creation":{"ephemeral_5m_input_tokens":5,"ephemeral_1h_input_tokens":7},"cache_creation_input_tokens":15}}`,
want: usage.Tokens{UncachedInput: 80, CacheWrite5M: 5, CacheWrite1H: 7, Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticInconsistentTotal},
delta: usageInt64(3),
},
{
name: "aggregate shortfall records negative delta",
body: `{"usage":{"input_tokens":80,"output_tokens":30,"cache_creation":{"ephemeral_5m_input_tokens":5,"ephemeral_1h_input_tokens":7},"cache_creation_input_tokens":10}}`,
want: usage.Tokens{UncachedInput: 80, CacheWrite5M: 5, CacheWrite1H: 7, Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticInconsistentTotal},
delta: usageInt64(-2),
},
{
name: "detail sum overflow is diagnosed without wrapping",
body: `{"usage":{"input_tokens":80,"output_tokens":30,"cache_creation":{"ephemeral_5m_input_tokens":9223372036854775807,"ephemeral_1h_input_tokens":1},"cache_creation_input_tokens":9223372036854775807}}`,
want: usage.Tokens{UncachedInput: 80, CacheWrite5M: math.MaxInt64, CacheWrite1H: 1, Output: 30},
diagnostics: []usage.DiagnosticCode{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)
}
requireUsageDiagnostics(t, result.Diagnostics, tt.diagnostics...)
if tt.delta == nil {
if _, ok := result.Diagnostics.TotalDelta(); ok {
t.Fatalf("TotalDelta() unexpectedly present: %#v", result.Diagnostics)
}
} else if delta, ok := result.Diagnostics.TotalDelta(); !ok || delta != *tt.delta {
t.Fatalf("TotalDelta() = %d, %t, want %d, true", delta, ok, *tt.delta)
}
})
}
}
func TestUsageAnthropicStrictNumbersAndRequiredFields(t *testing.T) {
extractor := NewAnthropic(http.DefaultClient)
tests := []struct {
name string
body string
want usage.Tokens
diagnostics []usage.DiagnosticCode
}{
{
name: "all required and optional values may be zero",
body: `{"usage":{"input_tokens":0,"output_tokens":0,"cache_read_input_tokens":0,"cache_creation":{"ephemeral_5m_input_tokens":0,"ephemeral_1h_input_tokens":0},"cache_creation_input_tokens":0}}`,
},
{
name: "input missing preserves output",
body: `{"usage":{"output_tokens":30}}`,
want: usage.Tokens{Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticMissingRequiredField},
},
{
name: "output null preserves input",
body: `{"usage":{"input_tokens":80,"output_tokens":null}}`,
want: usage.Tokens{UncachedInput: 80},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticMissingRequiredField},
},
{
name: "negative required value is safely zeroed",
body: `{"usage":{"input_tokens":-80,"output_tokens":30}}`,
want: usage.Tokens{Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticNegativeValue},
},
{
name: "optional missing and null values are zero",
body: `{"usage":{"input_tokens":80,"output_tokens":30,"cache_read_input_tokens":null,"cache_creation_input_tokens":null}}`,
want: usage.Tokens{UncachedInput: 80, Output: 30},
},
{
name: "negative optional aggregate is not a defaulted aggregate",
body: `{"usage":{"input_tokens":80,"output_tokens":30,"cache_creation_input_tokens":-1}}`,
want: usage.Tokens{UncachedInput: 80, Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticNegativeValue},
},
{
name: "fractional required value is rejected",
body: `{"usage":{"input_tokens":1.5,"output_tokens":30}}`,
want: usage.Tokens{Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticInvalidNumber},
},
{
name: "exponent required value is rejected",
body: `{"usage":{"input_tokens":1e2,"output_tokens":30}}`,
want: usage.Tokens{Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticInvalidNumber},
},
{
name: "string optional value is rejected",
body: `{"usage":{"input_tokens":80,"output_tokens":30,"cache_read_input_tokens":"20"}}`,
want: usage.Tokens{UncachedInput: 80, Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticInvalidNumber},
},
{
name: "fractional optional aggregate is rejected",
body: `{"usage":{"input_tokens":80,"output_tokens":30,"cache_creation_input_tokens":1.5}}`,
want: usage.Tokens{UncachedInput: 80, Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticInvalidNumber},
},
{
name: "overflow required value is rejected",
body: `{"usage":{"input_tokens":9223372036854775808,"output_tokens":30}}`,
want: usage.Tokens{Output: 30},
diagnostics: []usage.DiagnosticCode{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)
}
requireUsageDiagnostics(t, result.Diagnostics, tt.diagnostics...)
})
}
}
func TestUsageAnthropicStreamPatchState(t *testing.T) {
extractor := NewAnthropic(http.DefaultClient)
tests := []struct {
name string
steps []string
state usage.State
want usage.Tokens
diagnostics []usage.DiagnosticCode
}{
{
name: "start only remains partial",
steps: []string{`{"type":"message_start","message":{"usage":{"input_tokens":80,"cache_read_input_tokens":20,"cache_creation":{"ephemeral_5m_input_tokens":5,"ephemeral_1h_input_tokens":7}}}}`},
state: usage.StatePartial,
want: usage.Tokens{UncachedInput: 80, CacheRead: 20, CacheWrite5M: 5, CacheWrite1H: 7},
},
{
name: "delta only preserves output but remains partial",
steps: []string{`{"type":"message_delta","usage":{"output_tokens":30}}`},
state: usage.StatePartial,
want: usage.Tokens{Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticInvalidEventSequence},
},
{
name: "start then missing output remains partial",
steps: []string{
`{"type":"message_start","message":{"usage":{"input_tokens":80}}}`,
`{"type":"message_delta","usage":{}}`,
},
state: usage.StatePartial,
want: usage.Tokens{UncachedInput: 80},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticMissingRequiredField},
},
{
name: "start then invalid output remains partial",
steps: []string{
`{"type":"message_start","message":{"usage":{"input_tokens":80}}}`,
`{"type":"message_delta","usage":{"output_tokens":1.5}}`,
},
state: usage.StatePartial,
want: usage.Tokens{UncachedInput: 80},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticInvalidNumber},
},
{
name: "repeated delta overwrites output",
steps: []string{
`{"type":"message_start","message":{"usage":{"input_tokens":80}}}`,
`{"type":"message_delta","usage":{"output_tokens":10}}`,
`{"type":"message_delta","usage":{"output_tokens":30}}`,
},
state: usage.StateComplete,
want: usage.Tokens{UncachedInput: 80, Output: 30},
},
{
name: "delta before start cannot become complete retroactively",
steps: []string{
`{"type":"message_delta","usage":{"output_tokens":30}}`,
`{"type":"message_start","message":{"usage":{"input_tokens":80}}}`,
},
state: usage.StatePartial,
want: usage.Tokens{UncachedInput: 80, Output: 30},
diagnostics: []usage.DiagnosticCode{usage.DiagnosticInvalidEventSequence},
},
}
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, finalized := stream.Finalize()
if !finalized || result.State != tt.state || result.Tokens != tt.want {
t.Fatalf("Finalize() = %#v, %t, want state %q tokens %#v", result, finalized, tt.state, tt.want)
}
requireUsageDiagnostics(t, result.Diagnostics, tt.diagnostics...)
})
}
}
func TestUsageAnthropicMalformedPayloadPreservesStateAndDoesNotLeak(t *testing.T) {
stream := NewAnthropic(http.DefaultClient).NewUsageStreamExtractor()
if err := stream.Observe([]byte(`{"type":"message_start","message":{"usage":{"input_tokens":80}}}`)); 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(`{"type":"message_delta","usage":{"output_tokens":30}}`)); err != nil {
t.Fatal(err)
}
result, finalized := stream.Finalize()
if !finalized || result.State != usage.StateComplete || result.Tokens != (usage.Tokens{UncachedInput: 80, Output: 30}) {
t.Fatalf("Finalize() = %#v, %t", result, finalized)
}
requireUsageDiagnostics(t, result.Diagnostics, usage.DiagnosticInvalidPayload)
}
func TestUsageAnthropicInvalidUsageObjectIsDiagnosed(t *testing.T) {
extractor := NewAnthropic(http.DefaultClient)
result, err := extractor.ExtractUsage([]byte(`{"usage":[]}`))
if err != nil || result.State != usage.StateMissing || result.Tokens != (usage.Tokens{}) {
t.Fatalf("invalid non-stream usage = %#v, %v", result, err)
}
requireUsageDiagnostics(t, result.Diagnostics, usage.DiagnosticInvalidNumber)
stream := extractor.NewUsageStreamExtractor()
if err := stream.Observe([]byte(`{"type":"message_start","message":{"usage":[]}}`)); err != nil {
t.Fatal(err)
}
result, finalized := stream.Finalize()
if !finalized || result.State != usage.StateMissing || result.Tokens != (usage.Tokens{}) {
t.Fatalf("invalid stream usage = %#v, %t", result, finalized)
}
requireUsageDiagnostics(t, result.Diagnostics, usage.DiagnosticInvalidNumber)
}
func TestUsageAnthropicMalformedBodiesReturnSanitizedErrors(t *testing.T) {
extractor := NewAnthropic(http.DefaultClient)
for _, body := range [][]byte{
[]byte(`[]`),
[]byte(`{"usage":{"input_tokens":80,"output_tokens":30}} {}`),
[]byte(`{"usage-secret-CANARY":"do-not-leak",`),
} {
if _, err := extractor.ExtractUsage(body); err == nil {
t.Fatalf("ExtractUsage(%q) succeeded", body)
} else if strings.Contains(err.Error(), "usage-secret-CANARY") || strings.Contains(err.Error(), "do-not-leak") {
t.Fatalf("malformed error leaked payload: %q", err)
}
}
}
func usageInt64(value int64) *int64 {
return &value
}
@@ -0,0 +1,13 @@
{
"type": "message",
"usage": {
"input_tokens": 80,
"output_tokens": 30,
"cache_read_input_tokens": 20,
"cache_creation": {
"ephemeral_5m_input_tokens": 5,
"ephemeral_1h_input_tokens": 7
},
"cache_creation_input_tokens": 12
}
}
@@ -0,0 +1,4 @@
{"type":"message_start","message":{"id":"msg_usage_fixture","usage":{"input_tokens":80,"cache_read_input_tokens":20,"cache_creation":{"ephemeral_5m_input_tokens":5,"ephemeral_1h_input_tokens":7},"cache_creation_input_tokens":12}}}
{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}
{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":30}}
+1
View File
@@ -15,3 +15,4 @@ type UsageStreamExtractor interface {
}
var _ UsageExtractor = (*OpenAI)(nil)
var _ UsageExtractor = (*Anthropic)(nil)