feat(dialect): 实现 Anthropic 用量抽取
This commit is contained in:
5 files changed
+540
No files matched your search
@@ -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
|
||||
}
|
||||
@@ -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}}
|
||||
@@ -15,3 +15,4 @@ type UsageStreamExtractor interface {
|
||||
}
|
||||
|
||||
var _ UsageExtractor = (*OpenAI)(nil)
|
||||
var _ UsageExtractor = (*Anthropic)(nil)
|
||||
Reference in new issue
Block a user