diff --git a/CLAUDE.md b/CLAUDE.md index d97a39c1..a61c50f9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -62,7 +62,7 @@ Client → Echo Middleware (logger → recover → body limit → audit log → - `internal/observability/metrics.go` — Prometheus metrics via hooks injected at factory level: `gomodel_requests_total`, `gomodel_request_duration_seconds`, `gomodel_requests_in_flight`. - `internal/cache/` — Local file or Redis cache backends for model registry. -**Startup:** Config load (defaults → YAML → env vars) → Register providers with factory → Init providers (cache → async model load → background refresh → router) → Init audit logging → Init usage tracking (shares storage if same backend) → Build guardrails pipeline → Create server → Start listening +**Startup:** Config load (defaults → YAML → env vars) → Register providers with factory → Init providers (cache → async model load → background refresh → router) → Register cost mappings (`RegisterCostMappings`) → Init audit logging → Init usage tracking (shares storage if same backend) → Build guardrails pipeline → Create server → Start listening **Shutdown (in order):** HTTP server (stop accepting) → Providers (stop refresh + close cache) → Usage tracking (flush buffer) → Audit logging (flush buffer) @@ -107,6 +107,8 @@ helm/ # Kubernetes Helm charts 1. Create `internal/providers/{name}/` implementing `core.Provider` 2. Export a `Registration` variable: `var Registration = providers.Registration{Type: "{name}", New: New}` + - Optionally add `CostMappings: []core.TokenCostMapping{...}` for provider-specific token cost fields (cached tokens, reasoning tokens, etc.) + - Optionally add `InformationalFields: []string{...}` for known token breakdown fields that don't need separate pricing 3. Register in `cmd/gomodel/main.go` via `factory.Add({name}.Registration)` 4. Add API key env var to `.env.template` and to `knownProviders` in `config/config.go` diff --git a/internal/app/app.go b/internal/app/app.go index f2b65313..41f81def 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -72,6 +72,10 @@ func New(ctx context.Context, cfg Config) (*App, error) { } app.providers = providerResult + // Register provider cost mappings for usage tracking + costMappings, informationalFields := cfg.Factory.CostRegistry() + usage.RegisterCostMappings(costMappings, informationalFields) + // Initialize audit logging auditResult, err := auditlog.New(ctx, appCfg) if err != nil { @@ -212,8 +216,9 @@ func (a *App) Start(addr string) error { // Shutdown gracefully shuts down all components in the correct order. // It ensures proper cleanup of resources: // 1. HTTP server (stop accepting new requests) -// 2. Background refresh goroutine and cache -// 3. Audit logging +// 2. Providers (stop background refresh goroutine and close cache) +// 3. Usage tracking (flush pending entries) +// 4. Audit logging (flush pending logs) // // Safe to call multiple times; subsequent calls are no-ops. func (a *App) Shutdown(ctx context.Context) error { diff --git a/internal/core/cost_mapping.go b/internal/core/cost_mapping.go new file mode 100644 index 00000000..0e35aa82 --- /dev/null +++ b/internal/core/cost_mapping.go @@ -0,0 +1,32 @@ +package core + +// CostSide indicates whether a token cost contributes to input or output. +type CostSide int + +const ( + CostSideUnknown CostSide = iota // zero value; must not be used in mappings + CostSideInput + CostSideOutput +) + +// CostUnit indicates how the pricing field is applied. +type CostUnit int + +const ( + CostUnitUnknown CostUnit = iota // zero value; must not be used in mappings + CostUnitPerMtok // divide token count by 1M, multiply by rate + CostUnitPerItem // multiply count directly by rate +) + +// TokenCostMapping maps a provider-specific RawData key to a pricing field and cost side. +type TokenCostMapping struct { + // RawDataKey is the key in the usage RawData map (e.g. "cached_tokens"). + RawDataKey string + // PricingField returns a pointer to the relevant rate from ModelPricing, or nil + // if the base rate already covers this token type. + PricingField func(p *ModelPricing) *float64 + // Side indicates whether this cost contributes to input or output. + Side CostSide + // Unit indicates the pricing unit (per million tokens or per item). + Unit CostUnit +} diff --git a/internal/providers/anthropic/anthropic.go b/internal/providers/anthropic/anthropic.go index 0bcc85b5..7dbe8149 100644 --- a/internal/providers/anthropic/anthropic.go +++ b/internal/providers/anthropic/anthropic.go @@ -24,6 +24,10 @@ import ( var Registration = providers.Registration{ Type: "anthropic", New: New, + CostMappings: []core.TokenCostMapping{ + {RawDataKey: "cache_read_input_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, Side: core.CostSideInput, Unit: core.CostUnitPerMtok}, + {RawDataKey: "cache_creation_input_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.CacheWritePerMtok }, Side: core.CostSideInput, Unit: core.CostUnitPerMtok}, + }, } const ( @@ -498,11 +502,13 @@ func (sc *streamConverter) convertEvent(event *anthropicStreamEvent) string { } // Include usage data if present (OpenAI format) if event.Usage != nil { - chunk["usage"] = map[string]interface{}{ + usageMap := map[string]interface{}{ "prompt_tokens": event.Usage.InputTokens, "completion_tokens": event.Usage.OutputTokens, "total_tokens": event.Usage.InputTokens + event.Usage.OutputTokens, } + appendCacheFields(usageMap, event.Usage) + chunk["usage"] = usageMap } jsonData, err := json.Marshal(chunk) if err != nil { @@ -668,6 +674,16 @@ func convertAnthropicResponseToResponses(resp *anthropicResponse, model string) } } +// appendCacheFields adds non-zero cache token fields from anthropicUsage to the given map. +func appendCacheFields(m map[string]interface{}, u *anthropicUsage) { + if u.CacheReadInputTokens > 0 { + m["cache_read_input_tokens"] = u.CacheReadInputTokens + } + if u.CacheCreationInputTokens > 0 { + m["cache_creation_input_tokens"] = u.CacheCreationInputTokens + } +} + // buildAnthropicRawUsage extracts cache fields from anthropicUsage into a RawData map. func buildAnthropicRawUsage(u anthropicUsage) map[string]any { raw := make(map[string]any) @@ -784,11 +800,13 @@ func (sc *responsesStreamConverter) Read(p []byte) (n int, err error) { } // Include usage data if captured from message_delta if sc.cachedUsage != nil { - responseData["usage"] = map[string]interface{}{ + usageMap := map[string]interface{}{ "input_tokens": sc.cachedUsage.InputTokens, "output_tokens": sc.cachedUsage.OutputTokens, "total_tokens": sc.cachedUsage.InputTokens + sc.cachedUsage.OutputTokens, } + appendCacheFields(usageMap, sc.cachedUsage) + responseData["usage"] = usageMap } doneEvent := map[string]interface{}{ "type": "response.completed", diff --git a/internal/providers/factory.go b/internal/providers/factory.go index ff78bc60..983b5a9a 100644 --- a/internal/providers/factory.go +++ b/internal/providers/factory.go @@ -3,6 +3,7 @@ package providers import ( "fmt" + "sort" "sync" "gomodel/config" @@ -21,21 +22,23 @@ type ProviderConstructor func(apiKey string, opts ProviderOptions) core.Provider // Registration contains metadata for registering a provider with the factory. type Registration struct { - Type string - New ProviderConstructor + Type string + New ProviderConstructor + CostMappings []core.TokenCostMapping // optional: provider-specific token cost mappings + InformationalFields []string // optional: known breakdown fields that need no separate pricing } // ProviderFactory manages provider registration and creation. type ProviderFactory struct { - mu sync.RWMutex - builders map[string]ProviderConstructor - hooks llmclient.Hooks + mu sync.RWMutex + registrations map[string]Registration + hooks llmclient.Hooks } // NewProviderFactory creates a new provider factory instance. func NewProviderFactory() *ProviderFactory { return &ProviderFactory{ - builders: make(map[string]ProviderConstructor), + registrations: make(map[string]Registration), } } @@ -58,13 +61,13 @@ func (f *ProviderFactory) Add(reg Registration) { } f.mu.Lock() defer f.mu.Unlock() - f.builders[reg.Type] = reg.New + f.registrations[reg.Type] = reg } // Create instantiates a provider based on its resolved configuration. func (f *ProviderFactory) Create(cfg ProviderConfig) (core.Provider, error) { f.mu.RLock() - builder, ok := f.builders[cfg.Type] + reg, ok := f.registrations[cfg.Type] hooks := f.hooks f.mu.RUnlock() @@ -77,7 +80,7 @@ func (f *ProviderFactory) Create(cfg ProviderConfig) (core.Provider, error) { Resilience: cfg.Resilience, } - p := builder(cfg.APIKey, opts) + p := reg.New(cfg.APIKey, opts) if cfg.BaseURL != "" { if setter, ok := p.(interface{ SetBaseURL(string) }); ok { @@ -93,9 +96,34 @@ func (f *ProviderFactory) RegisteredTypes() []string { f.mu.RLock() defer f.mu.RUnlock() - types := make([]string, 0, len(f.builders)) - for t := range f.builders { + types := make([]string, 0, len(f.registrations)) + for t := range f.registrations { types = append(types, t) } return types } + +// CostRegistry returns aggregated cost mappings and informational fields from all +// registered providers. The returned map is keyed by provider type. +func (f *ProviderFactory) CostRegistry() (mappings map[string][]core.TokenCostMapping, informationalFields []string) { + f.mu.RLock() + defer f.mu.RUnlock() + + mappings = make(map[string][]core.TokenCostMapping) + seen := make(map[string]struct{}) + + for _, reg := range f.registrations { + if len(reg.CostMappings) > 0 { + mappings[reg.Type] = reg.CostMappings + } + for _, field := range reg.InformationalFields { + if _, ok := seen[field]; !ok { + seen[field] = struct{}{} + informationalFields = append(informationalFields, field) + } + } + } + + sort.Strings(informationalFields) + return mappings, informationalFields +} diff --git a/internal/providers/factory_test.go b/internal/providers/factory_test.go index 03f9200c..4c0251ce 100644 --- a/internal/providers/factory_test.go +++ b/internal/providers/factory_test.go @@ -3,6 +3,7 @@ package providers import ( "context" "io" + "sort" "testing" "time" @@ -343,3 +344,56 @@ func TestProviderFactory_Create_PassesResilienceConfig(t *testing.T) { t.Errorf("JitterFactor = %f, want 0.5", r.JitterFactor) } } + +func TestProviderFactory_CostRegistry(t *testing.T) { + factory := NewProviderFactory() + + factory.Add(Registration{ + Type: "provider-a", + New: func(_ string, _ ProviderOptions) core.Provider { return &factoryMockProvider{} }, + CostMappings: []core.TokenCostMapping{ + {RawDataKey: "cached_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, Side: core.CostSideInput, Unit: core.CostUnitPerMtok}, + }, + InformationalFields: []string{"prompt_text_tokens"}, + }) + + factory.Add(Registration{ + Type: "provider-b", + New: func(_ string, _ ProviderOptions) core.Provider { return &factoryMockProvider{} }, + CostMappings: []core.TokenCostMapping{ + {RawDataKey: "thought_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, Side: core.CostSideOutput, Unit: core.CostUnitPerMtok}, + }, + InformationalFields: []string{"prompt_text_tokens", "prompt_image_tokens"}, + }) + + // Provider with no cost mappings + factory.Add(Registration{ + Type: "provider-c", + New: func(_ string, _ ProviderOptions) core.Provider { return &factoryMockProvider{} }, + }) + + mappings, informational := factory.CostRegistry() + + // Check mappings + if len(mappings) != 2 { + t.Fatalf("expected 2 provider mappings, got %d", len(mappings)) + } + if len(mappings["provider-a"]) != 1 { + t.Errorf("provider-a: expected 1 mapping, got %d", len(mappings["provider-a"])) + } + if len(mappings["provider-b"]) != 1 { + t.Errorf("provider-b: expected 1 mapping, got %d", len(mappings["provider-b"])) + } + if _, ok := mappings["provider-c"]; ok { + t.Error("provider-c should not have mappings") + } + + // Check informational fields are deduplicated + sort.Strings(informational) + if len(informational) != 2 { + t.Fatalf("expected 2 informational fields (deduplicated), got %d: %v", len(informational), informational) + } + if informational[0] != "prompt_image_tokens" || informational[1] != "prompt_text_tokens" { + t.Errorf("unexpected informational fields: %v", informational) + } +} diff --git a/internal/providers/gemini/gemini.go b/internal/providers/gemini/gemini.go index d47e2394..e7a25c21 100644 --- a/internal/providers/gemini/gemini.go +++ b/internal/providers/gemini/gemini.go @@ -19,6 +19,10 @@ import ( var Registration = providers.Registration{ Type: "gemini", New: New, + CostMappings: []core.TokenCostMapping{ + {RawDataKey: "cached_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, Side: core.CostSideInput, Unit: core.CostUnitPerMtok}, + {RawDataKey: "thought_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, Side: core.CostSideOutput, Unit: core.CostUnitPerMtok}, + }, } const ( diff --git a/internal/providers/openai/openai.go b/internal/providers/openai/openai.go index 2d43d247..043b24bd 100644 --- a/internal/providers/openai/openai.go +++ b/internal/providers/openai/openai.go @@ -16,6 +16,20 @@ import ( var Registration = providers.Registration{ Type: "openai", New: New, + CostMappings: []core.TokenCostMapping{ + {RawDataKey: "cached_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, Side: core.CostSideInput, Unit: core.CostUnitPerMtok}, + {RawDataKey: "prompt_cached_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, Side: core.CostSideInput, Unit: core.CostUnitPerMtok}, + {RawDataKey: "reasoning_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, Side: core.CostSideOutput, Unit: core.CostUnitPerMtok}, + {RawDataKey: "completion_reasoning_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, Side: core.CostSideOutput, Unit: core.CostUnitPerMtok}, + {RawDataKey: "prompt_audio_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.AudioInputPerMtok }, Side: core.CostSideInput, Unit: core.CostUnitPerMtok}, + {RawDataKey: "completion_audio_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.AudioOutputPerMtok }, Side: core.CostSideOutput, Unit: core.CostUnitPerMtok}, + }, + InformationalFields: []string{ + "prompt_text_tokens", + "prompt_image_tokens", + "completion_accepted_prediction_tokens", + "completion_rejected_prediction_tokens", + }, } const ( diff --git a/internal/providers/xai/xai.go b/internal/providers/xai/xai.go index 6b435c95..cdfd7983 100644 --- a/internal/providers/xai/xai.go +++ b/internal/providers/xai/xai.go @@ -15,6 +15,13 @@ import ( var Registration = providers.Registration{ Type: "xai", New: New, + CostMappings: []core.TokenCostMapping{ + {RawDataKey: "cached_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, Side: core.CostSideInput, Unit: core.CostUnitPerMtok}, + {RawDataKey: "prompt_cached_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, Side: core.CostSideInput, Unit: core.CostUnitPerMtok}, + {RawDataKey: "reasoning_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, Side: core.CostSideOutput, Unit: core.CostUnitPerMtok}, + {RawDataKey: "completion_reasoning_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, Side: core.CostSideOutput, Unit: core.CostUnitPerMtok}, + {RawDataKey: "image_tokens", PricingField: func(p *core.ModelPricing) *float64 { return p.InputPerImage }, Side: core.CostSideInput, Unit: core.CostUnitPerItem}, + }, } const ( diff --git a/internal/usage/cost.go b/internal/usage/cost.go index ecb8ab68..e0c375c3 100644 --- a/internal/usage/cost.go +++ b/internal/usage/cost.go @@ -4,6 +4,7 @@ import ( "fmt" "sort" "strings" + "sync/atomic" "gomodel/internal/core" ) @@ -16,79 +17,52 @@ type CostResult struct { Caveat string } -// costSide indicates whether a token cost contributes to input or output. -type costSide int - -const ( - sideInput costSide = iota - sideOutput -) +// costRegistry holds provider-specific cost mappings and informational fields, +// populated at startup via RegisterCostMappings. +type costRegistry struct { + providerMappings map[string][]core.TokenCostMapping + informationalFields map[string]struct{} + extendedFieldSet map[string]struct{} +} -// costUnit indicates how the pricing field is applied. -type costUnit int +// costRegistryPtr is the package-level registry used by CalculateGranularCost +// and stream_wrapper.go. Published atomically by RegisterCostMappings. +var costRegistryPtr atomic.Pointer[costRegistry] -const ( - unitPerMtok costUnit = iota // divide token count by 1M, multiply by rate - unitPerItem // multiply count directly by rate -) - -// tokenCostMapping maps a RawData key to a pricing field and cost side. -type tokenCostMapping struct { - rawDataKey string - pricingField func(p *core.ModelPricing) *float64 - side costSide - unit costUnit +func init() { + costRegistryPtr.Store(&costRegistry{ + providerMappings: make(map[string][]core.TokenCostMapping), + informationalFields: make(map[string]struct{}), + extendedFieldSet: make(map[string]struct{}), + }) } -// providerMappings defines the per-provider RawData key to pricing field mappings. -var providerMappings = map[string][]tokenCostMapping{ - "openai": { - {rawDataKey: "cached_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, side: sideInput, unit: unitPerMtok}, - {rawDataKey: "prompt_cached_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, side: sideInput, unit: unitPerMtok}, - {rawDataKey: "reasoning_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, side: sideOutput, unit: unitPerMtok}, - {rawDataKey: "completion_reasoning_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, side: sideOutput, unit: unitPerMtok}, - {rawDataKey: "prompt_audio_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.AudioInputPerMtok }, side: sideInput, unit: unitPerMtok}, - {rawDataKey: "completion_audio_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.AudioOutputPerMtok }, side: sideOutput, unit: unitPerMtok}, - }, - "anthropic": { - {rawDataKey: "cache_read_input_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, side: sideInput, unit: unitPerMtok}, - {rawDataKey: "cache_creation_input_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.CacheWritePerMtok }, side: sideInput, unit: unitPerMtok}, - }, - "gemini": { - {rawDataKey: "cached_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, side: sideInput, unit: unitPerMtok}, - {rawDataKey: "thought_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, side: sideOutput, unit: unitPerMtok}, - }, - "xai": { - {rawDataKey: "cached_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, side: sideInput, unit: unitPerMtok}, - {rawDataKey: "prompt_cached_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.CachedInputPerMtok }, side: sideInput, unit: unitPerMtok}, - {rawDataKey: "reasoning_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, side: sideOutput, unit: unitPerMtok}, - {rawDataKey: "completion_reasoning_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.ReasoningOutputPerMtok }, side: sideOutput, unit: unitPerMtok}, - {rawDataKey: "image_tokens", pricingField: func(p *core.ModelPricing) *float64 { return p.InputPerImage }, side: sideInput, unit: unitPerItem}, - }, +// loadCostRegistry returns the current cost registry. Never returns nil. +func loadCostRegistry() *costRegistry { + return costRegistryPtr.Load() } -// informationalFields are token fields that are known breakdowns of the base -// input/output counts. They never need separate pricing and should not trigger -// "unmapped token field" caveats. -var informationalFields = map[string]struct{}{ - "prompt_text_tokens": {}, - "prompt_image_tokens": {}, - "completion_accepted_prediction_tokens": {}, - "completion_rejected_prediction_tokens": {}, -} +// RegisterCostMappings populates the cost registry with provider-specific mappings +// and informational fields. Called once at startup after providers are registered. +func RegisterCostMappings(mappings map[string][]core.TokenCostMapping, informational []string) { + reg := &costRegistry{ + providerMappings: mappings, + informationalFields: make(map[string]struct{}, len(informational)), + extendedFieldSet: make(map[string]struct{}), + } -// extendedFieldSet is derived from providerMappings and contains all RawData keys -// that providers may report. Used by stream_wrapper.go to extract extended fields -// from SSE usage data without maintaining a separate hard-coded list. -var extendedFieldSet = func() map[string]struct{} { - set := make(map[string]struct{}) - for _, mappings := range providerMappings { - for _, m := range mappings { - set[m.rawDataKey] = struct{}{} + for _, f := range informational { + reg.informationalFields[f] = struct{}{} + } + + for _, ms := range mappings { + for _, m := range ms { + reg.extendedFieldSet[m.RawDataKey] = struct{}{} } } - return set -}() + + costRegistryPtr.Store(reg) +} // CalculateGranularCost computes input, output, and total costs from token counts, // raw provider-specific data, and pricing information. It accounts for cached tokens, @@ -101,6 +75,8 @@ func CalculateGranularCost(inputTokens, outputTokens int, rawData map[string]any return CostResult{} } + reg := loadCostRegistry() + var inputCost, outputCost float64 var hasInput, hasOutput bool var caveats []string @@ -125,15 +101,15 @@ func CalculateGranularCost(inputTokens, outputTokens int, rawData map[string]any // rawData keys map to the same pricing field (e.g. cached_tokens and prompt_cached_tokens // both map to CachedInputPerMtok). appliedFields := make(map[*float64]bool) - if mappings, ok := providerMappings[providerType]; ok { + if mappings, ok := reg.providerMappings[providerType]; ok { for _, m := range mappings { - count := extractInt(rawData, m.rawDataKey) + count := extractInt(rawData, m.RawDataKey) if count == 0 { continue } - mappedKeys[m.rawDataKey] = true + mappedKeys[m.RawDataKey] = true - rate := m.pricingField(pricing) + rate := m.PricingField(pricing) if rate == nil { continue // Base rate covers this token type; no adjustment needed } @@ -144,20 +120,25 @@ func CalculateGranularCost(inputTokens, outputTokens int, rawData map[string]any appliedFields[rate] = true var cost float64 - switch m.unit { - case unitPerMtok: + switch m.Unit { + case core.CostUnitPerMtok: cost = float64(count) * *rate / 1_000_000 - case unitPerItem: + case core.CostUnitPerItem: cost = float64(count) * *rate + default: + caveats = append(caveats, fmt.Sprintf("unknown cost unit %d for field %s", int(m.Unit), m.RawDataKey)) + continue } - switch m.side { - case sideInput: + switch m.Side { + case core.CostSideInput: inputCost += cost hasInput = true - case sideOutput: + case core.CostSideOutput: outputCost += cost hasOutput = true + default: + caveats = append(caveats, fmt.Sprintf("unknown cost side %d for field %s", int(m.Side), m.RawDataKey)) } } } @@ -167,7 +148,7 @@ func CalculateGranularCost(inputTokens, outputTokens int, rawData map[string]any if mappedKeys[key] { continue } - if _, ok := informationalFields[key]; ok { + if _, ok := reg.informationalFields[key]; ok { continue // Known breakdown of base counts, not separately priced } if isTokenField(key) { diff --git a/internal/usage/setup_test.go b/internal/usage/setup_test.go new file mode 100644 index 00000000..55d7d54d --- /dev/null +++ b/internal/usage/setup_test.go @@ -0,0 +1,27 @@ +package usage + +import ( + "os" + "testing" + + "gomodel/internal/providers" + "gomodel/internal/providers/anthropic" + "gomodel/internal/providers/gemini" + "gomodel/internal/providers/openai" + "gomodel/internal/providers/xai" +) + +func TestMain(m *testing.M) { + // Build cost mappings and informational fields the same way production does: + // register all providers into a factory and use CostRegistry to aggregate. + factory := providers.NewProviderFactory() + factory.Add(openai.Registration) + factory.Add(anthropic.Registration) + factory.Add(gemini.Registration) + factory.Add(xai.Registration) + + costMappings, informationalFields := factory.CostRegistry() + RegisterCostMappings(costMappings, informationalFields) + + os.Exit(m.Run()) +} diff --git a/internal/usage/stream_wrapper.go b/internal/usage/stream_wrapper.go index 8d5fb7fa..a732cf88 100644 --- a/internal/usage/stream_wrapper.go +++ b/internal/usage/stream_wrapper.go @@ -233,7 +233,7 @@ func (w *StreamUsageWrapper) extractUsageFromJSON(data []byte) *UsageEntry { // Extract extended usage data (provider-specific) using the field set // derived from providerMappings in cost.go (single source of truth). - for field := range extendedFieldSet { + for field := range loadCostRegistry().extendedFieldSet { if v, ok := usageMap[field].(float64); ok && v > 0 { rawData[field] = int(v) }