diff options
Diffstat (limited to '')
| -rw-r--r-- | internal/usage/observer.go | 200 | ||||
| -rw-r--r-- | internal/usage/observer_test.go | 26 |
2 files changed, 226 insertions, 0 deletions
diff --git a/internal/usage/observer.go b/internal/usage/observer.go new file mode 100644 index 0000000..cc8b52f --- /dev/null +++ b/internal/usage/observer.go @@ -0,0 +1,200 @@ +package usage + +import ( + "bytes" + "encoding/json" + "strings" + + "aigw/internal/domain" +) + +const maxCaptureBytes = 64 << 10 + +type Observer struct { + protocol domain.Protocol + stream bool + buffer []byte + line []byte + usage domain.Usage + found bool + explicitTotal bool +} + +func NewObserver(protocol domain.Protocol, stream bool) *Observer { + return &Observer{protocol: protocol, stream: stream} +} + +func (o *Observer) Write(p []byte) (int, error) { + if o.stream { + o.observeSSE(p) + } else { + o.captureTail(p) + } + return len(p), nil +} + +func (o *Observer) Usage() domain.Usage { + if o.stream { + if len(o.line) > 0 { + o.parseSSELine(o.line) + } + return o.usage + } + o.parseJSON(o.buffer) + return o.usage +} + +func (o *Observer) captureTail(p []byte) { + if len(p) >= maxCaptureBytes { + o.buffer = append(o.buffer[:0], p[len(p)-maxCaptureBytes:]...) + return + } + if len(o.buffer)+len(p) > maxCaptureBytes { + drop := len(o.buffer) + len(p) - maxCaptureBytes + copy(o.buffer, o.buffer[drop:]) + o.buffer = o.buffer[:len(o.buffer)-drop] + } + o.buffer = append(o.buffer, p...) +} + +func (o *Observer) observeSSE(p []byte) { + o.line = append(o.line, p...) + for { + index := bytes.IndexByte(o.line, '\n') + if index < 0 { + if len(o.line) > maxCaptureBytes { + o.line = append(o.line[:0], o.line[len(o.line)-maxCaptureBytes:]...) + } + return + } + line := bytes.TrimSpace(o.line[:index]) + o.parseSSELine(line) + o.line = o.line[index+1:] + } +} + +func (o *Observer) parseSSELine(line []byte) { + if !bytes.HasPrefix(line, []byte("data:")) || !bytes.Contains(line, []byte("\"usage\"")) { + return + } + payload := bytes.TrimSpace(bytes.TrimPrefix(line, []byte("data:"))) + if bytes.Equal(payload, []byte("[DONE]")) { + return + } + o.parseJSON(payload) +} + +type usageFields struct { + PromptTokens *int64 `json:"prompt_tokens"` + CompletionTokens *int64 `json:"completion_tokens"` + TotalTokens *int64 `json:"total_tokens"` + InputTokens *int64 `json:"input_tokens"` + OutputTokens *int64 `json:"output_tokens"` + CacheCreationInputTokens *int64 `json:"cache_creation_input_tokens"` + CacheReadInputTokens *int64 `json:"cache_read_input_tokens"` +} + +type responseEnvelope struct { + Usage *usageFields `json:"usage"` + Message *struct { + Usage *usageFields `json:"usage"` + } `json:"message"` +} + +func (o *Observer) parseJSON(payload []byte) { + var envelope responseEnvelope + if err := json.Unmarshal(payload, &envelope); err != nil { + payload = extractUsageObject(payload) + if len(payload) == 0 { + return + } + var fields usageFields + if json.Unmarshal(payload, &fields) == nil { + o.apply(&fields) + } + return + } + if envelope.Usage != nil { + o.apply(envelope.Usage) + } + if envelope.Message != nil && envelope.Message.Usage != nil { + o.apply(envelope.Message.Usage) + } +} + +func (o *Observer) apply(fields *usageFields) { + if fields.PromptTokens != nil { + o.usage.InputTokens = *fields.PromptTokens + o.found = true + } + if fields.InputTokens != nil { + o.usage.InputTokens = *fields.InputTokens + o.found = true + } + if fields.CompletionTokens != nil { + o.usage.OutputTokens = *fields.CompletionTokens + o.found = true + } + if fields.OutputTokens != nil { + o.usage.OutputTokens = *fields.OutputTokens + o.found = true + } + if fields.TotalTokens != nil { + o.usage.TotalTokens = *fields.TotalTokens + o.found = true + o.explicitTotal = true + } + if fields.CacheCreationInputTokens != nil { + o.usage.CacheCreationInputTokens = *fields.CacheCreationInputTokens + o.found = true + } + if fields.CacheReadInputTokens != nil { + o.usage.CacheReadInputTokens = *fields.CacheReadInputTokens + o.found = true + } + if !o.explicitTotal && o.found { + o.usage.TotalTokens = o.usage.InputTokens + o.usage.OutputTokens + } +} + +func extractUsageObject(payload []byte) []byte { + index := strings.LastIndex(string(payload), `"usage"`) + if index < 0 { + return nil + } + rest := payload[index+len(`"usage"`):] + start := bytes.IndexByte(rest, '{') + if start < 0 { + return nil + } + rest = rest[start:] + depth := 0 + inString := false + escaped := false + for i, b := range rest { + if inString { + if escaped { + escaped = false + continue + } + if b == '\\' { + escaped = true + } else if b == '"' { + inString = false + } + continue + } + switch b { + case '"': + inString = true + case '{': + depth++ + case '}': + depth-- + if depth == 0 { + return rest[:i+1] + } + } + } + return nil +} diff --git a/internal/usage/observer_test.go b/internal/usage/observer_test.go new file mode 100644 index 0000000..8fcf408 --- /dev/null +++ b/internal/usage/observer_test.go @@ -0,0 +1,26 @@ +package usage + +import ( + "testing" + + "aigw/internal/domain" +) + +func TestObserverReadsOpenAIJSONUsage(t *testing.T) { + observer := NewObserver(domain.ProtocolOpenAI, false) + _, _ = observer.Write([]byte(`{"choices":[],"usage":{"prompt_tokens":11,"completion_tokens":7,"total_tokens":18}}`)) + got := observer.Usage() + if got.InputTokens != 11 || got.OutputTokens != 7 || got.TotalTokens != 18 { + t.Fatalf("unexpected usage: %+v", got) + } +} + +func TestObserverCombinesAnthropicSSEUsage(t *testing.T) { + observer := NewObserver(domain.ProtocolAnthropic, true) + _, _ = observer.Write([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":12,\"output_tokens\":1}}}\n\n")) + _, _ = observer.Write([]byte("event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":8}}\n\n")) + got := observer.Usage() + if got.InputTokens != 12 || got.OutputTokens != 8 || got.TotalTokens != 20 { + t.Fatalf("unexpected usage: %+v", got) + } +} |