diff --git a/core/schemas/images.go b/core/schemas/images.go index fba3c2c08a5..07347df4b20 100644 --- a/core/schemas/images.go +++ b/core/schemas/images.go @@ -213,6 +213,25 @@ type ImageTokenDetails struct { TextTokens int `json:"text_tokens,omitempty"` } +// DeepCopy returns an independent copy of u with no shared pointer fields, +// safe for callers (e.g. cost calculation) that need to derive values +// without mutating the original response. Returns nil for a nil receiver. +func (u *ImageUsage) DeepCopy() *ImageUsage { + if u == nil { + return nil + } + out := *u + if u.InputTokensDetails != nil { + details := *u.InputTokensDetails + out.InputTokensDetails = &details + } + if u.OutputTokensDetails != nil { + details := *u.OutputTokensDetails + out.OutputTokensDetails = &details + } + return &out +} + // Streaming Response type BifrostImageGenerationStreamResponse struct { ID string `json:"id,omitempty"` diff --git a/framework/modelcatalog/capabilities_test.go b/framework/modelcatalog/capabilities_test.go deleted file mode 100644 index 3d188f775dd..00000000000 --- a/framework/modelcatalog/capabilities_test.go +++ /dev/null @@ -1,223 +0,0 @@ -package modelcatalog - -import ( - "testing" - - "github.com/maximhq/bifrost/core/schemas" - configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" -) - -func TestGetModelCapabilityEntryForModel_PrefersChatThenResponsesThenCompletion(t *testing.T) { - contextLengthChat := 128000 - maxInputTokensChat := 64000 - maxOutputTokensChat := 16000 - modality := "text" - - mc := &ModelCatalog{ - pricingData: map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "responses"): { - Model: "gpt-4o", - Provider: "openai", - Mode: "responses", - ContextLength: capabilityIntPtr(200000), - MaxInputTokens: capabilityIntPtr(100000), - MaxOutputTokens: capabilityIntPtr(32000), - }, - makeKey("gpt-4o", "openai", "chat"): { - Model: "gpt-4o", - Provider: "openai", - Mode: "chat", - ContextLength: &contextLengthChat, - MaxInputTokens: &maxInputTokensChat, - MaxOutputTokens: &maxOutputTokensChat, - Architecture: &schemas.Architecture{ - Modality: &modality, - }, - }, - }, - } - - entry := mc.GetModelCapabilityEntryForModel("gpt-4o", schemas.OpenAI) - if entry == nil { - t.Fatal("expected capability entry") - } - if entry.Mode != "chat" { - t.Fatalf("expected chat mode to win, got %q", entry.Mode) - } - if entry.ContextLength == nil || *entry.ContextLength != contextLengthChat { - t.Fatalf("expected context_length=%d, got %#v", contextLengthChat, entry.ContextLength) - } - if entry.MaxInputTokens == nil || *entry.MaxInputTokens != maxInputTokensChat { - t.Fatalf("expected max_input_tokens=%d, got %#v", maxInputTokensChat, entry.MaxInputTokens) - } - if entry.MaxOutputTokens == nil || *entry.MaxOutputTokens != maxOutputTokensChat { - t.Fatalf("expected max_output_tokens=%d, got %#v", maxOutputTokensChat, entry.MaxOutputTokens) - } - if entry.Architecture == nil || entry.Architecture.Modality == nil || *entry.Architecture.Modality != modality { - t.Fatalf("expected architecture modality=%q, got %#v", modality, entry.Architecture) - } -} - -func TestGetModelCapabilityEntryForModel_FallsBackToAnyModeDeterministically(t *testing.T) { - mc := &ModelCatalog{ - pricingData: map[string]configstoreTables.TableModelPricing{ - makeKey("imagen", "vertex", "image_generation"): { - Model: "imagen", - Provider: "vertex", - Mode: "image_generation", - ContextLength: capabilityIntPtr(4096), - MaxOutputTokens: capabilityIntPtr(1), - }, - }, - } - - entry := mc.GetModelCapabilityEntryForModel("imagen", schemas.Vertex) - if entry == nil { - t.Fatal("expected capability entry") - } - if entry.Mode != "image_generation" { - t.Fatalf("expected image_generation fallback, got %q", entry.Mode) - } -} - -func TestGetModelCapabilityEntryForModel_ResolvesAliasFamilyViaBaseModel(t *testing.T) { - contextLengthChat := 128000 - - mc := &ModelCatalog{ - pricingData: map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o-2024-08-06", "openai", "responses"): { - Model: "gpt-4o-2024-08-06", - BaseModel: "gpt-4o", - Provider: "openai", - Mode: "responses", - ContextLength: capabilityIntPtr(64000), - MaxOutputTokens: capabilityIntPtr(8000), - }, - makeKey("gpt-4o-2024-08-06", "openai", "chat"): { - Model: "gpt-4o-2024-08-06", - BaseModel: "gpt-4o", - Provider: "openai", - Mode: "chat", - ContextLength: &contextLengthChat, - MaxOutputTokens: capabilityIntPtr(16000), - }, - }, - baseModelIndex: map[string]string{ - "gpt-4o-2024-08-06": "gpt-4o", - }, - } - - entry := mc.GetModelCapabilityEntryForModel("gpt-4o", schemas.OpenAI) - if entry == nil { - t.Fatal("expected capability entry for base-model alias") - } - if entry.Mode != "chat" { - t.Fatalf("expected chat mode to win for alias family, got %q", entry.Mode) - } - if entry.ContextLength == nil || *entry.ContextLength != contextLengthChat { - t.Fatalf("expected alias family context_length=%d, got %#v", contextLengthChat, entry.ContextLength) - } -} - -func TestGetModelCapabilityEntryForModel_ResolvesProviderPrefixedAlias(t *testing.T) { - mc := &ModelCatalog{ - pricingData: map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o-2024-08-06", "openai", "chat"): { - Model: "gpt-4o-2024-08-06", - BaseModel: "gpt-4o", - Provider: "openai", - Mode: "chat", - ContextLength: capabilityIntPtr(128000), - MaxOutputTokens: capabilityIntPtr(16000), - }, - }, - baseModelIndex: map[string]string{ - "gpt-4o-2024-08-06": "gpt-4o", - }, - } - - entry := mc.GetModelCapabilityEntryForModel("openai/gpt-4o", schemas.OpenAI) - if entry == nil { - t.Fatal("expected capability entry for provider-prefixed alias") - } - if entry.Mode != "chat" { - t.Fatalf("expected chat mode for provider-prefixed alias, got %q", entry.Mode) - } -} - -func TestGetModelCapabilityEntryForModel_PrefersLiteralMatchOverAliasFamily(t *testing.T) { - literalContextLength := 32000 - aliasContextLength := 128000 - - mc := &ModelCatalog{ - pricingData: map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): { - Model: "gpt-4o", - BaseModel: "gpt-4o", - Provider: "openai", - Mode: "chat", - ContextLength: &literalContextLength, - MaxOutputTokens: capabilityIntPtr(4000), - }, - makeKey("gpt-4o-2024-08-06", "openai", "chat"): { - Model: "gpt-4o-2024-08-06", - BaseModel: "gpt-4o", - Provider: "openai", - Mode: "chat", - ContextLength: &aliasContextLength, - MaxOutputTokens: capabilityIntPtr(16000), - }, - }, - baseModelIndex: map[string]string{ - "gpt-4o": "gpt-4o", - "gpt-4o-2024-08-06": "gpt-4o", - }, - } - - entry := mc.GetModelCapabilityEntryForModel("gpt-4o", schemas.OpenAI) - if entry == nil { - t.Fatal("expected literal capability entry") - } - if entry.ContextLength == nil || *entry.ContextLength != literalContextLength { - t.Fatalf("expected literal match to win with context_length=%d, got %#v", literalContextLength, entry.ContextLength) - } -} - -func TestCapabilityFieldsRoundTripThroughPricingConversions(t *testing.T) { - modality := "text" - inputCost := float64(1) - outputCost := float64(2) - entry := PricingEntry{ - BaseModel: "gpt-4o", - Provider: "openai", - Mode: "chat", - PricingOptions: PricingOptions{ - InputCostPerToken: &inputCost, - OutputCostPerToken: &outputCost, - }, - ContextLength: capabilityIntPtr(128000), - MaxInputTokens: capabilityIntPtr(64000), - MaxOutputTokens: capabilityIntPtr(16000), - Architecture: &schemas.Architecture{ - Modality: &modality, - }, - } - - table := convertPricingDataToTableModelPricing("gpt-4o", entry) - roundTrip := convertTableModelPricingToPricingData(&table) - - if roundTrip.ContextLength == nil || *roundTrip.ContextLength != 128000 { - t.Fatalf("expected context_length to round-trip, got %#v", roundTrip.ContextLength) - } - if roundTrip.MaxInputTokens == nil || *roundTrip.MaxInputTokens != 64000 { - t.Fatalf("expected max_input_tokens to round-trip, got %#v", roundTrip.MaxInputTokens) - } - if roundTrip.MaxOutputTokens == nil || *roundTrip.MaxOutputTokens != 16000 { - t.Fatalf("expected max_output_tokens to round-trip, got %#v", roundTrip.MaxOutputTokens) - } - if roundTrip.Architecture == nil || roundTrip.Architecture.Modality == nil || *roundTrip.Architecture.Modality != modality { - t.Fatalf("expected architecture to round-trip, got %#v", roundTrip.Architecture) - } -} - -func capabilityIntPtr(v int) *int { return &v } diff --git a/framework/modelcatalog/config.go b/framework/modelcatalog/config.go index 80563eb1e00..50252398802 100644 --- a/framework/modelcatalog/config.go +++ b/framework/modelcatalog/config.go @@ -2,25 +2,18 @@ package modelcatalog import ( "time" + + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/modelcatalog/datasheet" + "github.com/maximhq/bifrost/framework/modelcatalog/keyconfig" ) const ( - DefaultSyncInterval = 24 * time.Hour + DefaultSyncInterval = datasheet.DefaultSyncInterval MinimumPricingSyncIntervalSec = int64(3600) - // syncWorkerTickerPeriod is the fixed interval at which the background sync worker - // wakes up to check whether a sync is due. This is independent of pricingSyncInterval — - // the ticker defines the check granularity, not the sync frequency. - // Kept well below MinimumPricingSyncIntervalSec so the threshold check is not - // defeated by ticker drift when pricingSyncInterval is set near the minimum. - syncWorkerTickerPeriod = 5 * time.Minute - - ConfigLastPricingSyncKey = "LastModelPricingSync" - ConfigLastParamsSyncKey = "LastModelParametersSync" - DefaultPricingURL = "https://getbifrost.ai/datasheet" - DefaultModelParametersURL = "https://getbifrost.ai/datasheet/model-parameters" - DefaultPricingTimeout = 45 * time.Second - DefaultModelParametersTimeout = 45 * time.Second + ConfigLastPricingSyncKey = "LastModelPricingSync" + ConfigLastParamsSyncKey = "LastModelParametersSync" ) // Config is the model pricing configuration. @@ -29,3 +22,53 @@ type Config struct { PricingSyncInterval *int64 `json:"pricing_sync_interval,omitempty"` // seconds ModelParametersURL *string `json:"model_parameters_url,omitempty"` } + +// Type re-exports so external callers can continue importing the legacy +// names (PricingEntry, PricingOptions, etc.) without changing imports. +// Internally these live in the datasheet / keyconfig subpackages. +type ( + PricingEntry = datasheet.Entry + PricingOptions = datasheet.Options + PricingOverride = datasheet.Override + PricingLookupScopes = datasheet.LookupScopes + ScopeKind = datasheet.ScopeKind + MatchType = datasheet.MatchType + + KeyConfigEntry = keyconfig.KeyEntry + AliasOwner = keyconfig.AliasOwner +) + +// Scope kind constants re-exported for callers that compare by value. +const ( + ScopeKindGlobal = datasheet.ScopeKindGlobal + ScopeKindProvider = datasheet.ScopeKindProvider + ScopeKindProviderKey = datasheet.ScopeKindProviderKey + ScopeKindVirtualKey = datasheet.ScopeKindVirtualKey + ScopeKindVirtualKeyProvider = datasheet.ScopeKindVirtualKeyProvider + ScopeKindVirtualKeyProviderKey = datasheet.ScopeKindVirtualKeyProviderKey + + MatchTypeExact = datasheet.MatchTypeExact + MatchTypeWildcard = datasheet.MatchTypeWildcard +) + +// PricingLookupScopesFromContext is re-exported so callers don't have to +// change their imports. +func PricingLookupScopesFromContext(ctx *schemas.BifrostContext, provider string) *PricingLookupScopes { + return datasheet.LookupScopesFromContext(ctx, provider) +} + +// Sync timing defaults re-exported from datasheet for consumers of the +// historical constants. +const ( + DefaultPricingURL = datasheet.DefaultURL + DefaultModelParametersURL = datasheet.DefaultModelParametersURL + DefaultPricingTimeout = datasheet.DefaultPricingTimeout + DefaultModelParametersTimeout = datasheet.DefaultModelParametersTimeout +) + +// syncWorkerTickerPeriod is the fixed interval at which the background sync worker +// wakes up to check whether a sync is due. This is independent of pricingSyncInterval — +// the ticker defines the check granularity, not the sync frequency. +// Kept well below MinimumPricingSyncIntervalSec so the threshold check is not +// defeated by ticker drift when pricingSyncInterval is set near the minimum. +const syncWorkerTickerPeriod = 5 * time.Minute diff --git a/framework/modelcatalog/datasheet/cost.go b/framework/modelcatalog/datasheet/cost.go index 5af169a9ae1..5d82e48703c 100644 --- a/framework/modelcatalog/datasheet/cost.go +++ b/framework/modelcatalog/datasheet/cost.go @@ -10,10 +10,10 @@ import ( configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" ) -// CalculateCost computes the dollar cost for a Bifrost response. Handles all -// request types, cache-debug billing, and tiered pricing. If scopes is nil, -// an empty LookupScopes is used; global and provider-scoped overrides may -// still apply since the provider is derived from the response. +// CalculateCost calculates the cost of a Bifrost response. +// It handles all request types, cache debug billing, and tiered pricing. +// If scopes is nil, an empty LookupScopes is used; global and provider-scoped +// overrides may still apply since the provider is derived from the response. func (s *Store) CalculateCost(result *schemas.BifrostResponse, scopes *LookupScopes) float64 { if result == nil { return 0 @@ -24,13 +24,16 @@ func (s *Store) CalculateCost(result *schemas.BifrostResponse, scopes *LookupSco lookupScopes = *scopes } + // Handle semantic cache billing cacheDebug := result.GetExtraFields().CacheDebug if cacheDebug != nil { return s.calculateCostWithCache(result, cacheDebug, lookupScopes) } + return s.calculateBaseCost(result, lookupScopes) } +// calculateCostWithCache handles cost calculation when semantic cache debug info is present. func (s *Store) calculateCostWithCache(result *schemas.BifrostResponse, cacheDebug *schemas.BifrostCacheDebug, scopes LookupScopes) float64 { if cacheDebug.CacheHit { // Direct cache hit — no LLM call, no cost @@ -43,12 +46,14 @@ func (s *Store) calculateCostWithCache(result *schemas.BifrostResponse, cacheDeb } return 0 } + // Cache miss — full LLM cost + embedding lookup cost baseCost := s.calculateBaseCost(result, scopes) embeddingCost := s.computeCacheEmbeddingCost(cacheDebug, scopes) return baseCost + embeddingCost } +// computeCacheEmbeddingCost calculates the embedding cost for a semantic cache lookup. func (s *Store) computeCacheEmbeddingCost(cacheDebug *schemas.BifrostCacheDebug, scopes LookupScopes) float64 { if cacheDebug == nil || cacheDebug.ProviderUsed == nil || cacheDebug.ModelUsed == nil || cacheDebug.InputTokens == nil { return 0 @@ -56,9 +61,9 @@ func (s *Store) computeCacheEmbeddingCost(cacheDebug *schemas.BifrostCacheDebug, if scopes.Provider == "" { scopes.Provider = *cacheDebug.ProviderUsed } - // Cache-debug pricing only carries a single model identifier (whatever the - // cache recorded). Maps to RoutingInfo.Model — no alias resolution context - // exists for the cache-replayed request. + // Cache-debug pricing has only a single model identifier (whatever the + // cache recorded). Maps to RoutingInfo.Model — no alias resolution + // context exists for the cache-replayed request. pricing := s.resolvePricing(schemas.RoutingInfo{ Provider: schemas.ModelProvider(*cacheDebug.ProviderUsed), Model: *cacheDebug.ModelUsed, @@ -69,6 +74,7 @@ func (s *Store) computeCacheEmbeddingCost(cacheDebug *schemas.BifrostCacheDebug, return float64(*cacheDebug.InputTokens) * tieredInputRate(pricing, *cacheDebug.InputTokens, serviceTier{}) } +// computeContainerCreationCost returns the cost for creating a container from an already-resolved pricing entry. func computeContainerCreationCost(pricing *configstoreTables.TableModelPricing) float64 { if pricing == nil || pricing.CodeInterpreterCostPerSession == nil { return 0 @@ -76,15 +82,18 @@ func computeContainerCreationCost(pricing *configstoreTables.TableModelPricing) return *pricing.CodeInterpreterCostPerSession } +// calculateBaseCost extracts usage from the response and routes to the appropriate compute function. func (s *Store) calculateBaseCost(result *schemas.BifrostResponse, scopes LookupScopes) float64 { extraFields := result.GetExtraFields() if extraFields == nil { return 0 } + // Read routing info populated by core.bifrost at request time. + // // Backward-compat fallback: when the caller (e.g. LoggerPlugin's - // RecalculateCosts replaying logs written before RoutingInfo existed, or - // third-party plugins still on the legacy ExtraFields shape) leaves + // RecalculateCosts replaying logs written before RoutingInfo existed, + // or third-party plugins still on the legacy ExtraFields shape) leaves // RoutingInfo empty, synthesise one from the deprecated triplet so // pricing keeps working. Triggered only when RoutingInfo is fully // unset — partial population is trusted as-is. @@ -98,28 +107,33 @@ func (s *Store) calculateBaseCost(result *schemas.BifrostResponse, scopes Lookup } requestType := extraFields.RequestType + // Extract usage data from the response (passthrough and native paths unified) input := extractCostInput(result) - // Provider-computed cost wins when present. + // If provider already computed cost, use it if input.usage != nil && input.usage.Cost != nil && input.usage.Cost.TotalCost > 0 { return input.usage.Cost.TotalCost } - // Nothing to price. + // If no usage data at all, nothing to price if input.usage == nil && input.audioSeconds == nil && input.audioTokenDetails == nil && input.imageUsage == nil && input.videoSeconds == nil && input.audioTextInputChars == 0 && input.ocrProcessedPages == nil && input.containerIdentifierString == "" { return 0 } if result.PassthroughResponse != nil { + // Infer request type from usage fields + path; passthrough bypasses stream normalization. requestType = inferPassthroughRequestType(routingInfo.Provider, extraFields.PassthroughPath, result.PassthroughResponse.PassthroughUsage) } else { + // Normalize stream request types to their base type for pricing lookup requestType = normalizeStreamRequestType(requestType) } // When a pricing model override is set (e.g. container creates always look // up "container"), it replaces the lookup hierarchy entirely. Build a // synthetic RoutingInfo that reuses Provider but pins the model fields to - // the container identifier so per-container overrides stay addressable. + // the container identifier — the lookup tries it as ModelName, the + // override key is the container identifier so per-container overrides + // stay addressable. if input.containerIdentifierString != "" { routingInfo = schemas.RoutingInfo{ Provider: routingInfo.Provider, @@ -132,8 +146,9 @@ func (s *Store) calculateBaseCost(result *schemas.BifrostResponse, scopes Lookup return 0 } + // Route to the appropriate compute function switch requestType { - case schemas.ChatCompletionRequest, schemas.TextCompletionRequest, schemas.ResponsesRequest, schemas.RealtimeRequest: + case schemas.ChatCompletionRequest, schemas.TextCompletionRequest, schemas.ResponsesRequest, schemas.RealtimeRequest, schemas.CompactionRequest: return computeTextCost(pricing, input.usage, input.tier) case schemas.EmbeddingRequest: return computeEmbeddingCost(pricing, input.usage, input.tier) @@ -178,6 +193,9 @@ func extractCostInput(result *schemas.BifrostResponse) costInput { input.usage = responsesUsageToBifrostUsage(result.ResponsesResponse.Usage) input.tier = tierFromString(result.ResponsesResponse.ServiceTier) + case result.CompactionResponse != nil && result.CompactionResponse.Usage != nil: + input.usage = responsesUsageToBifrostUsage(result.CompactionResponse.Usage) + case result.ResponsesStreamResponse != nil && result.ResponsesStreamResponse.Response != nil && result.ResponsesStreamResponse.Response.Usage != nil: input.usage = responsesUsageToBifrostUsage(result.ResponsesStreamResponse.Response.Usage) input.tier = tierFromString(result.ResponsesStreamResponse.Response.ServiceTier) @@ -203,9 +221,13 @@ func extractCostInput(result *schemas.BifrostResponse) costInput { input.usage, input.audioSeconds, input.audioTokenDetails = extractTranscriptionUsage(result.TranscriptionStreamResponse.Usage) case result.ImageGenerationResponse != nil: + // Defensive copy: populateOutputImageCount writes into imageUsage, + // and we must not mutate the caller's BifrostResponse during what is + // otherwise a pure read path. if result.ImageGenerationResponse.Usage != nil { - input.imageUsage = result.ImageGenerationResponse.Usage + input.imageUsage = result.ImageGenerationResponse.Usage.DeepCopy() } else { + // No usage data but response exists — default to empty so per-image pricing can apply input.imageUsage = &schemas.ImageUsage{} } populateOutputImageCount(input.imageUsage, len(result.ImageGenerationResponse.Data)) @@ -215,8 +237,12 @@ func extractCostInput(result *schemas.BifrostResponse) costInput { } case result.ImageGenerationStreamResponse != nil: + // Defensive copy mirrors the non-stream path so CalculateCost never + // aliases the caller's response — keeps the read-only invariant + // uniform and prevents accidental mutation if image-count derivation + // is later added on this branch. if result.ImageGenerationStreamResponse.Usage != nil { - input.imageUsage = result.ImageGenerationStreamResponse.Usage + input.imageUsage = result.ImageGenerationStreamResponse.Usage.DeepCopy() } else { input.imageUsage = &schemas.ImageUsage{} } @@ -256,6 +282,7 @@ func responsesUsageToBifrostUsage(u *schemas.ResponsesResponseUsage) *schemas.Bi TotalTokens: u.TotalTokens, Cost: u.Cost, } + // Map token details for cache and search query pricing if u.InputTokensDetails != nil { usage.PromptTokensDetails = &schemas.ChatPromptTokensDetails{ TextTokens: u.InputTokensDetails.TextTokens, @@ -299,6 +326,7 @@ func extractTranscriptionUsage(u *schemas.TranscriptionUsage) (*schemas.BifrostL } else { usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens } + var audioTokenDetails *schemas.TranscriptionUsageInputTokenDetails if u.InputTokenDetails != nil { audioTokenDetails = &schemas.TranscriptionUsageInputTokenDetails{ @@ -306,6 +334,7 @@ func extractTranscriptionUsage(u *schemas.TranscriptionUsage) (*schemas.BifrostL TextTokens: u.InputTokenDetails.TextTokens, } } + return usage, u.Seconds, audioTokenDetails } @@ -313,6 +342,7 @@ func extractTranscriptionUsage(u *schemas.TranscriptionUsage) (*schemas.BifrostL // Per-request-type cost computation // --------------------------------------------------------------------------- +// computeTextCost handles chat, text completion, and responses requests. func computeTextCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, tier serviceTier) float64 { if usage == nil { return 0 @@ -322,6 +352,7 @@ func computeTextCost(pricing *configstoreTables.TableModelPricing, usage *schema promptTokens := usage.PromptTokens completionTokens := usage.CompletionTokens + // Extract cached token counts cachedReadTokens := 0 cachedWriteTokens := 0 cachedWriteTokensAbove1hr := 0 @@ -339,22 +370,28 @@ func computeTextCost(pricing *configstoreTables.TableModelPricing, usage *schema cacheCreationInputRate := tieredCacheCreationInputTokenRate(pricing, totalTokens, tier) cacheCreationInputAbove1hrInputRate := tieredCacheCreationInputAbove1hrTokenRate(pricing, totalTokens, tier) - // Clamp cached token counts to avoid negative billing on malformed provider payloads. + // Clamp cached token counts to avoid negative billing on malformed provider payloads if cachedReadTokens > promptTokens { cachedReadTokens = promptTokens } if cachedWriteTokens > promptTokens-cachedReadTokens { cachedWriteTokens = promptTokens - cachedReadTokens } + // Should not happen, but just in case if cachedWriteTokensAbove1hr > cachedWriteTokens { cachedWriteTokensAbove1hr = cachedWriteTokens } + // Input cost: non-cached tokens at regular rate nonCachedPrompt := promptTokens - cachedReadTokens - cachedWriteTokens inputCost := float64(nonCachedPrompt) * inputRate + + // Add cached prompt tokens at cache read rate if cachedReadTokens > 0 { inputCost += float64(cachedReadTokens) * cacheReadInputRate } + + // Add cached write tokens at cache creation rate if cachedWriteTokens > 0 { if cachedWriteTokensAbove1hr > 0 { inputCost += float64(cachedWriteTokensAbove1hr) * cacheCreationInputAbove1hrInputRate @@ -364,8 +401,9 @@ func computeTextCost(pricing *configstoreTables.TableModelPricing, usage *schema outputCost := float64(completionTokens) * outputRate - // Audio token cost: when token details include audio tokens, price them at - // the dedicated audio rate and subtract from the text token costs above. + // Audio token cost: when token details include audio tokens, price them + // at the dedicated audio rate and subtract from the text token costs above. + // Realtime and audio-enabled chat models report audio tokens in details. audioCost := 0.0 inputAudioTokens := 0 outputAudioTokens := 0 @@ -386,12 +424,14 @@ func computeTextCost(pricing *configstoreTables.TableModelPricing, usage *schema outputAudioTokens = completionTokens } if inputAudioTokens > 0 && pricing.InputCostPerAudioToken != nil { + // Subtract audio tokens charged at text rate, add at audio rate. audioCost += float64(inputAudioTokens) * (*pricing.InputCostPerAudioToken - inputRate) } if outputAudioTokens > 0 && pricing.OutputCostPerAudioToken != nil { audioCost += float64(outputAudioTokens) * (*pricing.OutputCostPerAudioToken - outputRate) } + // Search query cost searchCost := 0.0 if pricing.SearchContextCostPerQuery != nil && usage.CompletionTokensDetails != nil && usage.CompletionTokensDetails.NumSearchQueries != nil { searchCost = float64(*usage.CompletionTokensDetails.NumSearchQueries) * *pricing.SearchContextCostPerQuery @@ -400,6 +440,7 @@ func computeTextCost(pricing *configstoreTables.TableModelPricing, usage *schema return inputCost + outputCost + audioCost + searchCost } +// computeEmbeddingCost handles embedding requests (input-only). func computeEmbeddingCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, tier serviceTier) float64 { if usage == nil { return 0 @@ -407,27 +448,34 @@ func computeEmbeddingCost(pricing *configstoreTables.TableModelPricing, usage *s return float64(usage.PromptTokens) * tieredInputRate(pricing, usage.TotalTokens, tier) } +// computeRerankCost handles rerank requests. func computeRerankCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, tier serviceTier) float64 { if usage == nil { return 0 } inputCost := float64(usage.PromptTokens) * tieredInputRate(pricing, usage.TotalTokens, tier) outputCost := float64(usage.CompletionTokens) * tieredOutputRate(pricing, usage.TotalTokens, tier) + searchCost := 0.0 if pricing.SearchContextCostPerQuery != nil && usage.CompletionTokensDetails != nil && usage.CompletionTokensDetails.NumSearchQueries != nil { searchCost = float64(*usage.CompletionTokensDetails.NumSearchQueries) * *pricing.SearchContextCostPerQuery } + return inputCost + outputCost + searchCost } -// computeSpeechCost handles speech (TTS) requests. Per-character pricing -// (InputCostPerCharacter) is first-class — providers like OpenAI TTS, -// ElevenLabs, and AWS Polly bill per character of input text. PromptTokens -// is treated as the character count since TTS providers report their -// billable unit in that field. Output falls back to per-second duration -// when no audio token rate is configured. +// computeSpeechCost handles speech (TTS) requests. +// Input is text (PromptTokens), output is audio (CompletionTokens). +// +// Per-character pricing (InputCostPerCharacter) is used as first-class support for TTS/audio +// models — providers such as OpenAI TTS, ElevenLabs, and AWS Polly bill per character of +// input text rather than per token. PromptTokens from usage is treated as the character count +// since TTS providers report their billable unit in that field. +// Output falls back to per-second duration when no audio token rate is configured. func computeSpeechCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, audioSeconds *int, audioTextInputChars int, tier serviceTier) float64 { totalTokens := safeTotalTokens(usage) + + // Input: per-character rate takes precedence for TTS/audio models inputCost := 0.0 if audioTextInputChars > 0 { if pricing.InputCostPerCharacter != nil { @@ -438,63 +486,93 @@ func computeSpeechCost(pricing *configstoreTables.TableModelPricing, usage *sche } else if usage != nil && usage.PromptTokens > 0 { inputCost = float64(usage.PromptTokens) * tieredInputRate(pricing, totalTokens, tier) } + + // Output: audio tokens first, then per-second fallback outputCost := computeAudioOutputCost(pricing, usage, audioSeconds, totalTokens, tier) + return inputCost + outputCost } +// computeTranscriptionCost handles transcription (STT) requests. +// Input is audio, output is text (CompletionTokens). +// Input and output are calculated independently — tokens first, then per-second fallback. func computeTranscriptionCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, audioSeconds *int, audioTokenDetails *schemas.TranscriptionUsageInputTokenDetails, tier serviceTier) float64 { totalTokens := safeTotalTokens(usage) + + // Input: audio tokens/details first, then per-second fallback inputCost := computeAudioInputCost(pricing, usage, audioSeconds, audioTokenDetails, totalTokens, tier) + + // Output: text tokens outputCost := 0.0 if usage != nil && usage.CompletionTokens > 0 { outputCost = float64(usage.CompletionTokens) * tieredOutputRate(pricing, totalTokens, tier) } + return inputCost + outputCost } +// computeAudioInputCost calculates input cost for audio: audio token details first, +// then generic input tokens, then per-second duration fallback. func computeAudioInputCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, audioSeconds *int, audioTokenDetails *schemas.TranscriptionUsageInputTokenDetails, totalTokens int, tier serviceTier) float64 { + // Audio token detail pricing (audio + text token breakdown) if audioTokenDetails != nil && (audioTokenDetails.AudioTokens > 0 || audioTokenDetails.TextTokens > 0) { return float64(audioTokenDetails.AudioTokens)*tieredAudioTokenInputRate(pricing, totalTokens, tier) + float64(audioTokenDetails.TextTokens)*tieredInputRate(pricing, totalTokens, tier) } + + // Generic input tokens if usage != nil && usage.PromptTokens > 0 { return float64(usage.PromptTokens) * tieredInputRate(pricing, totalTokens, tier) } + + // Per-second duration fallback if audioSeconds != nil && *audioSeconds > 0 { if rate := tieredAudioInputPerSecondRate(pricing, totalTokens); rate > 0 { return float64(*audioSeconds) * rate } } + return 0 } +// computeAudioOutputCost calculates output cost for audio: audio tokens first, +// then generic output tokens, then per-second duration fallback. func computeAudioOutputCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, audioSeconds *int, totalTokens int, tier serviceTier) float64 { + // Audio-specific output tokens if usage != nil && usage.CompletionTokens > 0 { return float64(usage.CompletionTokens) * tieredAudioTokenOutputRate(pricing, totalTokens, tier) } + + // Per-second duration fallback if audioSeconds != nil && *audioSeconds > 0 { if pricing.OutputCostPerSecond != nil { return float64(*audioSeconds) * *pricing.OutputCostPerSecond } } + return 0 } -// computeImageCost handles image generation. Input and output are independent — -// each tries token-based pricing first, then per-pixel, then per-image fallback. -// imageQuality must be "low"/"medium"/"high"/"auto" to use quality-specific rates. +// computeImageCost handles image generation requests. +// Input and output are calculated independently — each tries token-based pricing first, +// then per-pixel pricing, falling back to per-image count pricing. +// imageQuality must be one of "low", "medium", "high", "auto" to use quality-specific rates; other values use base rates. func computeImageCost(pricing *configstoreTables.TableModelPricing, imageUsage *schemas.ImageUsage, imageSize string, imageQuality string, tier serviceTier) float64 { if imageUsage == nil { return 0 } + totalTokens := imageUsage.TotalTokens pixels := parseImagePixels(imageSize) inputCost := computeImageInputCost(pricing, imageUsage, totalTokens, pixels, tier) outputCost := computeImageOutputCost(pricing, imageUsage, totalTokens, pixels, imageQuality, tier) + return inputCost + outputCost } +// computeImageInputCost calculates input cost: tokens first, then per-pixel, then per-image count fallback. func computeImageInputCost(pricing *configstoreTables.TableModelPricing, imageUsage *schemas.ImageUsage, totalTokens int, pixels int, tier serviceTier) float64 { + // Try token-based pricing first var inputTextTokens, inputImageTokens int if imageUsage.InputTokensDetails != nil { inputImageTokens = imageUsage.InputTokensDetails.ImageTokens @@ -502,20 +580,29 @@ func computeImageInputCost(pricing *configstoreTables.TableModelPricing, imageUs } else { inputTextTokens = imageUsage.InputTokens } + if inputTextTokens > 0 || inputImageTokens > 0 { return float64(inputTextTokens)*tieredInputRate(pricing, totalTokens, tier) + float64(inputImageTokens)*tieredImageInputRate(pricing, totalTokens, tier) } + + // Per-pixel pricing fallback if pricing.InputCostPerPixel != nil && pixels > 0 && imageUsage.NumInputImages > 0 { return float64(pixels*imageUsage.NumInputImages) * *pricing.InputCostPerPixel } + + // Fall back to per-image count pricing if pricing.InputCostPerImage != nil && imageUsage.NumInputImages > 0 { return float64(imageUsage.NumInputImages) * *pricing.InputCostPerImage } + return 0 } +// computeImageOutputCost calculates output cost: tokens first, then per-pixel, then per-image count fallback. +// imageQuality: "low", "medium", "high", "auto" use quality-specific rates when available; other values use base/size-tier rates. func computeImageOutputCost(pricing *configstoreTables.TableModelPricing, imageUsage *schemas.ImageUsage, totalTokens int, pixels int, imageQuality string, tier serviceTier) float64 { + // Try token-based pricing first var outputTextTokens, outputImageTokens int if imageUsage.OutputTokensDetails != nil { outputImageTokens = imageUsage.OutputTokensDetails.ImageTokens @@ -523,10 +610,13 @@ func computeImageOutputCost(pricing *configstoreTables.TableModelPricing, imageU } else { outputImageTokens = imageUsage.OutputTokens } + if outputTextTokens > 0 || outputImageTokens > 0 { return float64(outputTextTokens)*tieredOutputRate(pricing, totalTokens, tier) + float64(outputImageTokens)*tieredImageOutputRate(pricing, totalTokens, tier) } + + // Per-pixel pricing fallback if pricing.OutputCostPerPixel != nil && pixels > 0 { numOutputImages := 1 if imageUsage.OutputTokensDetails != nil && imageUsage.OutputTokensDetails.NImages > 0 { @@ -535,6 +625,8 @@ func computeImageOutputCost(pricing *configstoreTables.TableModelPricing, imageU return float64(pixels*numOutputImages) * *pricing.OutputCostPerPixel } + // Fall back to per-image count pricing with size-tier selection + // TODO: handle premium image flag when it becomes available in imageUsage numOutputImages := 1 if imageUsage.OutputTokensDetails != nil && imageUsage.OutputTokensDetails.NImages > 0 { numOutputImages = imageUsage.OutputTokensDetails.NImages @@ -583,12 +675,16 @@ func computeImageOutputCost(pricing *configstoreTables.TableModelPricing, imageU if perImageRate != nil { return float64(numOutputImages) * *perImageRate } + return 0 } +// computeVideoCost handles video generation requests. +// Input and output are calculated independently — tokens first, then per-second fallback. func computeVideoCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, videoSeconds *int, tier serviceTier) float64 { totalTokens := safeTotalTokens(usage) + // Input: text prompt tokens first, then per-second fallback inputCost := 0.0 if usage != nil && usage.PromptTokens > 0 { inputCost = float64(usage.PromptTokens) * tieredInputRate(pricing, totalTokens, tier) @@ -598,6 +694,7 @@ func computeVideoCost(pricing *configstoreTables.TableModelPricing, usage *schem } } + // Output: completion tokens first, then per-second fallback outputCost := 0.0 if usage != nil && usage.CompletionTokens > 0 { outputCost = float64(usage.CompletionTokens) * tieredOutputRate(pricing, totalTokens, tier) @@ -608,9 +705,12 @@ func computeVideoCost(pricing *configstoreTables.TableModelPricing, usage *schem outputCost = float64(*videoSeconds) * *pricing.OutputCostPerSecond } } + return inputCost + outputCost } +// computeOCRCost handles OCR requests, billing per page processed. +// ocr_cost_per_page covers base processing; annotation_cost_per_page is added when set. func computeOCRCost(pricing *configstoreTables.TableModelPricing, ocrProcessedPages *int, ocrIsAnnotated *bool) float64 { if ocrProcessedPages == nil { return 0 @@ -627,9 +727,10 @@ func computeOCRCost(pricing *configstoreTables.TableModelPricing, ocrProcessedPa } // --------------------------------------------------------------------------- -// Tier resolution and rate selectors +// Helpers // --------------------------------------------------------------------------- +// tierFromString constructs a serviceTier from an OpenAI service_tier response value. func tierFromString(s *schemas.BifrostServiceTier) serviceTier { if s == nil { return serviceTier{} @@ -644,6 +745,8 @@ func tierFromString(s *schemas.BifrostServiceTier) serviceTier { } } +// tieredInputRate returns the effective per-token input rate based on total token count. +// Flex applies a flat rate. Priority-specific tier rates are preferred where available. func tieredInputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { if tier.isFlex && pricing.InputCostPerTokenFlex != nil { return *pricing.InputCostPerTokenFlex @@ -676,6 +779,8 @@ func tieredInputRate(pricing *configstoreTables.TableModelPricing, totalTokens i return 0 } +// tieredOutputRate returns the effective per-token output rate based on total token count. +// Flex applies a flat rate. Priority-specific tier rates are preferred where available. func tieredOutputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { if tier.isFlex && pricing.OutputCostPerTokenFlex != nil { return *pricing.OutputCostPerTokenFlex @@ -699,15 +804,20 @@ func tieredOutputRate(pricing *configstoreTables.TableModelPricing, totalTokens if totalTokens > TokenTierAbove128K && pricing.OutputCostPerTokenAbove128kTokens != nil { return *pricing.OutputCostPerTokenAbove128kTokens } + if tier.isPriority && pricing.OutputCostPerTokenPriority != nil { return *pricing.OutputCostPerTokenPriority } + if pricing.OutputCostPerToken != nil { return *pricing.OutputCostPerToken } + return 0 } +// tieredImageInputRate returns the effective rate for image tokens on the input side. +// Falls back to the general tieredInputRate when no image-specific rate is configured. func tieredImageInputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { if totalTokens > TokenTierAbove128K && pricing.InputCostPerImageAbove128kTokens != nil { return *pricing.InputCostPerImageAbove128kTokens @@ -718,6 +828,8 @@ func tieredImageInputRate(pricing *configstoreTables.TableModelPricing, totalTok return tieredInputRate(pricing, totalTokens, tier) } +// tieredImageOutputRate returns the effective rate for image tokens on the output side. +// Falls back to the general tieredOutputRate when no image-specific rate is configured. func tieredImageOutputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { if pricing.OutputCostPerImageToken != nil { return *pricing.OutputCostPerImageToken @@ -725,6 +837,7 @@ func tieredImageOutputRate(pricing *configstoreTables.TableModelPricing, totalTo return tieredOutputRate(pricing, totalTokens, tier) } +// tieredAudioInputPerSecondRate returns the effective per-second rate for audio input. func tieredAudioInputPerSecondRate(pricing *configstoreTables.TableModelPricing, totalTokens int) float64 { if totalTokens > TokenTierAbove128K && pricing.InputCostPerAudioPerSecondAbove128kTokens != nil { return *pricing.InputCostPerAudioPerSecondAbove128kTokens @@ -738,6 +851,7 @@ func tieredAudioInputPerSecondRate(pricing *configstoreTables.TableModelPricing, return 0 } +// tieredVideoInputPerSecondRate returns the effective per-second rate for video input. func tieredVideoInputPerSecondRate(pricing *configstoreTables.TableModelPricing, totalTokens int) float64 { if totalTokens > TokenTierAbove128K && pricing.InputCostPerVideoPerSecondAbove128kTokens != nil { return *pricing.InputCostPerVideoPerSecondAbove128kTokens @@ -748,6 +862,8 @@ func tieredVideoInputPerSecondRate(pricing *configstoreTables.TableModelPricing, return 0 } +// tieredAudioTokenInputRate returns the effective per-token rate for audio input tokens. +// Falls back to the general tieredInputRate when no audio-specific rate is configured. func tieredAudioTokenInputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { if pricing.InputCostPerAudioToken != nil { return *pricing.InputCostPerAudioToken @@ -755,6 +871,8 @@ func tieredAudioTokenInputRate(pricing *configstoreTables.TableModelPricing, tot return tieredInputRate(pricing, totalTokens, tier) } +// tieredAudioTokenOutputRate returns the effective per-token rate for audio output tokens. +// Falls back to the general tieredOutputRate when no audio-specific rate is configured. func tieredAudioTokenOutputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { if pricing.OutputCostPerAudioToken != nil { return *pricing.OutputCostPerAudioToken @@ -791,8 +909,8 @@ func tieredCacheReadInputTokenRate(pricing *configstoreTables.TableModelPricing, return tieredInputRate(pricing, totalTokens, tier) } -// Note: flex tier is not checked here because cache creation isn't a concept -// in OpenAI's pricing model (the only flex-tier provider). Only cache read +// Note: flex tier is not checked here because cache creation is not a concept in +// OpenAI's pricing model (the only provider that uses flex tier). Only cache read // has a flex-specific rate. func tieredCacheCreationInputTokenRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { if totalTokens > TokenTierAbove200K && pricing.CacheCreationInputTokenCostAbove200kTokens != nil { @@ -821,6 +939,8 @@ func safeTotalTokens(usage *schemas.BifrostLLMUsage) int { return usage.TotalTokens } +// parseImagePixels parses a size string like "1024x1024" into total pixel count. +// Returns 0 if the size string is empty or malformed. func parseImagePixels(size string) int { if size == "" { return 0 @@ -840,6 +960,8 @@ func parseImagePixels(size string) int { return w * h } +// populateOutputImageCount sets the output image count on ImageUsage from len(Data) +// when OutputTokensDetails.NImages is not already populated. func populateOutputImageCount(imageUsage *schemas.ImageUsage, dataLen int) { if imageUsage == nil || dataLen == 0 { return @@ -888,9 +1010,7 @@ func (s *Store) resolvePricing(routingInfo schemas.RoutingInfo, requestType sche if overrideKey == "" { overrideKey = routingInfo.Model } - if s.logger != nil { - s.logger.Debug("looking up pricing for wire model %s and provider %s of request type %s", overrideKey, provider, normalizeRequestType(requestType)) - } + s.logger.Debug("looking up pricing for wire model %s and provider %s of request type %s", overrideKey, provider, normalizeRequestType(requestType)) if scopes.Provider == "" { scopes.Provider = provider @@ -905,35 +1025,39 @@ func (s *Store) resolvePricing(routingInfo schemas.RoutingInfo, requestType sche result, _ := s.applyPricingOverrides(overrideKey, requestType, *base, scopes) return &result } - if s.logger != nil { - s.logger.Debug("pricing not found for %s, trying next candidate", candidate) - } + s.logger.Debug("pricing not found for %s, trying next candidate", candidate) } - // No base catalog entry found — still try overrides in case the user - // defined override-only pricing for a model outside the built-in catalog. - if s.logger != nil { - s.logger.Debug("pricing not found for any candidate (provider %s), trying override-only pricing keyed by %s", provider, overrideKey) - } + // No base catalog entry found; still try overrides in case the user defined + // override-only pricing for a model not in the built-in catalog. + s.logger.Debug("pricing not found for any candidate (provider %s), trying override-only pricing keyed by %s", provider, overrideKey) result, applied := s.applyPricingOverrides(overrideKey, requestType, configstoreTables.TableModelPricing{}, scopes) if applied { return &result } - if s.logger != nil { - s.logger.Debug("no pricing found for wire model %s and provider %s, skipping cost calculation", overrideKey, provider) - } + s.logger.Debug("no pricing found for wire model %s and provider %s, skipping cost calculation", overrideKey, provider) return nil } -// getBasePricing looks up catalog pricing for (model, provider, requestType) -// with provider-specific fallback chains: +// getBasePricing looks up catalog pricing for the given model, provider, and request type. +// It applies a provider-specific fallback chain when an exact match is not found: +// +// - Gemini: retries under the "vertex" provider, then falls back to chat mode for Responses requests. +// - Vertex: strips the "provider/model" prefix and retries, then falls back to chat mode for Responses requests. +// - Bedrock: prepends the "anthropic." namespace for Claude models, then falls back to chat mode for Responses requests. +// - All providers: for Responses/ResponsesStream requests, retries the lookup in chat mode. +// - All providers: for ImageEdit/ImageVariation requests, retries the lookup in image-generation mode. +// +// The method acquires a read lock for the duration of the lookup. +// +// Input: model — exact model name to look up. // -// - Gemini: retries under "vertex", then chat-mode fallback for Responses. -// - Vertex: strips "provider/model" prefix and retries, then chat-mode for Responses. -// - Bedrock: prepends "anthropic." for Claude models, then chat-mode for Responses. -// - All providers: Responses/ResponsesStream falls back to chat mode. -// - All providers: ImageEdit/ImageVariation falls back to image-generation mode. -// - ContainerCreate: chat mode for the model, then base "container" entry. +// provider — provider identifier (e.g. "openai", "anthropic"). +// requestType — the request type used to derive the pricing mode. +// +// Output: TableModelPricing — the matched pricing row (zero value when not found). +// +// bool — true when a pricing entry was found, false otherwise. func (s *Store) getBasePricing(model, provider string, requestType schemas.RequestType) (*configstoreTables.TableModelPricing, bool) { s.mu.RLock() defer s.mu.RUnlock() @@ -945,18 +1069,17 @@ func (s *Store) getBasePricing(model, provider string, requestType schemas.Reque return &pricing, true } + // Lookup in vertex if gemini not found if provider == string(schemas.Gemini) { - if s.logger != nil { - s.logger.Debug("primary lookup failed, trying vertex provider for the same model") - } + s.logger.Debug("primary lookup failed, trying vertex provider for the same model") pricing, ok = s.pricingData[makeKey(model, "vertex", mode)] if ok { return &pricing, true } - if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest { - if s.logger != nil { - s.logger.Debug("secondary lookup failed, trying vertex provider for the same model in chat completion") - } + + // Lookup in chat if responses not found + if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest || requestType == schemas.CompactionRequest { + s.logger.Debug("secondary lookup failed, trying vertex provider for the same model in chat completion") pricing, ok = s.pricingData[makeKey(model, "vertex", normalizeRequestType(schemas.ChatCompletionRequest))] if ok { return &pricing, true @@ -965,21 +1088,18 @@ func (s *Store) getBasePricing(model, provider string, requestType schemas.Reque } if provider == string(schemas.Vertex) { - // Vertex models can be of the form "provider/model" — try without the - // provider prefix, keeping the original provider. + // Vertex models can be of the form "provider/model", so try to lookup the model without the provider prefix and keep the original provider if strings.Contains(model, "/") { modelWithoutProvider := strings.SplitN(model, "/", 2)[1] - if s.logger != nil { - s.logger.Debug("primary lookup failed, trying vertex provider for model with provider/model format %s", modelWithoutProvider) - } + s.logger.Debug("primary lookup failed, trying vertex provider for the same model with provider/model format %s", modelWithoutProvider) pricing, ok = s.pricingData[makeKey(modelWithoutProvider, "vertex", mode)] if ok { return &pricing, true } - if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest { - if s.logger != nil { - s.logger.Debug("secondary lookup failed, trying vertex provider for the same model in chat completion") - } + + // Lookup in chat if responses not found + if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest || requestType == schemas.CompactionRequest { + s.logger.Debug("secondary lookup failed, trying vertex provider for the same model in chat completion") pricing, ok = s.pricingData[makeKey(modelWithoutProvider, "vertex", normalizeRequestType(schemas.ChatCompletionRequest))] if ok { return &pricing, true @@ -989,18 +1109,17 @@ func (s *Store) getBasePricing(model, provider string, requestType schemas.Reque } if provider == string(schemas.Bedrock) { + // If model is claude without "anthropic." prefix, try with "anthropic." prefix if !strings.Contains(model, "anthropic.") && schemas.IsAnthropicModel(model) { - if s.logger != nil { - s.logger.Debug("primary lookup failed, trying with anthropic. prefix for the same model") - } + s.logger.Debug("primary lookup failed, trying with anthropic. prefix for the same model") pricing, ok = s.pricingData[makeKey("anthropic."+model, provider, mode)] if ok { return &pricing, true } - if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest { - if s.logger != nil { - s.logger.Debug("secondary lookup failed, trying chat provider for the same model in chat completion") - } + + // Lookup in chat if responses not found + if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest || requestType == schemas.CompactionRequest { + s.logger.Debug("secondary lookup failed, trying chat provider for the same model in chat completion") pricing, ok = s.pricingData[makeKey("anthropic."+model, provider, normalizeRequestType(schemas.ChatCompletionRequest))] if ok { return &pricing, true @@ -1009,40 +1128,37 @@ func (s *Store) getBasePricing(model, provider string, requestType schemas.Reque } } - if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest { - if s.logger != nil { - s.logger.Debug("primary lookup failed, trying chat provider for the same model in chat completion") - } + // Lookup in chat if responses/compaction not found + if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest || requestType == schemas.CompactionRequest { + s.logger.Debug("primary lookup failed, trying chat provider for the same model in chat completion") pricing, ok = s.pricingData[makeKey(model, provider, normalizeRequestType(schemas.ChatCompletionRequest))] if ok { return &pricing, true } } + // Lookup in image generation if image edit not found if requestType == schemas.ImageEditRequest || requestType == schemas.ImageEditStreamRequest || requestType == schemas.ImageVariationRequest { - if s.logger != nil { - s.logger.Debug("primary lookup failed, trying image generation provider for the same model") - } + s.logger.Debug("primary lookup failed, trying image generation provider for the same model") pricing, ok = s.pricingData[makeKey(model, provider, normalizeRequestType(schemas.ImageGenerationRequest))] if ok { return &pricing, true } } + // Lookup fallback chain for container_create: + // 1. Try chat mode for the same model (e.g. "container-1g" in chat mode) + // 2. Try the base "container" model in chat mode (default rate when no memory-specific entry exists) if requestType == schemas.ContainerCreateRequest { - if s.logger != nil { - s.logger.Debug("primary lookup failed, trying chat mode for container create pricing") - } + s.logger.Debug("primary lookup failed, trying chat mode for container create pricing") pricing, ok = s.pricingData[makeKey(model, provider, normalizeRequestType(schemas.ChatCompletionRequest))] if ok { return &pricing, true } if model != "container" { - if s.logger != nil { - s.logger.Debug("memory-specific container pricing not found, falling back to base container entry") - } + s.logger.Debug("memory-specific container pricing not found, falling back to base container entry") pricing, ok = s.pricingData[makeKey("container", provider, normalizeRequestType(schemas.ChatCompletionRequest))] if ok { return &pricing, true @@ -1054,10 +1170,10 @@ func (s *Store) getBasePricing(model, provider string, requestType schemas.Reque } // UpsertModelPricingAttributes writes the additional_attributes column for -// every pricing row matching (model, provider), then reloads the pricing -// cache so the new values are immediately visible. Returns the number of -// rows updated (0 = no such pricing row, which callers must surface as a -// validation error). An empty/nil attrs map clears the column. +// every pricing row that matches (model, provider), then reloads the pricing +// cache so the new values are immediately visible to list-models. Returns +// the number of rows updated (0 = no such pricing row, which callers must +// surface as a validation error). An empty/nil attrs map clears the column. func (s *Store) UpsertModelPricingAttributes(ctx context.Context, model string, provider schemas.ModelProvider, attrs map[string]string) (int64, error) { if s.configStore == nil { return 0, fmt.Errorf("model catalog requires a config store") @@ -1079,6 +1195,7 @@ func (s *Store) UpsertModelPricingAttributes(ctx context.Context, model string, // Passthrough pricing helpers // --------------------------------------------------------------------------- +// detectPassthroughRequestType maps a provider + stripped path to a RequestType. func detectPassthroughRequestType(provider schemas.ModelProvider, path string) schemas.RequestType { if idx := strings.IndexByte(path, '?'); idx >= 0 { path = path[:idx] @@ -1093,6 +1210,8 @@ func detectPassthroughRequestType(provider schemas.ModelProvider, path string) s return schemas.TextCompletionRequest case strings.HasSuffix(path, "/embeddings"): return schemas.EmbeddingRequest + case strings.HasSuffix(path, "/responses/compact"): + return schemas.CompactionRequest case strings.HasSuffix(path, "/responses"): return schemas.ResponsesRequest case strings.HasSuffix(path, "/images/generations"): @@ -1114,6 +1233,7 @@ func detectPassthroughRequestType(provider schemas.ModelProvider, path string) s return schemas.ChatCompletionRequest } case schemas.Gemini, schemas.Vertex: + // Interactions API paths carry no colon action suffix. if strings.Contains(path, "/interactions") { return schemas.ResponsesRequest } @@ -1149,9 +1269,8 @@ func detectPassthroughRequestType(provider schemas.ModelProvider, path string) s } } -// inferPassthroughRequestType determines the request type from usage fields -// (primary), falling back to path detection for text/embedding/responses -// where LLMUsage is ambiguous. +// inferPassthroughRequestType determines the request type from usage fields (primary) +// and falls back to path detection for text/embedding/responses where LLMUsage is ambiguous. func inferPassthroughRequestType(provider schemas.ModelProvider, path string, su *schemas.BifrostPassthroughUsage) schemas.RequestType { if su != nil { if su.ContainerIdentifier != "" { @@ -1173,6 +1292,7 @@ func inferPassthroughRequestType(provider schemas.ModelProvider, path string, su return detectPassthroughRequestType(provider, path) } +// passthroughUsageToCostInput converts BifrostPassthroughUsage into costInput. func passthroughUsageToCostInput(su *schemas.BifrostPassthroughUsage) costInput { var input costInput if su.LLMUsage != nil { diff --git a/framework/modelcatalog/datasheet/params.go b/framework/modelcatalog/datasheet/params.go index d3cf03f64c3..86e990f7bee 100644 --- a/framework/modelcatalog/datasheet/params.go +++ b/framework/modelcatalog/datasheet/params.go @@ -6,8 +6,11 @@ import ( "fmt" "io" "net/http" + "net/url" + "os" "slices" + bifrost "github.com/maximhq/bifrost/core" providerUtils "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" @@ -113,22 +116,43 @@ func (s *Store) LoadModelParamsFromURLIntoMemory(ctx context.Context) error { // loadModelParametersFromURL fetches and parses the model parameters // datasheet at the configured URL. func (s *Store) loadModelParametersFromURL(ctx context.Context) (map[string]json.RawMessage, error) { - client := &http.Client{Timeout: DefaultModelParametersTimeout} - req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.ModelParametersURL(), nil) - if err != nil { - return nil, fmt.Errorf("failed to create HTTP request: %w", err) - } - resp, err := client.Do(req) + s.syncCfgMu.RLock() + rawURL := s.modelParametersURL + s.syncCfgMu.RUnlock() + + parsed, err := url.Parse(rawURL) if err != nil { - return nil, fmt.Errorf("failed to download model parameters data: %w", err) + return nil, fmt.Errorf("failed to parse model parameters URL: %w", err) } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("failed to download model parameters data: HTTP %d", resp.StatusCode) - } - data, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read model parameters response: %w", err) + + var data []byte + + if parsed.Scheme == "file" { + data, err = os.ReadFile(parsed.Path) + if err != nil { + return nil, fmt.Errorf("failed to read model parameters file: %w", err) + } + } else { + if err := bifrost.ValidateExternalURL(rawURL, true); err != nil { + return nil, fmt.Errorf("model parameters URL validation failed: %w", err) + } + client := &http.Client{Timeout: DefaultModelParametersTimeout} + req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP request: %w", err) + } + resp, err := client.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to download model parameters data: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("failed to download model parameters data: HTTP %d", resp.StatusCode) + } + data, err = io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read model parameters response: %w", err) + } } var paramsData map[string]json.RawMessage if err := json.Unmarshal(data, ¶msData); err != nil { diff --git a/framework/modelcatalog/datasheet/store.go b/framework/modelcatalog/datasheet/store.go index 826011baf81..4c3ee8b6dc7 100644 --- a/framework/modelcatalog/datasheet/store.go +++ b/framework/modelcatalog/datasheet/store.go @@ -6,6 +6,7 @@ import ( "sync" "time" + bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore" configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" @@ -400,6 +401,24 @@ func (s *Store) selectCapabilityEntryFromKeysUnsafe(matchingKeys []string) *Entr return convertTablePricingToEntry(&pricing) } +// NewTestStore constructs a minimal Store for unit tests without I/O. +// Optionally seed baseModelIndex so BaseModelName lookups resolve. A no-op +// logger is wired so cost / pricing paths (which assume Store.logger is +// non-nil) don't panic from external test code. +func NewTestStore(baseModelIndex map[string]string) *Store { + if baseModelIndex == nil { + baseModelIndex = make(map[string]string) + } + return &Store{ + logger: bifrost.NewNoOpLogger(), + pricingData: make(map[string]configstoreTables.TableModelPricing), + baseModelIndex: baseModelIndex, + supportedResponseTypes: make(map[string][]string), + supportedParams: make(map[string][]string), + datasheetByProvider: make(map[schemas.ModelProvider][]string), + } +} + // --- Internal: rebuild the datasheet view from current pricingData --- // rebuildDatasheetViewUnsafe regenerates baseModelIndex and datasheetByProvider diff --git a/framework/modelcatalog/datasheet/sync.go b/framework/modelcatalog/datasheet/sync.go index 162a6efc6ff..4d3b9749c6c 100644 --- a/framework/modelcatalog/datasheet/sync.go +++ b/framework/modelcatalog/datasheet/sync.go @@ -6,8 +6,11 @@ import ( "fmt" "io" "net/http" + "net/url" + "os" "time" + bifrost "github.com/maximhq/bifrost/core" providerUtils "github.com/maximhq/bifrost/core/providers/utils" configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" "gorm.io/gorm" @@ -145,22 +148,43 @@ func (s *Store) applyPricingData(pricingData map[string]Entry) { // loadPricingFromURL fetches and parses the pricing datasheet at the // configured URL. Honors ctx for cancellation. func (s *Store) loadPricingFromURL(ctx context.Context) (map[string]Entry, error) { - client := &http.Client{Timeout: DefaultPricingTimeout} - req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.URL(), nil) - if err != nil { - return nil, fmt.Errorf("failed to create HTTP request: %w", err) - } - resp, err := client.Do(req) + s.syncCfgMu.RLock() + rawURL := s.url + s.syncCfgMu.RUnlock() + + parsed, err := url.Parse(rawURL) if err != nil { - return nil, fmt.Errorf("failed to download pricing data: %w", err) + return nil, fmt.Errorf("failed to parse pricing URL: %w", err) } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("failed to download pricing data: HTTP %d", resp.StatusCode) - } - data, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read pricing data response: %w", err) + + var data []byte + + if parsed.Scheme == "file" { + data, err = os.ReadFile(parsed.Path) + if err != nil { + return nil, fmt.Errorf("failed to read pricing file: %w", err) + } + } else { + if err := bifrost.ValidateExternalURL(rawURL, true); err != nil { + return nil, fmt.Errorf("pricing URL validation failed: %w", err) + } + client := &http.Client{Timeout: DefaultPricingTimeout} + req, err := http.NewRequestWithContext(ctx, http.MethodGet, s.URL(), nil) + if err != nil { + return nil, fmt.Errorf("failed to create HTTP request: %w", err) + } + resp, err := client.Do(req) + if err != nil { + return nil, fmt.Errorf("failed to download pricing data: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("failed to download pricing data: HTTP %d", resp.StatusCode) + } + data, err = io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("failed to read pricing data response: %w", err) + } } var pricingData map[string]Entry if err := json.Unmarshal(data, &pricingData); err != nil { @@ -169,6 +193,7 @@ func (s *Store) loadPricingFromURL(ctx context.Context) (map[string]Entry, error if s.logger != nil { s.logger.Debug("successfully downloaded and parsed %d pricing records", len(pricingData)) } + return pricingData, nil } diff --git a/framework/modelcatalog/datasheet/types.go b/framework/modelcatalog/datasheet/types.go index e025a6a6302..fb4e20757fe 100644 --- a/framework/modelcatalog/datasheet/types.go +++ b/framework/modelcatalog/datasheet/types.go @@ -341,7 +341,7 @@ func normalizeRequestType(reqType schemas.RequestType) string { return "completion" case schemas.ChatCompletionRequest, schemas.ChatCompletionStreamRequest: return "chat" - case schemas.ResponsesRequest, schemas.ResponsesStreamRequest, schemas.WebSocketResponsesRequest, schemas.RealtimeRequest: + case schemas.ResponsesRequest, schemas.ResponsesStreamRequest, schemas.WebSocketResponsesRequest, schemas.RealtimeRequest, schemas.CompactionRequest: return "responses" case schemas.EmbeddingRequest: return "embedding" diff --git a/framework/modelcatalog/main.go b/framework/modelcatalog/main.go index d8bb9446ec3..b7c4effc482 100644 --- a/framework/modelcatalog/main.go +++ b/framework/modelcatalog/main.go @@ -1,113 +1,87 @@ -// Package modelcatalog provides a pricing manager for the framework. +// Package modelcatalog composes three subpackages — datasheet (pricing + +// model parameters + capabilities), live (per-(provider, keyID) list-models +// cache), and keyconfig (per-provider allow/block/aliases derived from +// keys) — into the ModelCatalog facade that consumers (governance, +// telemetry, logging, server, etc.) use. +// +// The composer owns I/O orchestration: the hourly pricing sync ticker, the +// distributed lock used during sync, and the gossip after-sync hook. +// Subpackages perform no I/O directly — they expose Load/Sync methods the +// composer calls. package modelcatalog import ( "context" "encoding/json" "fmt" - "slices" "sync" "time" providerUtils "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore" - configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/maximhq/bifrost/framework/modelcatalog/datasheet" + "github.com/maximhq/bifrost/framework/modelcatalog/keyconfig" + "github.com/maximhq/bifrost/framework/modelcatalog/live" ) type ModelCatalog struct { configStore configstore.ConfigStore distributedLockManager *configstore.DistributedLockManager + logger schemas.Logger - logger schemas.Logger - - // Configuration fields (protected by syncMu) - pricingURL string - modelParametersURL string - syncInterval time.Duration - lastSyncedAt time.Time - syncMu sync.RWMutex + datasheet *datasheet.Store + live *live.Store + keyconf *keyconfig.Store shouldSyncGate func(ctx context.Context) bool afterSyncHook func(ctx context.Context) - // In-memory cache for fast access - direct map for O(1) lookups - pricingData map[string]configstoreTables.TableModelPricing - mu sync.RWMutex - - // rawOverrides is the canonical list of all active overrides. It exists solely - // to support incremental mutations: UpsertPricingOverrides and DeletePricingOverride - // iterate over it to rebuild the list, then derive customPricing from it. - // customPricing is the actual lookup structure used at query time. - rawOverrides []PricingOverride - customPricing *customPricingData - overridesMu sync.RWMutex - - modelPool map[schemas.ModelProvider][]string - unfilteredModelPool map[schemas.ModelProvider][]string // model pool without allowed models filtering - baseModelIndex map[string]string // model string → canonical base model name - - // Pre-parsed supported response types index (keyed by model name) - // Values are normalized response types: "chat_completion", "responses", "text_completion" - supportedResponseTypes map[string][]string - - // Pre-parsed supported parameters index (keyed by model name, populated from model parameters supported_parameters) - // Values are parameter names the model accepts (e.g., "temperature", "top_p", "tools") - supportedParams map[string][]string - - // Background sync worker + // Background sync orchestration. The ticker, distributed lock, and gossip + // hook live at this level — datasheet.Store has no internal scheduler. syncTicker *time.Ticker - done chan struct{} - wg sync.WaitGroup syncCtx context.Context syncCancel context.CancelFunc + done chan struct{} + wg sync.WaitGroup } -// Init initializes the model catalog func Init(ctx context.Context, config *Config, configStore configstore.ConfigStore, logger schemas.Logger) (*ModelCatalog, error) { - // Initialize pricing URL and sync interval pricingURL := DefaultPricingURL - if config.PricingURL != nil { + if config != nil && config.PricingURL != nil { pricingURL = *config.PricingURL } modelParametersURL := DefaultModelParametersURL - if config.ModelParametersURL != nil && *config.ModelParametersURL != "" { + if config != nil && config.ModelParametersURL != nil && *config.ModelParametersURL != "" { modelParametersURL = *config.ModelParametersURL } syncInterval := DefaultSyncInterval - if config.PricingSyncInterval != nil { + if config != nil && config.PricingSyncInterval != nil { syncInterval = time.Duration(*config.PricingSyncInterval) * time.Second } // Log the active interval and the scheduler's actual check frequency so operators // are not surprised that setting interval=1h does not mean checks happen every second. - // Actual syncs occur when: (1) the 1-hour ticker fires AND (2) time.Since(lastSync) >= pricingSyncInterval. logger.Info("pricing sync interval set to %v (scheduler checks every %v)", syncInterval, syncWorkerTickerPeriod) mc := &ModelCatalog{ - pricingURL: pricingURL, - modelParametersURL: modelParametersURL, - syncInterval: syncInterval, configStore: configStore, logger: logger, - pricingData: make(map[string]configstoreTables.TableModelPricing), - modelPool: make(map[schemas.ModelProvider][]string), - unfilteredModelPool: make(map[schemas.ModelProvider][]string), - baseModelIndex: make(map[string]string), - supportedResponseTypes: make(map[string][]string), - supportedParams: make(map[string][]string), - done: make(chan struct{}), distributedLockManager: configstore.NewDistributedLockManager(configStore, logger, configstore.WithDefaultTTL(30*time.Second)), + datasheet: datasheet.New(configStore, logger, datasheet.Config{ + URL: pricingURL, + ModelParametersURL: modelParametersURL, + SyncInterval: syncInterval, + }), + live: live.New(logger), + keyconf: keyconfig.New(logger), + done: make(chan struct{}), } - - // Initialize syncCtx early so background startup goroutines can use it and - // Cleanup() can cancel them. startSyncWorker is still called at the end after - // cold-start paths have completed. mc.syncCtx, mc.syncCancel = context.WithCancel(ctx) // If Init returns an error the caller never owns mc and will never call - // Cleanup(), so cancel syncCtx to stop any background goroutines that were - // already spawned before the failure. + // Cleanup(), so cancel syncCtx to stop any background goroutines that + // were already spawned before the failure. initSucceeded := false defer func() { if !initSucceeded { @@ -117,13 +91,13 @@ func Init(ctx context.Context, config *Config, configStore configstore.ConfigSto logger.Info("initializing model catalog...") if configStore != nil { - // Per-model lazy load when the in-memory cache misses (eviction, new models, or if - // startup bulk load was skipped). loadModelParametersFromDatabase still bulk-warms - // the cache on init and on ReloadFromDB so common paths avoid a DB read per model. + // Lazy load on cache miss: providers may need params for models not + // covered by the startup bulk load (e.g. just-uploaded models). The + // bulk load still warms the common case so this only fires on misses. providerUtils.SetCacheMissHandler(func(model string) *providerUtils.ModelParams { missCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() - params, err := configStore.GetModelParametersByModel(missCtx, model) + params, err := mc.datasheet.GetModelParametersByModel(missCtx, model) if err != nil || params == nil { return nil } @@ -142,34 +116,32 @@ func Init(ctx context.Context, config *Config, configStore configstore.ConfigSto IsVertexMultiRegionOnly: p.VertexMultiRegionOnly, } }) + var wg sync.WaitGroup var pricingErr, paramsErr error wg.Add(2) go func() { defer wg.Done() - if err := mc.loadPricingFromDatabase(ctx); err != nil { + if err := mc.datasheet.LoadFromDB(ctx); err != nil { pricingErr = fmt.Errorf("failed to load initial pricing data: %w", err) return } - mc.mu.RLock() - hasPricingData := len(mc.pricingData) > 0 - mc.mu.RUnlock() - if hasPricingData { - mc.logger.Info("existing pricing data found in database, syncing from URL in background") + if mc.hasPricingData() { + logger.Info("existing pricing data found in database, syncing from URL in background") mc.wg.Add(1) go func() { defer mc.wg.Done() if err := mc.withDistributedLock(mc.syncCtx, "model_catalog_pricing_startup_sync", 10, func() error { - return mc.syncPricing(mc.syncCtx) + return mc.runPricingSync(mc.syncCtx) }); err != nil { - mc.logger.Warn("background startup pricing sync failed: %v", err) + logger.Warn("background startup pricing sync failed: %v", err) } else { - mc.logger.Info("background startup pricing sync completed successfully") + logger.Info("background startup pricing sync completed successfully") } }() } else { if err := mc.withDistributedLock(ctx, "model_catalog_pricing_startup_sync", 10, func() error { - return mc.syncPricing(ctx) + return mc.runPricingSync(ctx) }); err != nil { pricingErr = fmt.Errorf("failed to sync pricing data: %w", err) } @@ -177,27 +149,27 @@ func Init(ctx context.Context, config *Config, configStore configstore.ConfigSto }() go func() { defer wg.Done() - n, err := mc.loadModelParametersFromDatabase(ctx) + n, err := mc.datasheet.LoadModelParamsFromDB(ctx) if err != nil { paramsErr = fmt.Errorf("failed to load initial model parameters: %w", err) return } if n > 0 { - mc.logger.Info("existing model parameters found in database (%d records), syncing from URL in background", n) + logger.Info("existing model parameters found in database (%d records), syncing from URL in background", n) mc.wg.Add(1) go func() { defer mc.wg.Done() if err := mc.withDistributedLock(mc.syncCtx, "model_catalog_params_startup_sync", 10, func() error { - return mc.syncModelParameters(mc.syncCtx) + return mc.runParamsSync(mc.syncCtx) }); err != nil { - mc.logger.Warn("background startup model parameters sync failed: %v", err) + logger.Warn("background startup model parameters sync failed: %v", err) } else { - mc.logger.Info("background startup model parameters sync completed successfully") + logger.Info("background startup model parameters sync completed successfully") } }() } else { if err := mc.withDistributedLock(ctx, "model_catalog_params_startup_sync", 10, func() error { - return mc.syncModelParameters(ctx) + return mc.runParamsSync(ctx) }); err != nil { paramsErr = fmt.Errorf("failed to sync model parameters data: %w", err) } @@ -211,60 +183,57 @@ func Init(ctx context.Context, config *Config, configStore configstore.ConfigSto return nil, paramsErr } } else { - // Load pricing and model parameters from URL into memory (no config store) - if err := mc.loadPricingIntoMemoryFromURL(ctx); err != nil { - return nil, fmt.Errorf("failed to load pricing data from config memory: %w", err) + if err := mc.datasheet.LoadFromURLIntoMemory(ctx); err != nil { + return nil, fmt.Errorf("failed to load pricing data into memory: %w", err) } - if err := mc.loadModelParametersIntoMemoryFromURL(ctx); err != nil { + if err := mc.datasheet.LoadModelParamsFromURLIntoMemory(ctx); err != nil { return nil, fmt.Errorf("failed to load model parameters from URL: %w", err) } } - mc.syncMu.Lock() - mc.lastSyncedAt = time.Now() - mc.syncMu.Unlock() - - // Populate model pool with normalized providers from pricing data - mc.populateModelPoolFromPricingData() + mc.datasheet.MarkSynced(time.Now()) - if err := mc.loadPricingOverridesFromStore(ctx); err != nil { + if err := mc.datasheet.LoadOverridesFromStore(ctx); err != nil { return nil, fmt.Errorf("failed to load pricing overrides: %w", err) } - // Start background sync worker mc.startSyncWorker(mc.syncCtx) initSucceeded = true return mc, nil } -func (mc *ModelCatalog) SetShouldSyncGate(shouldSyncGate func(ctx context.Context) bool) { - mc.shouldSyncGate = shouldSyncGate +func (mc *ModelCatalog) SetShouldSyncGate(fn func(ctx context.Context) bool) { + mc.shouldSyncGate = fn } -// SetAfterSyncHook registers a callback invoked after every successful URL → DB pricing sync. -// In enterprise this is used to broadcast a gossip message so other pods reload from DB. +// SetAfterSyncHook registers a callback invoked after every successful +// URL → DB pricing sync. In enterprise this broadcasts a gossip message so +// other pods reload from DB. func (mc *ModelCatalog) SetAfterSyncHook(fn func(ctx context.Context)) { mc.afterSyncHook = fn } -// ReloadFromDB reloads the in-memory pricing cache and model-parameters provider cache from the database. -// In enterprise this is called on non-leader pods when they receive a gossip sync notification. +// ReloadFromDB reloads pricing + model-parameters caches from the database. +// Gossip handler on non-leader pods. func (mc *ModelCatalog) ReloadFromDB(ctx context.Context) error { - if err := mc.loadPricingFromDatabase(ctx); err != nil { + if err := mc.datasheet.LoadFromDB(ctx); err != nil { return err } - mc.populateModelPoolFromPricingData() - _, err := mc.loadModelParametersFromDatabase(ctx) + _, err := mc.datasheet.LoadModelParamsFromDB(ctx) return err } -// UpdateSyncConfig updates the pricing URL and sync interval, restarts the background sync worker, -// then delegates to ForceReloadPricing for a full sync cycle. -func (mc *ModelCatalog) UpdateSyncConfig(ctx context.Context, config *Config) error { - // Acquire pricing mutex to update configuration atomically - mc.syncMu.Lock() +// ReloadPricing re-reads the pricing table into the in-memory cache. The +// management API uses this after a batched write so the new attributes are +// observable immediately. The 24-hour ticker still owns refreshing pricing +// fields from the upstream datasheet; this just refreshes the cache. +func (mc *ModelCatalog) ReloadPricing(ctx context.Context) error { + return mc.datasheet.LoadFromDB(ctx) +} - // Stop existing sync worker before updating configuration +// UpdateSyncConfig updates the pricing/params URLs and sync interval, +// restarts the background sync worker, then runs a full sync cycle. +func (mc *ModelCatalog) UpdateSyncConfig(ctx context.Context, config *Config) error { if mc.syncCancel != nil { mc.syncCancel() } @@ -272,70 +241,64 @@ func (mc *ModelCatalog) UpdateSyncConfig(ctx context.Context, config *Config) er mc.syncTicker.Stop() } - // Update pricing configuration - mc.pricingURL = DefaultPricingURL - if config.PricingURL != nil { - mc.pricingURL = *config.PricingURL + pricingURL := DefaultPricingURL + if config != nil && config.PricingURL != nil { + pricingURL = *config.PricingURL } - - mc.modelParametersURL = DefaultModelParametersURL - if config.ModelParametersURL != nil && *config.ModelParametersURL != "" { - mc.modelParametersURL = *config.ModelParametersURL + modelParametersURL := DefaultModelParametersURL + if config != nil && config.ModelParametersURL != nil && *config.ModelParametersURL != "" { + modelParametersURL = *config.ModelParametersURL } - - mc.syncInterval = DefaultSyncInterval - if config.PricingSyncInterval != nil { - mc.syncInterval = time.Duration(*config.PricingSyncInterval) * time.Second + syncInterval := DefaultSyncInterval + if config != nil && config.PricingSyncInterval != nil { + syncInterval = time.Duration(*config.PricingSyncInterval) * time.Second } + mc.datasheet.UpdateSyncConfig(datasheet.Config{ + URL: pricingURL, + ModelParametersURL: modelParametersURL, + SyncInterval: syncInterval, + }) - // Create new sync worker with updated configuration mc.syncCtx, mc.syncCancel = context.WithCancel(ctx) mc.startSyncWorker(mc.syncCtx) - mc.syncMu.Unlock() - - // Delegate to ForceReloadPricing for a complete sync cycle return mc.ForceReloadPricing(ctx) } +// ForceReloadPricing triggers an immediate URL→DB→memory sync for pricing +// and model parameters in parallel, fires the gossip hook, and resets the +// ticker so the next scheduled sync waits a full interval from now. +// +// Behavior change from pre-refactor: this no longer touches the live +// list-models cache. List-models refresh is now driven by key/provider +// edits, not by pricing reloads. func (mc *ModelCatalog) ForceReloadPricing(ctx context.Context) error { - timeout := DefaultPricingTimeout + timeout := datasheet.DefaultPricingTimeout if timeout > 0 { var cancel context.CancelFunc ctx, cancel = context.WithTimeout(ctx, timeout) defer cancel() } - // Run pricing sync and model parameters sync in parallel var wg sync.WaitGroup var pricingErr, paramsErr error - - wg.Add(1) + wg.Add(2) go func() { defer wg.Done() - if err := mc.syncPricing(ctx); err != nil { + if err := mc.runPricingSync(ctx); err != nil { pricingErr = fmt.Errorf("failed to sync pricing data: %w", err) return } - - // Rebuild model pool from updated pricing data - mc.populateModelPoolFromPricingData() - - if err := mc.loadPricingOverridesFromStore(ctx); err != nil { + if err := mc.datasheet.LoadOverridesFromStore(ctx); err != nil { pricingErr = fmt.Errorf("failed to load pricing overrides: %w", err) - return } }() - - wg.Add(1) go func() { defer wg.Done() - if err := mc.syncModelParameters(ctx); err != nil { + if err := mc.runParamsSync(ctx); err != nil { paramsErr = fmt.Errorf("failed to sync model parameters: %w", err) - return } }() - wg.Wait() if pricingErr != nil { return pricingErr @@ -348,182 +311,182 @@ func (mc *ModelCatalog) ForceReloadPricing(ctx context.Context) error { mc.afterSyncHook(ctx) } - mc.syncMu.Lock() - // Reset the ticker so the next scheduled sync waits a full interval from now if mc.syncTicker != nil { - mc.syncTicker.Reset(mc.syncInterval) + mc.syncTicker.Reset(mc.datasheet.SyncInterval()) } - mc.syncMu.Unlock() + return nil +} +func (mc *ModelCatalog) Cleanup() error { + if mc.syncCancel != nil { + mc.syncCancel() + } + if mc.syncTicker != nil { + mc.syncTicker.Stop() + } + close(mc.done) + mc.wg.Wait() return nil } -// getPricingURL returns a copy of the pricing URL under mutex protection -func (mc *ModelCatalog) getPricingURL() string { - mc.syncMu.RLock() - defer mc.syncMu.RUnlock() - return mc.pricingURL +// --- Sync ticker (orchestrates datasheet.Store sync methods) --- + +func (mc *ModelCatalog) startSyncWorker(ctx context.Context) { + // IMPORTANT scheduling model: + // + // The sync worker wakes on a fixed ticker (syncWorkerTickerPeriod = 1h). + // On each wake it checks time.Since(LastSyncedAt) >= SyncInterval. + // This means SyncInterval defines the *minimum elapsed time* between syncs, + // and the actual frequency = max(syncWorkerTickerPeriod, SyncInterval). + // Setting SyncInterval below the ticker period has no effect — the hourly + // ticker is the hard lower bound on check granularity. + mc.syncTicker = time.NewTicker(syncWorkerTickerPeriod) + mc.wg.Add(1) + go mc.syncWorker(ctx) } -func (mc *ModelCatalog) getModelParametersURL() string { - mc.syncMu.RLock() - defer mc.syncMu.RUnlock() - return mc.modelParametersURL +func (mc *ModelCatalog) syncWorker(ctx context.Context) { + // Capture the ticker once so the select loop doesn't race with + // UpdateSyncConfig overwriting mc.syncTicker while this goroutine + // is still draining after mc.syncCancel(). + ticker := mc.syncTicker + defer mc.wg.Done() + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + mc.syncTick(ctx) + case <-mc.done: + return + } + } } -// IsRequestTypeSupported checks if a model supports chat completion. -// It checks the supportedResponseTypes index. -func (mc *ModelCatalog) IsRequestTypeSupported(model string, provider schemas.ModelProvider, requestType schemas.RequestType) bool { - mc.mu.RLock() - defer mc.mu.RUnlock() - outputs, ok := mc.supportedResponseTypes[model] - return ok && slices.Contains(outputs, string(requestType)) +func (mc *ModelCatalog) syncTick(ctx context.Context) { + if time.Since(mc.datasheet.LastSyncedAt()) < mc.datasheet.SyncInterval() { + return + } + mc.logger.Debug("starting model catalog background sync") + if err := mc.withDistributedLock(ctx, "model_catalog_pricing_sync", 10, func() error { + var wg sync.WaitGroup + var pricingErr, paramsErr error + wg.Add(2) + go func() { + defer wg.Done() + if err := mc.runPricingSync(ctx); err != nil { + mc.logger.Error("background pricing sync failed: %v", err) + pricingErr = err + } + }() + go func() { + defer wg.Done() + if err := mc.runParamsSync(ctx); err != nil { + mc.logger.Error("background model parameters sync failed: %v", err) + paramsErr = err + } + }() + wg.Wait() + if pricingErr == nil && paramsErr == nil { + if mc.afterSyncHook != nil { + mc.afterSyncHook(ctx) + } + mc.datasheet.MarkSynced(time.Now()) + } + if pricingErr != nil { + return pricingErr + } + return paramsErr + }); err != nil { + mc.logger.Error("failed to run model catalog sync: %v", err) + } + mc.logger.Debug("model catalog background sync completed") } -// GetSupportedParameters returns the list of supported parameter names for a model. -// Returns nil if the model is not found in the catalog. -func (mc *ModelCatalog) GetSupportedParameters(model string) []string { - mc.mu.RLock() - params, ok := mc.supportedParams[model] - mc.mu.RUnlock() - if !ok { +// runPricingSync wraps the datasheet pricing sync with the gate check. +func (mc *ModelCatalog) runPricingSync(ctx context.Context) error { + if mc.shouldSyncGate != nil && !mc.shouldSyncGate(ctx) { return nil } - // Return a copy to prevent external modification - result := make([]string, len(params)) - copy(result, params) - return result + return mc.datasheet.SyncFromURL(ctx) } -// populateModelPool populates the model pool with all available models per provider (thread-safe). -// -// This function is the only path that resets modelPool / unfilteredModelPool / -// baseModelIndex from upstream pricing data. It is called both at init (where -// the pool is empty) and on every reload (gossip ReloadFromDB, manual -// ForceReloadPricing). To avoid drift on reload — where a naive wipe would -// drop everything contributed by per-provider list-models output and key -// allowed_models — the pre-wipe pool is snapshotted and unioned back in after -// the pricing rebuild. baseModelIndex is intentionally not preserved: aliases -// outside the pricing sheet have no canonical base-model entry, and -// getBaseModelNameUnsafe falls through to algorithmic stripping for them. -func (mc *ModelCatalog) populateModelPoolFromPricingData() { - // Acquire write lock for the entire rebuild operation - mc.mu.Lock() - defer mc.mu.Unlock() - - // Snapshot the pre-wipe pool so non-pricing contributions (list-models - // output, allowed_models) survive the rebuild. - previousModelPool := make(map[schemas.ModelProvider][]string, len(mc.modelPool)) - for provider, models := range mc.modelPool { - copied := make([]string, len(models)) - copy(copied, models) - previousModelPool[provider] = copied - } - previousUnfilteredModelPool := make(map[schemas.ModelProvider][]string, len(mc.unfilteredModelPool)) - for provider, models := range mc.unfilteredModelPool { - copied := make([]string, len(models)) - copy(copied, models) - previousUnfilteredModelPool[provider] = copied +// runParamsSync wraps the datasheet params sync with the gate check. +func (mc *ModelCatalog) runParamsSync(ctx context.Context) error { + if mc.shouldSyncGate != nil && !mc.shouldSyncGate(ctx) { + mc.logger.Debug("model parameters sync cancelled by custom gate") + return nil } + return mc.datasheet.SyncModelParamsFromURL(ctx) +} - // Clear existing model pool and base model index - mc.modelPool = make(map[schemas.ModelProvider][]string) - mc.unfilteredModelPool = make(map[schemas.ModelProvider][]string) - mc.baseModelIndex = make(map[string]string) - - // Map to track unique models per provider - providerModels := make(map[schemas.ModelProvider]map[string]bool) - - // Iterate through all pricing data to collect models per provider - for _, pricing := range mc.pricingData { - // Normalize provider before adding to model pool - normalizedProvider := schemas.ModelProvider(normalizeProvider(pricing.Provider)) - - // Initialize map for this provider if not exists - if providerModels[normalizedProvider] == nil { - providerModels[normalizedProvider] = make(map[string]bool) - } - - // Add model to the provider's model set (using map for deduplication) - providerModels[normalizedProvider][pricing.Model] = true - - // Build base model index from pre-computed base_model field - if pricing.BaseModel != "" { - mc.baseModelIndex[pricing.Model] = pricing.BaseModel - } +// withDistributedLock acquires a named distributed lock and runs fn under +// it. retries=0 blocks until acquired; retries>0 uses LockWithRetry. The +// unlock uses a fresh context so cancelled work contexts don't leak the +// lock until TTL expiry. +func (mc *ModelCatalog) withDistributedLock(ctx context.Context, key string, retries int, fn func() error) error { + lock, err := mc.distributedLockManager.NewLock(key) + if err != nil { + return fmt.Errorf("failed to create lock %q: %w", key, err) } - - // Convert sets to slices and assign to modelPool - for provider, modelSet := range providerModels { - models := make([]string, 0, len(modelSet)) - for model := range modelSet { - models = append(models, model) + if retries > 0 { + if err := lock.LockWithRetry(ctx, retries); err != nil { + return fmt.Errorf("failed to acquire lock %q: %w", key, err) } - mc.modelPool[provider] = models - mc.unfilteredModelPool[provider] = models - } - - // Union the pre-wipe snapshot back in. Anything previously added by - // UpsertModelDataForProvider / UpsertUnfilteredModelDataForProvider — - // list-models output, allowed_models aliases — is restored. Pricing - // entries from the rebuild win on duplicates (already in place above); - // removals via DeleteModelDataForProvider are respected because that - // method strips the provider from the live map before this runs. - for provider, models := range previousModelPool { - for _, m := range models { - if !slices.Contains(mc.modelPool[provider], m) { - mc.modelPool[provider] = append(mc.modelPool[provider], m) - } + } else { + if err := lock.Lock(ctx); err != nil { + return fmt.Errorf("failed to acquire lock %q: %w", key, err) } } - for provider, models := range previousUnfilteredModelPool { - for _, m := range models { - if !slices.Contains(mc.unfilteredModelPool[provider], m) { - mc.unfilteredModelPool[provider] = append(mc.unfilteredModelPool[provider], m) - } + defer func() { + if err := lock.Unlock(context.Background()); err != nil { + mc.logger.Warn("failed to release distributed lock %q: %v", key, err) } - } + }() + return fn() +} - // Log the populated model pool for debugging - totalModels := 0 - for provider, models := range mc.modelPool { - totalModels += len(models) - mc.logger.Debug("populated %d models for provider %s", len(models), string(provider)) - } - mc.logger.Info("populated model pool with %d models across %d providers", totalModels, len(mc.modelPool)) +// hasPricingData reports whether the datasheet store currently has any +// pricing rows in memory. Used during Init to decide between blocking sync +// and background sync. +func (mc *ModelCatalog) hasPricingData() bool { + return len(mc.datasheet.DatasheetProviders()) > 0 } -// Cleanup cleans up the model catalog -func (mc *ModelCatalog) Cleanup() error { - if mc.syncCancel != nil { - mc.syncCancel() +// knownProviders returns the union of providers seen by any store. Used by +// GetProvidersForModel (models.go) to enumerate candidates. +func (mc *ModelCatalog) knownProviders() []schemas.ModelProvider { + seen := make(map[schemas.ModelProvider]struct{}) + out := make([]schemas.ModelProvider, 0) + for _, p := range mc.datasheet.DatasheetProviders() { + if _, ok := seen[p]; !ok { + seen[p] = struct{}{} + out = append(out, p) + } } - - mc.syncMu.Lock() - if mc.syncTicker != nil { - mc.syncTicker.Stop() + for k := range mc.live.Snapshot() { + if _, ok := seen[k.Provider]; !ok { + seen[k.Provider] = struct{}{} + out = append(out, k.Provider) + } } - mc.syncMu.Unlock() - - close(mc.done) - mc.wg.Wait() - - return nil + for _, p := range mc.keyconf.Providers() { + if _, ok := seen[p]; !ok { + seen[p] = struct{}{} + out = append(out, p) + } + } + return out } -// NewTestCatalog creates a minimal ModelCatalog for testing purposes. -// It does not start background sync workers or connect to external services. +// NewTestCatalog constructs a minimal ModelCatalog for unit tests. Does not +// start background workers or hit external services. func NewTestCatalog(baseModelIndex map[string]string) *ModelCatalog { - if baseModelIndex == nil { - baseModelIndex = make(map[string]string) - } return &ModelCatalog{ - modelPool: make(map[schemas.ModelProvider][]string), - unfilteredModelPool: make(map[schemas.ModelProvider][]string), - baseModelIndex: baseModelIndex, - pricingData: make(map[string]configstoreTables.TableModelPricing), - supportedResponseTypes: make(map[string][]string), - supportedParams: make(map[string][]string), - done: make(chan struct{}), + datasheet: datasheet.NewTestStore(baseModelIndex), + live: live.New(nil), + keyconf: keyconfig.New(nil), + done: make(chan struct{}), } } diff --git a/framework/modelcatalog/main_test.go b/framework/modelcatalog/main_test.go deleted file mode 100644 index 3f1120ab09d..00000000000 --- a/framework/modelcatalog/main_test.go +++ /dev/null @@ -1,209 +0,0 @@ -package modelcatalog - -import ( - "testing" - - "github.com/maximhq/bifrost/core/schemas" - "github.com/maximhq/bifrost/framework/configstore" - configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" - "github.com/stretchr/testify/assert" -) - -// newTestCatalog creates a minimal ModelCatalog for testing within the package. -func newTestCatalog(modelPool map[schemas.ModelProvider][]string, baseModelIndex map[string]string) *ModelCatalog { - if modelPool == nil { - modelPool = make(map[schemas.ModelProvider][]string) - } - if baseModelIndex == nil { - baseModelIndex = make(map[string]string) - } - return &ModelCatalog{ - modelPool: modelPool, - baseModelIndex: baseModelIndex, - pricingData: make(map[string]configstoreTables.TableModelPricing), - } -} - -// --- GetBaseModelName tests --- - -func TestGetBaseModelName_Simple(t *testing.T) { - mc := newTestCatalog(nil, nil) - // No catalog data, no prefix — returns as-is (no date suffix to strip either) - assert.Equal(t, "gpt-4o", mc.GetBaseModelName("gpt-4o")) -} - -func TestGetBaseModelName_Prefixed(t *testing.T) { - mc := newTestCatalog(nil, nil) - // Provider prefix stripped, no catalog — algorithmic fallback returns base - assert.Equal(t, "gpt-4o", mc.GetBaseModelName("openai/gpt-4o")) -} - -func TestGetBaseModelName_PrefixedAnthropic(t *testing.T) { - mc := newTestCatalog(nil, nil) - assert.Equal(t, "claude-3-5-sonnet", mc.GetBaseModelName("anthropic/claude-3-5-sonnet")) -} - -func TestGetBaseModelName_FromCatalog(t *testing.T) { - // Model has a pre-computed base_model in the catalog - mc := newTestCatalog(nil, map[string]string{ - "gpt-4o": "gpt-4o", - "gpt-4o-2024-08-06": "gpt-4o", - }) - assert.Equal(t, "gpt-4o", mc.GetBaseModelName("gpt-4o")) - assert.Equal(t, "gpt-4o", mc.GetBaseModelName("gpt-4o-2024-08-06")) -} - -func TestGetBaseModelName_ProviderPrefixWithCatalog(t *testing.T) { - // Model has provider prefix — strip prefix, then find in catalog - mc := newTestCatalog(nil, map[string]string{ - "gpt-4o": "gpt-4o", - }) - assert.Equal(t, "gpt-4o", mc.GetBaseModelName("openai/gpt-4o")) -} - -func TestGetBaseModelName_FallbackAlgorithmic(t *testing.T) { - // Model NOT in catalog — falls back to schemas.BaseModelName (date stripping) - mc := newTestCatalog(nil, nil) - // Anthropic-style date suffix - assert.Equal(t, "claude-sonnet-4", mc.GetBaseModelName("claude-sonnet-4-20250514")) - // OpenAI-style date suffix - assert.Equal(t, "gpt-4o", mc.GetBaseModelName("gpt-4o-2024-08-06")) -} - -func TestGetBaseModelName_FallbackAlgorithmicWithPrefix(t *testing.T) { - // Provider prefix + not in catalog — strip prefix, then algorithmic fallback - mc := newTestCatalog(nil, nil) - assert.Equal(t, "claude-sonnet-4", mc.GetBaseModelName("anthropic/claude-sonnet-4-20250514")) -} - -func TestGetBaseModelName_UnknownModel(t *testing.T) { - mc := newTestCatalog(nil, nil) - assert.Equal(t, "some-random-model", mc.GetBaseModelName("some-random-model")) -} - -func TestGetBaseModelName_CatalogTakesPrecedence(t *testing.T) { - // If catalog says the base_model is X, use it even if algorithmic would give Y - mc := newTestCatalog(nil, map[string]string{ - "my-custom-model-20250101": "my-custom-model-20250101", // catalog says keep the date - }) - assert.Equal(t, "my-custom-model-20250101", mc.GetBaseModelName("my-custom-model-20250101")) -} - -// --- IsSameModel tests --- - -func TestIsSameModel_DirectMatch(t *testing.T) { - mc := newTestCatalog(nil, nil) - assert.True(t, mc.IsSameModel("gpt-4o", "gpt-4o")) -} - -func TestIsSameModel_ProviderPrefix(t *testing.T) { - mc := newTestCatalog(nil, nil) - assert.True(t, mc.IsSameModel("openai/gpt-4o", "gpt-4o")) - assert.True(t, mc.IsSameModel("gpt-4o", "openai/gpt-4o")) -} - -func TestIsSameModel_BothPrefixed(t *testing.T) { - mc := newTestCatalog(nil, nil) - assert.True(t, mc.IsSameModel("openai/gpt-4o", "openai/gpt-4o")) -} - -func TestIsSameModel_DifferentProvidersSameBase(t *testing.T) { - mc := newTestCatalog(nil, nil) - // Both have the same base model after stripping different provider prefixes - assert.True(t, mc.IsSameModel("openai/gpt-4o", "azure/gpt-4o")) -} - -func TestIsSameModel_DifferentModels(t *testing.T) { - mc := newTestCatalog(nil, nil) - assert.False(t, mc.IsSameModel("gpt-4o", "claude-3-5-sonnet")) -} - -func TestIsSameModel_DifferentModelsBothPrefixed(t *testing.T) { - mc := newTestCatalog(nil, nil) - assert.False(t, mc.IsSameModel("openai/gpt-4o", "anthropic/claude-3-5-sonnet")) -} - -func TestIsSameModel_CatalogBacked(t *testing.T) { - // Two model strings that look different but the catalog says they have the same base_model - mc := newTestCatalog(nil, map[string]string{ - "claude-3-5-sonnet": "claude-3-5-sonnet", - "claude-3-5-sonnet-20241022": "claude-3-5-sonnet", - }) - assert.True(t, mc.IsSameModel("claude-3-5-sonnet", "claude-3-5-sonnet-20241022")) - assert.True(t, mc.IsSameModel("claude-3-5-sonnet-20241022", "claude-3-5-sonnet")) -} - -func TestIsSameModel_AlgorithmicFallback(t *testing.T) { - // Models not in catalog — use algorithmic date stripping - mc := newTestCatalog(nil, nil) - assert.True(t, mc.IsSameModel("custom-model-20250101", "custom-model")) -} - -func TestIsSameModel_EmptyStrings(t *testing.T) { - mc := newTestCatalog(nil, nil) - assert.True(t, mc.IsSameModel("", "")) - assert.False(t, mc.IsSameModel("gpt-4o", "")) - assert.False(t, mc.IsSameModel("", "gpt-4o")) -} - -func TestIsModelAllowedForProvider_PrefixedAllowedModelInCatalog(t *testing.T) { - mc := newTestCatalog( - map[schemas.ModelProvider][]string{ - schemas.OpenRouter: {"openai/gpt-4o"}, - }, - nil, - ) - - providerConfig := configstore.ProviderConfig{} - - assert.True(t, mc.IsModelAllowedForProvider(schemas.OpenRouter, "gpt-4o", &providerConfig, []string{"openai/gpt-4o"})) -} - -func TestIsModelAllowedForProvider_CustomProviderListModelsDisabled(t *testing.T) { - mc := newTestCatalog(nil, nil) - - // Custom provider with list-models disabled + ["*"] → should return true - providerConfig := configstore.ProviderConfig{ - CustomProviderConfig: &schemas.CustomProviderConfig{ - AllowedRequests: &schemas.AllowedRequests{ - ListModels: false, - }, - }, - } - assert.True(t, mc.IsModelAllowedForProvider("custom-provider", "any-model", &providerConfig, []string{"*"})) -} - -func TestIsModelAllowedForProvider_CustomProviderListModelsEnabled(t *testing.T) { - mc := newTestCatalog( - map[schemas.ModelProvider][]string{ - "custom-provider": {"model-a"}, - }, - nil, - ) - - // Custom provider with list-models enabled + ["*"] → should go through catalog - providerConfig := configstore.ProviderConfig{ - CustomProviderConfig: &schemas.CustomProviderConfig{ - AllowedRequests: &schemas.AllowedRequests{ - ListModels: true, - }, - }, - } - // model-a is in catalog → allowed - assert.True(t, mc.IsModelAllowedForProvider("custom-provider", "model-a", &providerConfig, []string{"*"})) - // model-b is NOT in catalog → denied - assert.False(t, mc.IsModelAllowedForProvider("custom-provider", "model-b", &providerConfig, []string{"*"})) -} - -func TestIsModelAllowedForProvider_NilProviderConfig(t *testing.T) { - mc := newTestCatalog( - map[schemas.ModelProvider][]string{ - "some-provider": {"model-x"}, - }, - nil, - ) - - // nil providerConfig + ["*"] → should go through catalog (not bypass) - assert.True(t, mc.IsModelAllowedForProvider("some-provider", "model-x", nil, []string{"*"})) - assert.False(t, mc.IsModelAllowedForProvider("some-provider", "model-y", nil, []string{"*"})) -} diff --git a/framework/modelcatalog/models.go b/framework/modelcatalog/models.go index d91b337d983..547679a80b6 100644 --- a/framework/modelcatalog/models.go +++ b/framework/modelcatalog/models.go @@ -7,228 +7,225 @@ import ( "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore" - configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" ) -// GetModelCapabilityEntryForModel returns capability metadata for a model/provider pair. -// It prefers chat, then responses, then text-completion entries; if none exist, -// it falls back to the lexicographically first available mode for deterministic behavior. -func (mc *ModelCatalog) GetModelCapabilityEntryForModel(model string, provider schemas.ModelProvider) *PricingEntry { - mc.mu.RLock() - defer mc.mu.RUnlock() - - if entry := mc.getCapabilityEntryForExactModelUnsafe(model, provider); entry != nil { - return entry - } - - baseModel := mc.getBaseModelNameUnsafe(model) - if baseModel != model { - if entry := mc.getCapabilityEntryForExactModelUnsafe(baseModel, provider); entry != nil { - return entry +// GetModelsForProvider returns the effective allowed model set for the +// provider. Filtered live entries are authoritative when present (they were +// pre-gated by ListModelsPipeline against the key's allow/block/aliases); +// otherwise the datasheet view is filtered by the keyconfig aggregates. +func (mc *ModelCatalog) GetModelsForProvider(provider schemas.ModelProvider) []string { + blacklisted := mc.keyconf.BlacklistedFor(provider) + allowed := mc.keyconf.AllowedFor(provider) + + var out []string + if liveModels := mc.live.ModelsForProvider(provider); len(liveModels) > 0 { + out = liveModels + } else if datasheetModels := mc.datasheet.DatasheetModelsForProvider(provider); len(datasheetModels) > 0 && allowed != nil { + out = make([]string, 0, len(datasheetModels)) + for _, m := range datasheetModels { + if blacklisted.IsBlocked(m) { + continue + } + if allowed.IsAllowed(m) { + out = append(out, m) + } } + } else { + out = []string{} } - if entry := mc.getCapabilityEntryForModelFamilyUnsafe(baseModel, provider); entry != nil { - return entry + seen := make(map[string]struct{}, len(out)) + for _, m := range out { + seen[m] = struct{}{} } - - return nil -} - -// GetModelsForProvider returns all available models for a given provider (thread-safe) -func (mc *ModelCatalog) GetModelsForProvider(provider schemas.ModelProvider) []string { - mc.mu.RLock() - defer mc.mu.RUnlock() - - models, exists := mc.modelPool[provider] - if !exists { - return []string{} + for _, e := range mc.keyconf.EntriesFor(provider) { + if !e.Enabled { + continue + } + for alias := range e.Aliases { + if blacklisted.IsBlocked(alias) { + continue + } + if allowed == nil || !allowed.IsAllowed(alias) { + continue + } + if _, ok := seen[alias]; ok { + continue + } + seen[alias] = struct{}{} + out = append(out, alias) + } + for _, m := range e.Allowed { + if m == "*" || blacklisted.IsBlocked(m) { + continue + } + if _, ok := seen[m]; ok { + continue + } + seen[m] = struct{}{} + out = append(out, m) + } } - - // Return a copy to prevent external modification - result := make([]string, len(models)) - copy(result, models) - return result + return out } -// GetUnfilteredModelsForProvider returns all available models for a given provider (thread-safe) +// GetUnfilteredModelsForProvider returns the raw catalog view (no gate +// applied): union of live unfiltered entries and the datasheet view. func (mc *ModelCatalog) GetUnfilteredModelsForProvider(provider schemas.ModelProvider) []string { - mc.mu.RLock() - defer mc.mu.RUnlock() - - models, exists := mc.unfilteredModelPool[provider] - if !exists { - return []string{} + liveModels := mc.live.UnfilteredModelsForProvider(provider) + datasheetModels := mc.datasheet.DatasheetModelsForProvider(provider) + if len(liveModels) == 0 { + return datasheetModels } - - // Return a copy to prevent external modification - result := make([]string, len(models)) - copy(result, models) - return result + if len(datasheetModels) == 0 { + return liveModels + } + seen := make(map[string]struct{}, len(liveModels)+len(datasheetModels)) + out := make([]string, 0, len(liveModels)+len(datasheetModels)) + for _, m := range liveModels { + if _, ok := seen[m]; !ok { + seen[m] = struct{}{} + out = append(out, m) + } + } + for _, m := range datasheetModels { + if _, ok := seen[m]; !ok { + seen[m] = struct{}{} + out = append(out, m) + } + } + slices.Sort(out) + return out } -// GetDistinctBaseModelNames returns all unique base model names from the catalog (thread-safe). -// This is used for governance model selection when no specific provider is chosen. +// GetDistinctBaseModelNames returns all unique base model names from the +// datasheet. Used by governance for cross-provider model selection. func (mc *ModelCatalog) GetDistinctBaseModelNames() []string { - mc.mu.RLock() - defer mc.mu.RUnlock() - - seen := make(map[string]bool) - for _, baseName := range mc.baseModelIndex { - seen[baseName] = true - } - - result := make([]string, 0, len(seen)) - for name := range seen { - result = append(result, name) - } - return result + return mc.datasheet.DistinctBaseModelNames() } -// GetProvidersForModel returns all providers for a given model (thread-safe) +// GetProvidersForModel returns every provider that can serve the model. +// Composes across stores and applies the cross-provider special cases +// (openrouter / vertex / groq-gpt / bedrock-claude) preserved verbatim from +// the pre-refactor implementation. func (mc *ModelCatalog) GetProvidersForModel(model string) []schemas.ModelProvider { - mc.mu.RLock() - defer mc.mu.RUnlock() + baseModel := mc.datasheet.BaseModelName(model) providers := make([]schemas.ModelProvider, 0) - for provider, models := range mc.modelPool { - isModelMatch := false + seen := make(map[schemas.ModelProvider]struct{}) + for _, p := range mc.knownProviders() { + models := mc.GetModelsForProvider(p) + matched := false for _, m := range models { - if m == model || mc.getBaseModelNameUnsafe(m) == mc.getBaseModelNameUnsafe(model) { - isModelMatch = true + if m == model || mc.datasheet.BaseModelName(m) == baseModel { + matched = true break } } - if isModelMatch { - providers = append(providers, provider) + if matched { + if _, ok := seen[p]; !ok { + providers = append(providers, p) + seen[p] = struct{}{} + } } } - // Handler special provider cases - // 1. Handler openrouter models - if !slices.Contains(providers, schemas.OpenRouter) { - for _, provider := range providers { - if openRouterModels, ok := mc.modelPool[schemas.OpenRouter]; ok { - if slices.Contains(openRouterModels, string(provider)+"/"+model) { - providers = append(providers, schemas.OpenRouter) - } + // Cross-provider special cases (preserved from pre-refactor models.go). + if _, ok := seen[schemas.OpenRouter]; !ok { + openRouterModels := mc.GetModelsForProvider(schemas.OpenRouter) + for _, p := range providers { + if slices.Contains(openRouterModels, string(p)+"/"+model) { + providers = append(providers, schemas.OpenRouter) + seen[schemas.OpenRouter] = struct{}{} + break } } } - - // 2. Handle vertex models - if !slices.Contains(providers, schemas.Vertex) { - for _, provider := range providers { - if vertexModels, ok := mc.modelPool[schemas.Vertex]; ok { - if slices.Contains(vertexModels, string(provider)+"/"+model) { - providers = append(providers, schemas.Vertex) - } + if _, ok := seen[schemas.Vertex]; !ok { + vertexModels := mc.GetModelsForProvider(schemas.Vertex) + for _, p := range providers { + if slices.Contains(vertexModels, string(p)+"/"+model) { + providers = append(providers, schemas.Vertex) + seen[schemas.Vertex] = struct{}{} + break } } } - - // 3. Handle openai models for groq - if !slices.Contains(providers, schemas.Groq) && strings.Contains(model, "gpt-") { - if groqModels, ok := mc.modelPool[schemas.Groq]; ok { - if slices.Contains(groqModels, "openai/"+model) { - providers = append(providers, schemas.Groq) + if _, ok := seen[schemas.Groq]; !ok && strings.Contains(model, "gpt-") { + if slices.Contains(mc.GetModelsForProvider(schemas.Groq), "openai/"+model) { + providers = append(providers, schemas.Groq) + } + } + if _, ok := seen[schemas.Bedrock]; !ok && strings.Contains(model, "claude") { + for _, bedrockModel := range mc.GetModelsForProvider(schemas.Bedrock) { + if strings.Contains(bedrockModel, model) { + providers = append(providers, schemas.Bedrock) + break } } } - // 4. Handle anthropic models for bedrock - if !slices.Contains(providers, schemas.Bedrock) && strings.Contains(model, "claude") { - if bedrockModels, ok := mc.modelPool[schemas.Bedrock]; ok { - for _, bedrockModel := range bedrockModels { - if strings.Contains(bedrockModel, model) { - providers = append(providers, schemas.Bedrock) - break - } - } + for _, p := range mc.keyconf.Providers() { + if _, ok := seen[p]; ok { + continue + } + if mc.keyconf.BlacklistedFor(p).IsBlocked(model) { + continue + } + allowed := mc.keyconf.AllowedFor(p) + matched := false + if _, hit := mc.keyconf.ResolveAlias(p, model); hit && allowed.IsAllowed(model) { + matched = true + } else if allowed.Contains(model) { + matched = true + } else if allowed.IsUnrestricted() && + len(mc.datasheet.DatasheetModelsForProvider(p)) == 0 && + len(mc.live.UnfilteredModelsForProvider(p)) == 0 { + matched = true + } + if matched { + providers = append(providers, p) + seen[p] = struct{}{} } } return providers } -// IsModelAllowedForProvider checks if a model is allowed for a specific provider -// based on the allowed models list and catalog data. It handles all cross-provider -// logic including provider-prefixed models and special routing rules. -// -// Parameters: -// - provider: The provider to check against -// - model: The model name (without provider prefix, e.g., "gpt-4o" or "claude-3-5-sonnet") -// - allowedModels: List of allowed model names (can be empty, can include provider prefixes) -// -// Behavior: -// - If allowedModels is ["*"]: Uses model catalog to check if provider supports the model -// (delegates to GetProvidersForModel which handles all cross-provider logic) -// - If allowedModels is empty ([]): Deny-by-default — returns false for any provider/model pair -// - If allowedModels is not empty: Checks if model matches any entry in the list -// Provider-specific validation: -// - Direct matches: "gpt-4o" in allowedModels for any provider -// - Prefixed matches: Only if the prefixed model exists in provider's catalog -// (e.g., "openai/gpt-4o" in allowedModels only matches if openrouter's catalog -// contains "openai/gpt-4o" AND the model part matches the request) -// -// Returns: -// - bool: true if the model is allowed for the provider, false otherwise +// IsModelAllowedForProvider checks whether the model is allowed for the +// provider given an explicit allowedModels list (used by VK governance +// checks, not by the static keyconfig allow set). // -// Examples: -// -// // Wildcard allowedModels - uses catalog to check provider support -// mc.IsModelAllowedForProvider("openrouter", "claude-3-5-sonnet", []string{"*"}) -// // Returns: true (catalog knows openrouter has "anthropic/claude-3-5-sonnet") -// -// // Empty allowedModels - deny all (deny-by-default) -// mc.IsModelAllowedForProvider("openrouter", "claude-3-5-sonnet", []string{}) -// // Returns: false (no models are permitted) -// -// // Explicit allowedModels with prefix - validates against catalog -// mc.IsModelAllowedForProvider("openrouter", "gpt-4o", []string{"openai/gpt-4o"}) -// // Returns: true (openrouter's catalog contains "openai/gpt-4o" AND model part is "gpt-4o") -// -// // Explicit allowedModels with prefix - wrong model -// mc.IsModelAllowedForProvider("openrouter", "claude-3-5-sonnet", []string{"openai/gpt-4o"}) -// // Returns: false (model part "gpt-4o" doesn't match request "claude-3-5-sonnet") -// -// // Explicit allowedModels without prefix -// mc.IsModelAllowedForProvider("openai", "gpt-4o", []string{"gpt-4o"}) -// // Returns: true (direct match) +// - allowedModels=["*"]: defer to GetProvidersForModel (with custom-provider +// fast path when list-models is disabled). +// - allowedModels=[]: deny-by-default. +// - explicit allowedModels: direct or provider-prefixed match against the +// provider's catalog. func (mc *ModelCatalog) IsModelAllowedForProvider(provider schemas.ModelProvider, model string, providerConfig *configstore.ProviderConfig, allowedModels schemas.WhiteList) bool { - // Case 1: ["*"] = allow all models; use catalog to determine support - // Empty allowedModels = deny all (fail-safe deny-by-default) + isCustomProvider := false + hasListModelsEndpointDisabled := false + if providerConfig != nil && providerConfig.CustomProviderConfig != nil { + isCustomProvider = true + hasListModelsEndpointDisabled = !providerConfig.CustomProviderConfig.IsOperationAllowed(schemas.ListModelsRequest) + } + if allowedModels.IsUnrestricted() { - // Providers whose models the catalog cannot enumerate (custom without list-models, - // or keyless self-hosted vLLM/Ollama/SGL) cannot be cross-checked against catalog - // membership, so a wildcard allow-list permits any model for them. - if mc.IsCatalogOpaqueProvider(provider, providerConfig) { + if isCustomProvider && hasListModelsEndpointDisabled { return true } - supportedProviders := mc.GetProvidersForModel(model) - return slices.Contains(supportedProviders, provider) + return slices.Contains(mc.GetProvidersForModel(model), provider) } if allowedModels.IsEmpty() { return false } - // Case 2: Explicit allowedModels = check if model matches any entry - // Get provider's catalog models for validation of prefixed entries providerCatalogModels := mc.GetModelsForProvider(provider) - for _, allowedModel := range allowedModels { - // Direct match: "gpt-4o" == "gpt-4o" if allowedModel == model { return true } - - // Provider-prefixed match: verify it exists in provider's catalog first - // This ensures we only allow provider-specific model combinations that are actually supported if strings.Contains(allowedModel, "/") { - // Check if this exact prefixed model exists in the provider's catalog - // e.g., for openrouter, check if "openai/gpt-4o" is in its catalog if slices.Contains(providerCatalogModels, allowedModel) { - // Extract the model part and compare with request _, modelPart := schemas.ParseModelString(allowedModel, "") if modelPart == model { return true @@ -236,211 +233,21 @@ func (mc *ModelCatalog) IsModelAllowedForProvider(provider schemas.ModelProvider } } } - return false } -// IsCatalogOpaqueProvider reports whether the catalog cannot enumerate the models a provider -// serves, so a wildcard ("*") allow-list must be honored as allow-all rather than cross-checked -// against catalog membership. True for custom providers and for native providers the catalog has -// no model list for (keyless self-hosted vLLM/Ollama/SGL, or providers without list-models -// support). Shared by OSS governance and enterprise load-balancing so the rule has one definition. -func (mc *ModelCatalog) IsCatalogOpaqueProvider(provider schemas.ModelProvider, providerConfig *configstore.ProviderConfig) bool { - if providerConfig != nil && providerConfig.CustomProviderConfig != nil { - // A custom provider is opaque only when it cannot list its models. If it supports the - // list-models endpoint, the catalog can enumerate its models, so it is NOT opaque. - return !providerConfig.CustomProviderConfig.IsOperationAllowed(schemas.ListModelsRequest) - } - if mc == nil { - return false - } - // Only an emptiness check is needed, so read modelPool directly under the - // read lock rather than calling GetModelsForProvider, which allocates and - // copies the full slice on this request-path check. - mc.mu.RLock() - defer mc.mu.RUnlock() - return len(mc.modelPool[provider]) == 0 -} - -// GetBaseModelName returns the canonical base model name for a given model string. -// It uses the pre-computed base_model from the pricing catalog when available, -// falling back to algorithmic date/version stripping for models not in the catalog. -// -// Examples: -// -// mc.GetBaseModelName("gpt-4o") // Returns: "gpt-4o" -// mc.GetBaseModelName("openai/gpt-4o") // Returns: "gpt-4o" -// mc.GetBaseModelName("gpt-4o-2024-08-06") // Returns: "gpt-4o" (algorithmic fallback) func (mc *ModelCatalog) GetBaseModelName(model string) string { - mc.mu.RLock() - defer mc.mu.RUnlock() - return mc.getBaseModelNameUnsafe(model) -} - -// getBaseModelNameUnsafe returns the canonical base model name for a given model string without locking. -// This is used to avoid locking overhead when getting the base model name for many models. -// Make sure the caller function is holding the read lock before calling this function. -// It is not safe to use this function when the model pool is being updated. -func (mc *ModelCatalog) getBaseModelNameUnsafe(model string) string { - // Step 1: Direct lookup in base model index - if base, ok := mc.baseModelIndex[model]; ok { - return base - } - - // Step 2: Strip provider prefix and try again - _, baseName := schemas.ParseModelString(model, "") - if baseName != model { - if base, ok := mc.baseModelIndex[baseName]; ok { - return base - } - } - - // Step 3: Fallback to algorithmic date/version stripping - // (for models not in the catalog, e.g., user-configured custom models) - return schemas.BaseModelName(baseName) + return mc.datasheet.BaseModelName(model) } -// IsSameModel checks if two model strings refer to the same underlying model. -// It compares the canonical base model names derived from the pricing catalog -// (or algorithmic fallback for models not in the catalog). -// -// Examples: -// -// mc.IsSameModel("gpt-4o", "gpt-4o") // true (direct match) -// mc.IsSameModel("openai/gpt-4o", "gpt-4o") // true (same base model) -// mc.IsSameModel("gpt-4o", "claude-3-5-sonnet") // false (different models) -// mc.IsSameModel("openai/gpt-4o", "anthropic/claude-3-5-sonnet") // false func (mc *ModelCatalog) IsSameModel(model1, model2 string) bool { - if model1 == model2 { - return true - } - return mc.GetBaseModelName(model1) == mc.GetBaseModelName(model2) + return mc.datasheet.IsSameModel(model1, model2) } -// DeleteModelDataForProvider deletes all model data from the pool for a given provider -func (mc *ModelCatalog) DeleteModelDataForProvider(provider schemas.ModelProvider) { - mc.mu.Lock() - defer mc.mu.Unlock() - - delete(mc.modelPool, provider) - delete(mc.unfilteredModelPool, provider) -} - -// UpsertModelDataForProvider upserts model data for a given provider -func (mc *ModelCatalog) UpsertModelDataForProvider(provider schemas.ModelProvider, modelData *schemas.BifrostListModelsResponse, allowedModels []schemas.Model) { - if modelData == nil { - return - } - mc.mu.Lock() - defer mc.mu.Unlock() - - // Populating models from pricing data for the given provider - // Provider models map - providerModels := []string{} - // Iterate through all pricing data to collect models per provider - for _, pricing := range mc.pricingData { - // Normalize provider before adding to model pool - normalizedProvider := schemas.ModelProvider(normalizeProvider(pricing.Provider)) - // We will only add models for the given provider - if normalizedProvider != provider { - continue - } - // Add model to the provider's model set (using map for deduplication) - if slices.Contains(providerModels, pricing.Model) { - continue - } - providerModels = append(providerModels, pricing.Model) - // Build base model index from pre-computed base_model field - if pricing.BaseModel != "" { - mc.baseModelIndex[pricing.Model] = pricing.BaseModel - } - } - // If modelData is empty, then we allow all models - if len(modelData.Data) == 0 && len(allowedModels) == 0 { - mc.modelPool[provider] = providerModels - return - } - // Here we make sure that we still keep the backup for model catalog intact - // So we start with a existing model pool and add the new models from incoming data - finalModelList := make([]string, 0) - seenModels := make(map[string]bool) - // Case where list models failed but we have allowed models from keys - if len(modelData.Data) == 0 && len(allowedModels) > 0 { - for _, allowedModel := range allowedModels { - parsedProvider, parsedModel := schemas.ParseModelString(allowedModel.ID, "") - if parsedProvider != provider { - continue - } - if !seenModels[parsedModel] { - seenModels[parsedModel] = true - finalModelList = append(finalModelList, parsedModel) - } - } - } - for _, model := range modelData.Data { - parsedProvider, parsedModel := schemas.ParseModelString(model.ID, "") - if parsedProvider != provider { - continue - } - if !seenModels[parsedModel] { - seenModels[parsedModel] = true - finalModelList = append(finalModelList, parsedModel) - } - } - - if len(allowedModels) == 0 { - for _, model := range providerModels { - if !seenModels[model] { - seenModels[model] = true - finalModelList = append(finalModelList, model) - } - } - } - mc.modelPool[provider] = finalModelList -} - -// UpsertUnfilteredModelDataForProvider upserts unfiltered model data for a given provider -func (mc *ModelCatalog) UpsertUnfilteredModelDataForProvider(provider schemas.ModelProvider, modelData *schemas.BifrostListModelsResponse) { - if modelData == nil { - return - } - mc.mu.Lock() - defer mc.mu.Unlock() - - // Populating models from pricing data for the given provider - providerModels := []string{} - seenModels := make(map[string]bool) - for _, pricing := range mc.pricingData { - normalizedProvider := schemas.ModelProvider(normalizeProvider(pricing.Provider)) - if normalizedProvider != provider { - continue - } - if !seenModels[pricing.Model] { - seenModels[pricing.Model] = true - providerModels = append(providerModels, pricing.Model) - } - } - for _, model := range modelData.Data { - parsedProvider, parsedModel := schemas.ParseModelString(model.ID, "") - if parsedProvider != provider { - continue - } - if !seenModels[parsedModel] { - seenModels[parsedModel] = true - providerModels = append(providerModels, parsedModel) - } - } - mc.unfilteredModelPool[provider] = providerModels -} - -// RefineModelForProvider refines the model for a given provider by performing a lookup -// in mc.modelPool and using schemas.ParseModelString to extract provider and model parts. -// e.g. "gpt-oss-120b" for groq provider -> "openai/gpt-oss-120b" -// -// Behavior: -// - When the provider's catalog (mc.modelPool) yields multiple matching models, returns an error -// - When exactly one match is found, returns the fully-qualified model (provider/model format) -// - When the provider is not handled or no refinement is needed, returns the original model unchanged +// RefineModelForProvider refines a model identifier for providers that need +// a leading "provider/" segment (Groq, Replicate). Returns the original +// model unchanged when no refinement applies, or an error when multiple +// catalog candidates match ambiguously. func (mc *ModelCatalog) RefineModelForProvider(provider schemas.ModelProvider, model string) (string, error) { switch provider { case schemas.Groq: @@ -454,179 +261,14 @@ func (mc *ModelCatalog) RefineModelForProvider(provider schemas.ModelProvider, m return model, nil } -// SetPricingOverrides replaces the full in-memory pricing override set. -func (mc *ModelCatalog) SetPricingOverrides(rows []configstoreTables.TablePricingOverride) error { - seen := make(map[string]int, len(rows)) - overrides := make([]PricingOverride, 0, len(rows)) - for i := range rows { - o, err := convertTablePricingOverrideToPricingOverride(&rows[i]) - if err != nil { - return err - } - if idx, exists := seen[o.ID]; exists { - overrides[idx] = o // last entry wins for duplicate IDs - } else { - seen[o.ID] = len(overrides) - overrides = append(overrides, o) - } - } - mc.overridesMu.Lock() - mc.rawOverrides = overrides - mc.customPricing = buildCustomPricingData(overrides) - mc.overridesMu.Unlock() - return nil -} - -// UpsertPricingOverrides inserts or replaces one or more pricing overrides in a single -// operation, rebuilding the lookup map only once at the end. -func (mc *ModelCatalog) UpsertPricingOverrides(rows ...*configstoreTables.TablePricingOverride) error { - // Deduplicate the input batch by ID (last entry wins) and build the - // incoming set for O(1) lookup when filtering existing rawOverrides. - seenIncoming := make(map[string]int, len(rows)) - overrides := make([]PricingOverride, 0, len(rows)) - for _, row := range rows { - o, err := convertTablePricingOverrideToPricingOverride(row) - if err != nil { - return err - } - if idx, exists := seenIncoming[o.ID]; exists { - overrides[idx] = o // last entry wins for duplicate IDs - } else { - seenIncoming[o.ID] = len(overrides) - overrides = append(overrides, o) - } - } - - mc.overridesMu.Lock() - defer mc.overridesMu.Unlock() - - updated := make([]PricingOverride, 0, len(mc.rawOverrides)+len(overrides)) - for _, o := range mc.rawOverrides { - if _, replacing := seenIncoming[o.ID]; !replacing { - updated = append(updated, o) - } - } - updated = append(updated, overrides...) - mc.rawOverrides = updated - mc.customPricing = buildCustomPricingData(updated) - return nil -} - -// DeletePricingOverride removes a pricing override by ID. -func (mc *ModelCatalog) DeletePricingOverride(id string) { - mc.overridesMu.Lock() - defer mc.overridesMu.Unlock() - - updated := make([]PricingOverride, 0, len(mc.rawOverrides)) - for _, o := range mc.rawOverrides { - if o.ID != id { - updated = append(updated, o) - } - } - mc.rawOverrides = updated - mc.customPricing = buildCustomPricingData(updated) -} - -// IsTextCompletionSupported checks if a model supports text completion for the given provider. -// Returns true if the model has pricing data for text completion ("text_completion"), -// false otherwise. This is used by the litellmcompat plugin to determine whether to -// convert text completion requests to chat completion requests. -func (mc *ModelCatalog) IsTextCompletionSupported(model string, provider schemas.ModelProvider) bool { - mc.mu.RLock() - defer mc.mu.RUnlock() - // Check for text completion mode in pricing data - key := makeKey(model, normalizeProvider(string(provider)), normalizeRequestType(schemas.TextCompletionRequest)) - _, ok := mc.pricingData[key] - return ok -} - -// HELPER FUNCTIONS - -func (mc *ModelCatalog) getCapabilityEntryForExactModelUnsafe(model string, provider schemas.ModelProvider) *PricingEntry { - preferredModes := []schemas.RequestType{ - schemas.ChatCompletionRequest, - schemas.ResponsesRequest, - schemas.TextCompletionRequest, - } - - for _, mode := range preferredModes { - key := makeKey(model, string(provider), normalizeRequestType(mode)) - pricing, ok := mc.pricingData[key] - if ok { - return convertTableModelPricingToPricingData(&pricing) - } - } - - prefix := model + "|" + string(provider) + "|" - matchingKeys := make([]string, 0) - for key := range mc.pricingData { - if strings.HasPrefix(key, prefix) { - matchingKeys = append(matchingKeys, key) - } - } - return mc.selectCapabilityEntryFromKeysUnsafe(matchingKeys) -} - -func (mc *ModelCatalog) getCapabilityEntryForModelFamilyUnsafe(baseModel string, provider schemas.ModelProvider) *PricingEntry { - if baseModel == "" { - return nil - } - - matchingKeys := make([]string, 0) - for key, pricing := range mc.pricingData { - if normalizeProvider(pricing.Provider) != string(provider) { - continue - } - if mc.getBaseModelNameUnsafe(pricing.Model) != baseModel { - continue - } - matchingKeys = append(matchingKeys, key) - } - return mc.selectCapabilityEntryFromKeysUnsafe(matchingKeys) -} - -func (mc *ModelCatalog) selectCapabilityEntryFromKeysUnsafe(matchingKeys []string) *PricingEntry { - if len(matchingKeys) == 0 { - return nil - } - - preferredModes := []string{ - normalizeRequestType(schemas.ChatCompletionRequest), - normalizeRequestType(schemas.ResponsesRequest), - normalizeRequestType(schemas.TextCompletionRequest), - } - - for _, mode := range preferredModes { - modeMatches := make([]string, 0) - for _, key := range matchingKeys { - parts := strings.SplitN(key, "|", 3) - if len(parts) != 3 || parts[2] != mode { - continue - } - modeMatches = append(modeMatches, key) - } - if len(modeMatches) == 0 { - continue - } - slices.Sort(modeMatches) - pricing := mc.pricingData[modeMatches[0]] - return convertTableModelPricingToPricingData(&pricing) - } - - slices.Sort(matchingKeys) - pricing := mc.pricingData[matchingKeys[0]] - return convertTableModelPricingToPricingData(&pricing) -} - // refineNestedProviderModel resolves provider-native model slugs such as -// "openai/gpt-5-nano" from a base model request like "gpt-5-nano". -// It only considers catalog entries whose leading segment is a known Bifrost provider, -// so Replicate owner/model identifiers like "meta/llama-3-8b" are left untouched. +// "openai/gpt-5-nano" from a base model request like "gpt-5-nano". Only +// considers catalog entries whose leading segment is a known Bifrost +// provider so Replicate owner/model identifiers like "meta/llama-3-8b" are +// left untouched. func (mc *ModelCatalog) refineNestedProviderModel(provider schemas.ModelProvider, model string) (string, error) { - mc.mu.RLock() - models, ok := mc.modelPool[provider] - mc.mu.RUnlock() - if !ok { + models := mc.GetModelsForProvider(provider) + if len(models) == 0 { return model, nil } @@ -637,7 +279,6 @@ func (mc *ModelCatalog) refineNestedProviderModel(provider schemas.ModelProvider if providerPart == "" || model != modelPart { continue } - candidate := string(providerPart) + "/" + modelPart if _, seen := seenCandidates[candidate]; seen { continue diff --git a/framework/modelcatalog/pool.go b/framework/modelcatalog/pool.go new file mode 100644 index 00000000000..f4e316f27f4 --- /dev/null +++ b/framework/modelcatalog/pool.go @@ -0,0 +1,101 @@ +// Editing the model pool. These methods are the push surface server.go +// orchestrates against — fetched list-models responses go into live, +// configstore key edits go into keyconfig, and the composed pool is what +// the read methods in models.go return. +package modelcatalog + +import ( + "github.com/maximhq/bifrost/core/schemas" +) + +// UpsertLive caches one (provider, keyID, unfiltered) list-models response. +func (mc *ModelCatalog) UpsertLive(provider schemas.ModelProvider, keyID string, unfiltered bool, models []string) { + mc.live.Upsert(provider, keyID, unfiltered, models) +} + +// InvalidateLive drops both filtered + unfiltered live entries for one key. +func (mc *ModelCatalog) InvalidateLive(provider schemas.ModelProvider, keyID string) { + mc.live.Invalidate(provider, keyID) +} + +// InvalidateLiveProvider drops all live entries for a provider. +func (mc *ModelCatalog) InvalidateLiveProvider(provider schemas.ModelProvider) { + mc.live.InvalidateProvider(provider) +} + +// SetKeyConfigForProvider replaces the keyconfig snapshot for one provider. +func (mc *ModelCatalog) SetKeyConfigForProvider(provider schemas.ModelProvider, keys []schemas.Key) { + mc.keyconf.SetProvider(provider, keys) +} + +// ReplaceKeyConfig atomically resets the keyconfig snapshot for all providers. +func (mc *ModelCatalog) ReplaceKeyConfig(snapshot map[schemas.ModelProvider][]schemas.Key) { + mc.keyconf.Replace(snapshot) +} + +// RemoveKeyConfigForProvider drops keyconfig state for the provider. +func (mc *ModelCatalog) RemoveKeyConfigForProvider(provider schemas.ModelProvider) { + mc.keyconf.RemoveProvider(provider) +} + +// KeyConfigEntries returns the per-key entries for one provider (used by +// orchestration to know which keys to fan list-models calls across). +func (mc *ModelCatalog) KeyConfigEntries(provider schemas.ModelProvider) []KeyConfigEntry { + return mc.keyconf.EntriesFor(provider) +} + +// ResolveAlias returns which key owns an alias on the provider and its +// AliasConfig. +func (mc *ModelCatalog) ResolveAlias(provider schemas.ModelProvider, model string) (AliasOwner, bool) { + return mc.keyconf.ResolveAlias(provider, model) +} + +// KeysAllowingModel returns the IDs of enabled keys that can serve the model. +func (mc *ModelCatalog) KeysAllowingModel(provider schemas.ModelProvider, model string) []string { + return mc.keyconf.KeysAllowingModel(provider, model) +} + +// AllowedModelsForProvider returns the aggregated whitelist for the +// provider (union of enabled keys' Models minus per-key blacklists, or +// ["*"] when any key is unrestricted). Used by the load balancer to know +// what each provider can serve without re-walking the configstore. +func (mc *ModelCatalog) AllowedModelsForProvider(provider schemas.ModelProvider) schemas.WhiteList { + return mc.keyconf.AllowedFor(provider) +} + +// BlacklistedModelsForProvider returns the intersection of enabled keys' +// BlacklistedModels for the provider — a model is included only when every +// enabled key blacklists it. +func (mc *ModelCatalog) BlacklistedModelsForProvider(provider schemas.ModelProvider) schemas.BlackList { + return mc.keyconf.BlacklistedFor(provider) +} + +// ConfiguredProviders returns every provider with at least one entry in +// keyconfig. Used by the load balancer's provider selection where the +// configured-provider set is the routing-eligible universe. +func (mc *ModelCatalog) ConfiguredProviders() []schemas.ModelProvider { + return mc.keyconf.Providers() +} + +// extractModelIDs flattens a list-models response into bare model +// identifiers, filtering entries whose ID prefix doesn't match the +// requested provider. +func extractModelIDs(resp *schemas.BifrostListModelsResponse, provider schemas.ModelProvider) []string { + if resp == nil { + return nil + } + seen := make(map[string]struct{}, len(resp.Data)) + out := make([]string, 0, len(resp.Data)) + for _, m := range resp.Data { + parsedProvider, parsedModel := schemas.ParseModelString(m.ID, "") + if parsedProvider != "" && parsedProvider != provider { + continue + } + if _, ok := seen[parsedModel]; ok { + continue + } + seen[parsedModel] = struct{}{} + out = append(out, parsedModel) + } + return out +} diff --git a/framework/modelcatalog/pricing.go b/framework/modelcatalog/pricing.go index d58ab529d52..f4e33ec4155 100644 --- a/framework/modelcatalog/pricing.go +++ b/framework/modelcatalog/pricing.go @@ -2,1528 +2,60 @@ package modelcatalog import ( "context" - "fmt" - "strconv" - "strings" - "github.com/bytedance/sonic" "github.com/maximhq/bifrost/core/schemas" configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/maximhq/bifrost/framework/modelcatalog/datasheet" ) -// Default sync interval and config key -const ( - TokenTierAbove272K = 272000 - TokenTierAbove200K = 200000 - TokenTierAbove128K = 128000 -) - -// PricingEntry represents a single model's pricing information. -// Field names and JSON tags match the datasheet schema exactly. -// AdditionalAttributes carries editorial metadata stored on the pricing row, never populated from the URL datasheet — only from DB reads via the management API. -type PricingEntry struct { - BaseModel string `json:"base_model,omitempty"` - Provider string `json:"provider"` - Mode string `json:"mode"` - - ContextLength *int `json:"context_length,omitempty"` - MaxInputTokens *int `json:"max_input_tokens,omitempty"` - MaxOutputTokens *int `json:"max_output_tokens,omitempty"` - Architecture *schemas.Architecture `json:"architecture,omitempty"` - - // AdditionalAttributes carries editorial metadata stored on the pricing - // row (e.g. description). Populated from the DB read path only; the - // json:"-" tag prevents URL datasheet payloads from ever feeding into - // this field via json.Unmarshal. - AdditionalAttributes map[string]string `json:"-"` - - PricingOptions +// GetModelCapabilityEntryForModel returns capability metadata for a +// (model, provider) pair. Prefers chat, then responses, then text-completion +// entries; falls back to the lexicographically first available mode for +// deterministic behavior. +func (mc *ModelCatalog) GetModelCapabilityEntryForModel(model string, provider schemas.ModelProvider) *PricingEntry { + return mc.datasheet.GetCapabilityEntry(model, provider) } -// UnmarshalJSON implements json.Unmarshaler for PricingEntry. -// It handles the special case where search_context_cost_per_query may arrive as either -// a plain float64 or a tiered object {"search_context_size_low":…, -// "search_context_size_medium":…, "search_context_size_high":…}. -func (p *PricingEntry) UnmarshalJSON(data []byte) error { - // Type alias breaks the UnmarshalJSON recursion while keeping all other fields. - type PricingEntryAlias PricingEntry - var raw struct { - PricingEntryAlias - SearchContextCostPerQuery *struct { - Low *float64 `json:"search_context_size_low"` - Medium *float64 `json:"search_context_size_medium"` - High *float64 `json:"search_context_size_high"` - } `json:"search_context_cost_per_query,omitempty"` - } - if err := sonic.Unmarshal(data, &raw); err != nil { - return err - } - *p = PricingEntry(raw.PricingEntryAlias) - - // search_context_cost_per_query arrives as a tiered object – all three values are - // equal for non-Perplexity providers; we prefer medium, then low, then high. - // Perplexity always returns a pre-computed total_cost so the per-query rate is - // never consumed for that provider. - if q := raw.SearchContextCostPerQuery; q != nil { - switch { - case q.Medium != nil: - p.SearchContextCostPerQuery = q.Medium - case q.Low != nil: - p.SearchContextCostPerQuery = q.Low - case q.High != nil: - p.SearchContextCostPerQuery = q.High - } - } - return nil +// IsRequestTypeSupported preserves the historical (model, provider, +// requestType) signature; provider is ignored (the underlying datasheet +// index is keyed by model only). +func (mc *ModelCatalog) IsRequestTypeSupported(model string, provider schemas.ModelProvider, requestType schemas.RequestType) bool { + return mc.datasheet.IsRequestTypeSupported(model, requestType) } -type PricingOptions struct { - // Costs - Text - InputCostPerToken *float64 `json:"input_cost_per_token,omitempty"` - OutputCostPerToken *float64 `json:"output_cost_per_token,omitempty"` - InputCostPerTokenBatches *float64 `json:"input_cost_per_token_batches,omitempty"` - OutputCostPerTokenBatches *float64 `json:"output_cost_per_token_batches,omitempty"` - InputCostPerTokenPriority *float64 `json:"input_cost_per_token_priority,omitempty"` - OutputCostPerTokenPriority *float64 `json:"output_cost_per_token_priority,omitempty"` - InputCostPerTokenFlex *float64 `json:"input_cost_per_token_flex,omitempty"` - OutputCostPerTokenFlex *float64 `json:"output_cost_per_token_flex,omitempty"` - InputCostPerCharacter *float64 `json:"input_cost_per_character,omitempty"` - // Costs - 128k Tier - InputCostPerTokenAbove128kTokens *float64 `json:"input_cost_per_token_above_128k_tokens,omitempty"` - InputCostPerImageAbove128kTokens *float64 `json:"input_cost_per_image_above_128k_tokens,omitempty"` - InputCostPerVideoPerSecondAbove128kTokens *float64 `json:"input_cost_per_video_per_second_above_128k_tokens,omitempty"` - InputCostPerAudioPerSecondAbove128kTokens *float64 `json:"input_cost_per_audio_per_second_above_128k_tokens,omitempty"` - OutputCostPerTokenAbove128kTokens *float64 `json:"output_cost_per_token_above_128k_tokens,omitempty"` - // Costs - 200k Tier - InputCostPerTokenAbove200kTokens *float64 `json:"input_cost_per_token_above_200k_tokens,omitempty"` - InputCostPerTokenAbove200kTokensPriority *float64 `json:"input_cost_per_token_above_200k_tokens_priority,omitempty"` - OutputCostPerTokenAbove200kTokens *float64 `json:"output_cost_per_token_above_200k_tokens,omitempty"` - OutputCostPerTokenAbove200kTokensPriority *float64 `json:"output_cost_per_token_above_200k_tokens_priority,omitempty"` - // Costs - 272k Tier - InputCostPerTokenAbove272kTokens *float64 `json:"input_cost_per_token_above_272k_tokens,omitempty"` - InputCostPerTokenAbove272kTokensPriority *float64 `json:"input_cost_per_token_above_272k_tokens_priority,omitempty"` - OutputCostPerTokenAbove272kTokens *float64 `json:"output_cost_per_token_above_272k_tokens,omitempty"` - OutputCostPerTokenAbove272kTokensPriority *float64 `json:"output_cost_per_token_above_272k_tokens_priority,omitempty"` - - // Costs - Cache - CacheCreationInputTokenCost *float64 `json:"cache_creation_input_token_cost,omitempty"` - CacheReadInputTokenCost *float64 `json:"cache_read_input_token_cost,omitempty"` - CacheCreationInputTokenCostAbove200kTokens *float64 `json:"cache_creation_input_token_cost_above_200k_tokens,omitempty"` - CacheReadInputTokenCostAbove200kTokens *float64 `json:"cache_read_input_token_cost_above_200k_tokens,omitempty"` - CacheReadInputTokenCostAbove200kTokensPriority *float64 `json:"cache_read_input_token_cost_above_200k_tokens_priority,omitempty"` - CacheCreationInputTokenCostAbove1hr *float64 `json:"cache_creation_input_token_cost_above_1hr,omitempty"` - CacheCreationInputTokenCostAbove1hrAbove200kTokens *float64 `json:"cache_creation_input_token_cost_above_1hr_above_200k_tokens,omitempty"` - CacheCreationInputAudioTokenCost *float64 `json:"cache_creation_input_audio_token_cost,omitempty"` - CacheReadInputTokenCostPriority *float64 `json:"cache_read_input_token_cost_priority,omitempty"` - CacheReadInputTokenCostFlex *float64 `json:"cache_read_input_token_cost_flex,omitempty"` - CacheReadInputImageTokenCost *float64 `json:"cache_read_input_image_token_cost,omitempty"` - CacheReadInputTokenCostAbove272kTokens *float64 `json:"cache_read_input_token_cost_above_272k_tokens,omitempty"` - CacheReadInputTokenCostAbove272kTokensPriority *float64 `json:"cache_read_input_token_cost_above_272k_tokens_priority,omitempty"` - - // Costs - Image - InputCostPerImage *float64 `json:"input_cost_per_image,omitempty"` - InputCostPerPixel *float64 `json:"input_cost_per_pixel,omitempty"` - OutputCostPerImage *float64 `json:"output_cost_per_image,omitempty"` - OutputCostPerPixel *float64 `json:"output_cost_per_pixel,omitempty"` - OutputCostPerImagePremiumImage *float64 `json:"output_cost_per_image_premium_image,omitempty"` - OutputCostPerImageAbove512x512Pixels *float64 `json:"output_cost_per_image_above_512_and_512_pixels,omitempty"` - OutputCostPerImageAbove512x512PixelsPremium *float64 `json:"output_cost_per_image_above_512_and_512_pixels_and_premium_image,omitempty"` - OutputCostPerImageAbove1024x1024Pixels *float64 `json:"output_cost_per_image_above_1024_and_1024_pixels,omitempty"` - OutputCostPerImageAbove1024x1024PixelsPremium *float64 `json:"output_cost_per_image_above_1024_and_1024_pixels_and_premium_image,omitempty"` - OutputCostPerImageAbove2048x2048Pixels *float64 `json:"output_cost_per_image_above_2048_and_2048_pixels,omitempty"` - OutputCostPerImageAbove4096x4096Pixels *float64 `json:"output_cost_per_image_above_4096_and_4096_pixels,omitempty"` - OutputCostPerImageLowQuality *float64 `json:"output_cost_per_image_low_quality,omitempty"` - OutputCostPerImageMediumQuality *float64 `json:"output_cost_per_image_medium_quality,omitempty"` - OutputCostPerImageHighQuality *float64 `json:"output_cost_per_image_high_quality,omitempty"` - OutputCostPerImageAutoQuality *float64 `json:"output_cost_per_image_auto_quality,omitempty"` - InputCostPerImageToken *float64 `json:"input_cost_per_image_token,omitempty"` - OutputCostPerImageToken *float64 `json:"output_cost_per_image_token,omitempty"` - - // Costs - Audio/Video - InputCostPerAudioToken *float64 `json:"input_cost_per_audio_token,omitempty"` - InputCostPerAudioPerSecond *float64 `json:"input_cost_per_audio_per_second,omitempty"` - InputCostPerSecond *float64 `json:"input_cost_per_second,omitempty"` - InputCostPerVideoPerSecond *float64 `json:"input_cost_per_video_per_second,omitempty"` - OutputCostPerAudioToken *float64 `json:"output_cost_per_audio_token,omitempty"` - OutputCostPerVideoPerSecond *float64 `json:"output_cost_per_video_per_second,omitempty"` - OutputCostPerSecond *float64 `json:"output_cost_per_second,omitempty"` - - // Costs - Other - // - // SearchContextCostPerQuery is stored as a single float64, but the pricing datasheet - // represents it as a tiered object with three keys: search_context_size_low, - // search_context_size_medium, and search_context_size_high. For every provider except - // Perplexity the three tier values are identical, so we collapse the object to its - // medium tier value (falling back to low then high). Perplexity always returns a - // pre-computed total_cost in its usage response, so the per-query rate is never - // consumed for that provider; the collapsed value is therefore correct in all cases. - // See UnmarshalJSON below for the custom decoding logic. - SearchContextCostPerQuery *float64 `json:"search_context_cost_per_query,omitempty"` - CodeInterpreterCostPerSession *float64 `json:"code_interpreter_cost_per_session,omitempty"` - - // Costs - OCR - OCRCostPerPage *float64 `json:"ocr_cost_per_page,omitempty"` - AnnotationCostPerPage *float64 `json:"annotation_cost_per_page,omitempty"` +func (mc *ModelCatalog) GetSupportedParameters(model string) []string { + return mc.datasheet.GetSupportedParameters(model) } -// serviceTier captures the OpenAI service_tier value from a response. -// Add new tier flags here as OpenAI introduces them. -type serviceTier struct { - isPriority bool // true when service_tier == "priority" - isFlex bool // true when service_tier == "flex" +func (mc *ModelCatalog) IsTextCompletionSupported(model string, provider schemas.ModelProvider) bool { + return mc.datasheet.IsTextCompletionSupported(model, provider) } -// costInput holds the extracted usage data from a BifrostResponse, -// normalized for the pricing engine. -type costInput struct { - usage *schemas.BifrostLLMUsage - audioTextInputChars int - audioSeconds *int - audioTokenDetails *schemas.TranscriptionUsageInputTokenDetails - imageUsage *schemas.ImageUsage - imageSize string // e.g. "1024x1024", used for per-pixel pricing - imageQuality string // "low", "medium", "high", "auto" (gpt-image-1.5); empty = use base rate - videoSeconds *int - ocrProcessedPages *int - ocrIsAnnotated *bool - // containerIdentifierString, when non-empty, replaces the actual requested/resolved - // model names during pricing lookup. Used for request types whose cost is not - // tied to a specific model. Currently only used for container creates. - containerIdentifierString string - tier serviceTier -} - -// GetPricingEntryForModel returns the pricing data +// GetPricingEntryForModel returns any pricing entry for the model across +// known modes. Used by the inference handler to enrich list-models responses. func (mc *ModelCatalog) GetPricingEntryForModel(model string, provider schemas.ModelProvider) *PricingEntry { - mc.mu.RLock() - defer mc.mu.RUnlock() - // Check all modes - for _, mode := range []schemas.RequestType{ - schemas.TextCompletionRequest, - schemas.ChatCompletionRequest, - schemas.ResponsesRequest, - schemas.EmbeddingRequest, - schemas.RerankRequest, - schemas.SpeechRequest, - schemas.TranscriptionRequest, - schemas.ImageGenerationRequest, - schemas.ImageEditRequest, - schemas.ImageVariationRequest, - schemas.VideoGenerationRequest, - schemas.OCRRequest, - } { - key := makeKey(model, string(provider), normalizeRequestType(mode)) - pricing, ok := mc.pricingData[key] - if ok { - return convertTableModelPricingToPricingData(&pricing) - } - } - return nil + return mc.datasheet.GetPricingEntryForModel(model, provider) } -// CalculateCost calculates the cost of a Bifrost response. -// It handles all request types, cache debug billing, and tiered pricing. -// If scopes is nil, an empty PricingLookupScopes is used; global and provider-scoped -// overrides may still apply since the provider is derived from the response. +// CalculateCost computes the dollar cost for a Bifrost response. func (mc *ModelCatalog) CalculateCost(result *schemas.BifrostResponse, scopes *PricingLookupScopes) float64 { - if result == nil { - return 0 - } - - var s PricingLookupScopes - if scopes != nil { - s = *scopes - } - - // Handle semantic cache billing - cacheDebug := result.GetExtraFields().CacheDebug - if cacheDebug != nil { - return mc.calculateCostWithCache(result, cacheDebug, s) - } - - return mc.calculateBaseCost(result, s) -} - -// calculateCostWithCache handles cost calculation when semantic cache debug info is present. -func (mc *ModelCatalog) calculateCostWithCache(result *schemas.BifrostResponse, cacheDebug *schemas.BifrostCacheDebug, scopes PricingLookupScopes) float64 { - if cacheDebug.CacheHit { - // Direct cache hit — no LLM call, no cost - if cacheDebug.HitType != nil && *cacheDebug.HitType == "direct" { - return 0 - } - // Semantic cache hit — only the embedding lookup cost - if cacheDebug.ProviderUsed != nil && cacheDebug.ModelUsed != nil && cacheDebug.InputTokens != nil { - return mc.computeCacheEmbeddingCost(cacheDebug, scopes) - } - return 0 - } - - // Cache miss — full LLM cost + embedding lookup cost - baseCost := mc.calculateBaseCost(result, scopes) - embeddingCost := mc.computeCacheEmbeddingCost(cacheDebug, scopes) - return baseCost + embeddingCost -} - -// computeCacheEmbeddingCost calculates the embedding cost for a semantic cache lookup. -func (mc *ModelCatalog) computeCacheEmbeddingCost(cacheDebug *schemas.BifrostCacheDebug, scopes PricingLookupScopes) float64 { - if cacheDebug == nil || cacheDebug.ProviderUsed == nil || cacheDebug.ModelUsed == nil || cacheDebug.InputTokens == nil { - return 0 - } - if scopes.Provider == "" { - scopes.Provider = *cacheDebug.ProviderUsed - } - // Cache-debug pricing has only a single model identifier (whatever the - // cache recorded). Maps to RoutingInfo.Model — no alias resolution - // context exists for the cache-replayed request. - pricing := mc.resolvePricing(schemas.RoutingInfo{ - Provider: schemas.ModelProvider(*cacheDebug.ProviderUsed), - Model: *cacheDebug.ModelUsed, - }, schemas.EmbeddingRequest, scopes) - if pricing == nil { - return 0 - } - return float64(*cacheDebug.InputTokens) * tieredInputRate(pricing, *cacheDebug.InputTokens, serviceTier{}) -} - -// computeContainerCreationCost returns the cost for creating a container from an already-resolved pricing entry. -func computeContainerCreationCost(pricing *configstoreTables.TableModelPricing) float64 { - if pricing == nil || pricing.CodeInterpreterCostPerSession == nil { - return 0 - } - return *pricing.CodeInterpreterCostPerSession -} - -// calculateBaseCost extracts usage from the response and routes to the appropriate compute function. -func (mc *ModelCatalog) calculateBaseCost(result *schemas.BifrostResponse, scopes PricingLookupScopes) float64 { - extraFields := result.GetExtraFields() - if extraFields == nil { - return 0 - } - - // Read routing info populated by core.bifrost at request time. - // - // Backward-compat fallback: when the caller (e.g. LoggerPlugin's - // RecalculateCosts replaying logs written before RoutingInfo existed, - // or third-party plugins still on the legacy ExtraFields shape) leaves - // RoutingInfo empty, synthesise one from the deprecated triplet so - // pricing keeps working. Triggered only when RoutingInfo is fully - // unset — partial population is trusted as-is. - routingInfo := extraFields.RoutingInfo - if routingInfo.Provider == "" && routingInfo.Model == "" && routingInfo.ResolvedKeyAlias == nil { - routingInfo.Provider = extraFields.Provider - routingInfo.Model = extraFields.OriginalModelRequested - if r := extraFields.ResolvedModelUsed; r != "" && r != extraFields.OriginalModelRequested { - routingInfo.ResolvedKeyAlias = &schemas.ResolvedKeyAlias{ModelID: r} - } - } - requestType := extraFields.RequestType - - // Extract usage data from the response (passthrough and native paths unified) - input := extractCostInput(result) - - // If provider already computed cost, use it - if input.usage != nil && input.usage.Cost != nil && input.usage.Cost.TotalCost > 0 { - return input.usage.Cost.TotalCost - } - - // If no usage data at all, nothing to price - if input.usage == nil && input.audioSeconds == nil && input.audioTokenDetails == nil && input.imageUsage == nil && input.videoSeconds == nil && input.audioTextInputChars == 0 && input.ocrProcessedPages == nil && input.containerIdentifierString == "" { - return 0 - } - - if result.PassthroughResponse != nil { - // Infer request type from usage fields + path; passthrough bypasses stream normalization. - requestType = inferPassthroughRequestType(routingInfo.Provider, extraFields.PassthroughPath, result.PassthroughResponse.PassthroughUsage) - } else { - // Normalize stream request types to their base type for pricing lookup - requestType = normalizeStreamRequestType(requestType) - } - - // When a pricing model override is set (e.g. container creates always look - // up "container"), it replaces the lookup hierarchy entirely. Build a - // synthetic RoutingInfo that reuses Provider but pins the model fields to - // the container identifier — the lookup tries it as ModelName, the - // override key is the container identifier so per-container overrides - // stay addressable. - if input.containerIdentifierString != "" { - routingInfo = schemas.RoutingInfo{ - Provider: routingInfo.Provider, - Model: input.containerIdentifierString, - } - } - - pricing := mc.resolvePricing(routingInfo, requestType, scopes) - if pricing == nil { - return 0 - } - - // Route to the appropriate compute function - switch requestType { - case schemas.ChatCompletionRequest, schemas.TextCompletionRequest, schemas.ResponsesRequest, schemas.RealtimeRequest, schemas.CompactionRequest: - return computeTextCost(pricing, input.usage, input.tier) - case schemas.EmbeddingRequest: - return computeEmbeddingCost(pricing, input.usage, input.tier) - case schemas.RerankRequest: - return computeRerankCost(pricing, input.usage, input.tier) - case schemas.SpeechRequest: - return computeSpeechCost(pricing, input.usage, input.audioSeconds, input.audioTextInputChars, input.tier) - case schemas.TranscriptionRequest: - return computeTranscriptionCost(pricing, input.usage, input.audioSeconds, input.audioTokenDetails, input.tier) - case schemas.ImageGenerationRequest, schemas.ImageEditRequest, schemas.ImageVariationRequest: - return computeImageCost(pricing, input.imageUsage, input.imageSize, input.imageQuality, input.tier) - case schemas.VideoGenerationRequest, schemas.VideoRemixRequest: - return computeVideoCost(pricing, input.usage, input.videoSeconds, input.tier) - case schemas.OCRRequest: - return computeOCRCost(pricing, input.ocrProcessedPages, input.ocrIsAnnotated) - case schemas.ContainerCreateRequest: - return computeContainerCreationCost(pricing) - default: - return 0 - } -} - -// --------------------------------------------------------------------------- -// Usage extraction -// --------------------------------------------------------------------------- - -func extractCostInput(result *schemas.BifrostResponse) costInput { - var input costInput - - switch { - case result.PassthroughResponse != nil && result.PassthroughResponse.PassthroughUsage != nil: - return passthroughUsageToCostInput(result.PassthroughResponse.PassthroughUsage) - - case result.TextCompletionResponse != nil && result.TextCompletionResponse.Usage != nil: - input.usage = result.TextCompletionResponse.Usage - - case result.ChatResponse != nil && result.ChatResponse.Usage != nil: - input.usage = result.ChatResponse.Usage - input.tier = tierFromString(result.ChatResponse.ServiceTier) - - case result.ResponsesResponse != nil && result.ResponsesResponse.Usage != nil: - input.usage = responsesUsageToBifrostUsage(result.ResponsesResponse.Usage) - input.tier = tierFromString(result.ResponsesResponse.ServiceTier) - - case result.CompactionResponse != nil && result.CompactionResponse.Usage != nil: - input.usage = responsesUsageToBifrostUsage(result.CompactionResponse.Usage) - - case result.ResponsesStreamResponse != nil && result.ResponsesStreamResponse.Response != nil && result.ResponsesStreamResponse.Response.Usage != nil: - input.usage = responsesUsageToBifrostUsage(result.ResponsesStreamResponse.Response.Usage) - input.tier = tierFromString(result.ResponsesStreamResponse.Response.ServiceTier) - - case result.EmbeddingResponse != nil && result.EmbeddingResponse.Usage != nil: - input.usage = result.EmbeddingResponse.Usage - - case result.RerankResponse != nil && result.RerankResponse.Usage != nil: - input.usage = result.RerankResponse.Usage - - case result.SpeechResponse != nil && result.SpeechResponse.Usage != nil: - input.usage = speechUsageToBifrostUsage(result.SpeechResponse.Usage) - input.audioTextInputChars = result.SpeechResponse.Usage.InputChars - - case result.SpeechStreamResponse != nil && result.SpeechStreamResponse.Usage != nil: - input.usage = speechUsageToBifrostUsage(result.SpeechStreamResponse.Usage) - input.audioTextInputChars = result.SpeechStreamResponse.Usage.InputChars - - case result.TranscriptionResponse != nil && result.TranscriptionResponse.Usage != nil: - input.usage, input.audioSeconds, input.audioTokenDetails = extractTranscriptionUsage(result.TranscriptionResponse.Usage) - - case result.TranscriptionStreamResponse != nil && result.TranscriptionStreamResponse.Usage != nil: - input.usage, input.audioSeconds, input.audioTokenDetails = extractTranscriptionUsage(result.TranscriptionStreamResponse.Usage) - - case result.ImageGenerationResponse != nil: - if result.ImageGenerationResponse.Usage != nil { - input.imageUsage = result.ImageGenerationResponse.Usage - } else { - // No usage data but response exists — default to empty so per-image pricing can apply - input.imageUsage = &schemas.ImageUsage{} - } - populateOutputImageCount(input.imageUsage, len(result.ImageGenerationResponse.Data)) - if result.ImageGenerationResponse.ImageGenerationResponseParameters != nil { - input.imageSize = result.ImageGenerationResponse.ImageGenerationResponseParameters.Size - input.imageQuality = result.ImageGenerationResponse.ImageGenerationResponseParameters.Quality - } - - case result.ImageGenerationStreamResponse != nil: - if result.ImageGenerationStreamResponse.Usage != nil { - input.imageUsage = result.ImageGenerationStreamResponse.Usage - } else { - input.imageUsage = &schemas.ImageUsage{} - } - input.imageSize = result.ImageGenerationStreamResponse.Size - input.imageQuality = result.ImageGenerationStreamResponse.Quality - - case result.VideoGenerationResponse != nil && result.VideoGenerationResponse.Seconds != nil: - seconds, err := strconv.Atoi(*result.VideoGenerationResponse.Seconds) - if err == nil { - input.videoSeconds = &seconds - } - - case result.OCRResponse != nil: - pages := len(result.OCRResponse.Pages) - if result.OCRResponse.UsageInfo != nil && result.OCRResponse.UsageInfo.PagesProcessed > 0 { - pages = result.OCRResponse.UsageInfo.PagesProcessed - } - input.ocrProcessedPages = &pages - isAnnotated := result.OCRResponse.DocumentAnnotation != nil && *result.OCRResponse.DocumentAnnotation != "" - input.ocrIsAnnotated = &isAnnotated - - case result.ContainerCreateResponse != nil: - if memLimit := result.ContainerCreateResponse.MemoryLimit; memLimit != "" { - input.containerIdentifierString = "container-" + memLimit - } else { - input.containerIdentifierString = "container" - } - } - - return input + return mc.datasheet.CalculateCost(result, (*datasheet.LookupScopes)(scopes)) } -func responsesUsageToBifrostUsage(u *schemas.ResponsesResponseUsage) *schemas.BifrostLLMUsage { - usage := &schemas.BifrostLLMUsage{ - PromptTokens: u.InputTokens, - CompletionTokens: u.OutputTokens, - TotalTokens: u.TotalTokens, - Cost: u.Cost, - } - // Map token details for cache and search query pricing - if u.InputTokensDetails != nil { - usage.PromptTokensDetails = &schemas.ChatPromptTokensDetails{ - TextTokens: u.InputTokensDetails.TextTokens, - AudioTokens: u.InputTokensDetails.AudioTokens, - ImageTokens: u.InputTokensDetails.ImageTokens, - CachedReadTokens: u.InputTokensDetails.CachedReadTokens, - CachedWriteTokens: u.InputTokensDetails.CachedWriteTokens, - CachedWriteTokenDetails: u.InputTokensDetails.CachedWriteTokenDetails, - } - } - if u.OutputTokensDetails != nil { - usage.CompletionTokensDetails = &schemas.ChatCompletionTokensDetails{ - ReasoningTokens: u.OutputTokensDetails.ReasoningTokens, - AudioTokens: u.OutputTokensDetails.AudioTokens, - } - if u.OutputTokensDetails.NumSearchQueries != nil { - usage.CompletionTokensDetails.NumSearchQueries = u.OutputTokensDetails.NumSearchQueries - } - } - return usage -} - -func speechUsageToBifrostUsage(u *schemas.SpeechUsage) *schemas.BifrostLLMUsage { - return &schemas.BifrostLLMUsage{ - PromptTokens: u.InputTokens, - CompletionTokens: u.OutputTokens, - TotalTokens: u.TotalTokens, - } -} - -func extractTranscriptionUsage(u *schemas.TranscriptionUsage) (*schemas.BifrostLLMUsage, *int, *schemas.TranscriptionUsageInputTokenDetails) { - usage := &schemas.BifrostLLMUsage{} - if u.InputTokens != nil { - usage.PromptTokens = *u.InputTokens - } - if u.OutputTokens != nil { - usage.CompletionTokens = *u.OutputTokens - } - if u.TotalTokens != nil { - usage.TotalTokens = *u.TotalTokens - } else { - usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens - } - - var audioTokenDetails *schemas.TranscriptionUsageInputTokenDetails - if u.InputTokenDetails != nil { - audioTokenDetails = &schemas.TranscriptionUsageInputTokenDetails{ - AudioTokens: u.InputTokenDetails.AudioTokens, - TextTokens: u.InputTokenDetails.TextTokens, - } - } - - return usage, u.Seconds, audioTokenDetails -} - -// --------------------------------------------------------------------------- -// Per-request-type cost computation -// --------------------------------------------------------------------------- - -// computeTextCost handles chat, text completion, and responses requests. -func computeTextCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, tier serviceTier) float64 { - if usage == nil { - return 0 - } - - totalTokens := usage.TotalTokens - promptTokens := usage.PromptTokens - completionTokens := usage.CompletionTokens - - // Extract cached token counts - cachedReadTokens := 0 - cachedWriteTokens := 0 - cachedWriteTokensAbove1hr := 0 - if usage.PromptTokensDetails != nil { - cachedReadTokens = usage.PromptTokensDetails.CachedReadTokens - cachedWriteTokens = usage.PromptTokensDetails.CachedWriteTokens - if usage.PromptTokensDetails.CachedWriteTokenDetails != nil { - cachedWriteTokensAbove1hr = usage.PromptTokensDetails.CachedWriteTokenDetails.CachedWriteTokens1h - } - } - - inputRate := tieredInputRate(pricing, totalTokens, tier) - outputRate := tieredOutputRate(pricing, totalTokens, tier) - cacheReadInputRate := tieredCacheReadInputTokenRate(pricing, totalTokens, tier) - cacheCreationInputRate := tieredCacheCreationInputTokenRate(pricing, totalTokens, tier) - cacheCreationInputAbove1hrInputRate := tieredCacheCreationInputAbove1hrTokenRate(pricing, totalTokens, tier) - - // Clamp cached token counts to avoid negative billing on malformed provider payloads - if cachedReadTokens > promptTokens { - cachedReadTokens = promptTokens - } - if cachedWriteTokens > promptTokens-cachedReadTokens { - cachedWriteTokens = promptTokens - cachedReadTokens - } - // Should not happen, but just in case - if cachedWriteTokensAbove1hr > cachedWriteTokens { - cachedWriteTokensAbove1hr = cachedWriteTokens - } - - // Input cost: non-cached tokens at regular rate - nonCachedPrompt := promptTokens - cachedReadTokens - cachedWriteTokens - inputCost := float64(nonCachedPrompt) * inputRate - - // Add cached prompt tokens at cache read rate - if cachedReadTokens > 0 { - inputCost += float64(cachedReadTokens) * cacheReadInputRate - } - - // Add cached write tokens at cache creation rate - if cachedWriteTokens > 0 { - if cachedWriteTokensAbove1hr > 0 { - inputCost += float64(cachedWriteTokensAbove1hr) * cacheCreationInputAbove1hrInputRate - } - inputCost += float64(cachedWriteTokens-cachedWriteTokensAbove1hr) * cacheCreationInputRate - } - - outputCost := float64(completionTokens) * outputRate - - // Audio token cost: when token details include audio tokens, price them - // at the dedicated audio rate and subtract from the text token costs above. - // Realtime and audio-enabled chat models report audio tokens in details. - audioCost := 0.0 - inputAudioTokens := 0 - outputAudioTokens := 0 - if usage.PromptTokensDetails != nil { - inputAudioTokens = usage.PromptTokensDetails.AudioTokens - } - if usage.CompletionTokensDetails != nil { - outputAudioTokens = usage.CompletionTokensDetails.AudioTokens - } - if inputAudioTokens < 0 { - inputAudioTokens = 0 - } else if inputAudioTokens > promptTokens { - inputAudioTokens = promptTokens - } - if outputAudioTokens < 0 { - outputAudioTokens = 0 - } else if outputAudioTokens > completionTokens { - outputAudioTokens = completionTokens - } - if inputAudioTokens > 0 && pricing.InputCostPerAudioToken != nil { - // Subtract audio tokens charged at text rate, add at audio rate. - audioCost += float64(inputAudioTokens) * (*pricing.InputCostPerAudioToken - inputRate) - } - if outputAudioTokens > 0 && pricing.OutputCostPerAudioToken != nil { - audioCost += float64(outputAudioTokens) * (*pricing.OutputCostPerAudioToken - outputRate) - } - - // Search query cost - searchCost := 0.0 - if pricing.SearchContextCostPerQuery != nil && usage.CompletionTokensDetails != nil && usage.CompletionTokensDetails.NumSearchQueries != nil { - searchCost = float64(*usage.CompletionTokensDetails.NumSearchQueries) * *pricing.SearchContextCostPerQuery - } - - return inputCost + outputCost + audioCost + searchCost -} - -// computeEmbeddingCost handles embedding requests (input-only). -func computeEmbeddingCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, tier serviceTier) float64 { - if usage == nil { - return 0 - } - return float64(usage.PromptTokens) * tieredInputRate(pricing, usage.TotalTokens, tier) -} - -// computeRerankCost handles rerank requests. -func computeRerankCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, tier serviceTier) float64 { - if usage == nil { - return 0 - } - inputCost := float64(usage.PromptTokens) * tieredInputRate(pricing, usage.TotalTokens, tier) - outputCost := float64(usage.CompletionTokens) * tieredOutputRate(pricing, usage.TotalTokens, tier) - - searchCost := 0.0 - if pricing.SearchContextCostPerQuery != nil && usage.CompletionTokensDetails != nil && usage.CompletionTokensDetails.NumSearchQueries != nil { - searchCost = float64(*usage.CompletionTokensDetails.NumSearchQueries) * *pricing.SearchContextCostPerQuery - } - - return inputCost + outputCost + searchCost -} - -// computeSpeechCost handles speech (TTS) requests. -// Input is text (PromptTokens), output is audio (CompletionTokens). -// -// Per-character pricing (InputCostPerCharacter) is used as first-class support for TTS/audio -// models — providers such as OpenAI TTS, ElevenLabs, and AWS Polly bill per character of -// input text rather than per token. PromptTokens from usage is treated as the character count -// since TTS providers report their billable unit in that field. -// Output falls back to per-second duration when no audio token rate is configured. -func computeSpeechCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, audioSeconds *int, audioTextInputChars int, tier serviceTier) float64 { - totalTokens := safeTotalTokens(usage) - - // Input: per-character rate takes precedence for TTS/audio models - inputCost := 0.0 - if audioTextInputChars > 0 { - if pricing.InputCostPerCharacter != nil { - inputCost = float64(audioTextInputChars) * *pricing.InputCostPerCharacter - } else { - inputCost = float64(audioTextInputChars) * tieredInputRate(pricing, totalTokens, tier) - } - } else if usage != nil && usage.PromptTokens > 0 { - inputCost = float64(usage.PromptTokens) * tieredInputRate(pricing, totalTokens, tier) - } - - // Output: audio tokens first, then per-second fallback - outputCost := computeAudioOutputCost(pricing, usage, audioSeconds, totalTokens, tier) - - return inputCost + outputCost -} - -// computeTranscriptionCost handles transcription (STT) requests. -// Input is audio, output is text (CompletionTokens). -// Input and output are calculated independently — tokens first, then per-second fallback. -func computeTranscriptionCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, audioSeconds *int, audioTokenDetails *schemas.TranscriptionUsageInputTokenDetails, tier serviceTier) float64 { - totalTokens := safeTotalTokens(usage) - - // Input: audio tokens/details first, then per-second fallback - inputCost := computeAudioInputCost(pricing, usage, audioSeconds, audioTokenDetails, totalTokens, tier) - - // Output: text tokens - outputCost := 0.0 - if usage != nil && usage.CompletionTokens > 0 { - outputCost = float64(usage.CompletionTokens) * tieredOutputRate(pricing, totalTokens, tier) - } - - return inputCost + outputCost -} - -// computeAudioInputCost calculates input cost for audio: audio token details first, -// then generic input tokens, then per-second duration fallback. -func computeAudioInputCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, audioSeconds *int, audioTokenDetails *schemas.TranscriptionUsageInputTokenDetails, totalTokens int, tier serviceTier) float64 { - // Audio token detail pricing (audio + text token breakdown) - if audioTokenDetails != nil && (audioTokenDetails.AudioTokens > 0 || audioTokenDetails.TextTokens > 0) { - return float64(audioTokenDetails.AudioTokens)*tieredAudioTokenInputRate(pricing, totalTokens, tier) + - float64(audioTokenDetails.TextTokens)*tieredInputRate(pricing, totalTokens, tier) - } - - // Generic input tokens - if usage != nil && usage.PromptTokens > 0 { - return float64(usage.PromptTokens) * tieredInputRate(pricing, totalTokens, tier) - } - - // Per-second duration fallback - if audioSeconds != nil && *audioSeconds > 0 { - if rate := tieredAudioInputPerSecondRate(pricing, totalTokens); rate > 0 { - return float64(*audioSeconds) * rate - } - } - - return 0 -} - -// computeAudioOutputCost calculates output cost for audio: audio tokens first, -// then generic output tokens, then per-second duration fallback. -func computeAudioOutputCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, audioSeconds *int, totalTokens int, tier serviceTier) float64 { - // Audio-specific output tokens - if usage != nil && usage.CompletionTokens > 0 { - return float64(usage.CompletionTokens) * tieredAudioTokenOutputRate(pricing, totalTokens, tier) - } - - // Per-second duration fallback - if audioSeconds != nil && *audioSeconds > 0 { - if pricing.OutputCostPerSecond != nil { - return float64(*audioSeconds) * *pricing.OutputCostPerSecond - } - } - - return 0 -} - -// computeImageCost handles image generation requests. -// Input and output are calculated independently — each tries token-based pricing first, -// then per-pixel pricing, falling back to per-image count pricing. -// imageQuality must be one of "low", "medium", "high", "auto" to use quality-specific rates; other values use base rates. -func computeImageCost(pricing *configstoreTables.TableModelPricing, imageUsage *schemas.ImageUsage, imageSize string, imageQuality string, tier serviceTier) float64 { - if imageUsage == nil { - return 0 - } - - totalTokens := imageUsage.TotalTokens - pixels := parseImagePixels(imageSize) - inputCost := computeImageInputCost(pricing, imageUsage, totalTokens, pixels, tier) - outputCost := computeImageOutputCost(pricing, imageUsage, totalTokens, pixels, imageQuality, tier) - - return inputCost + outputCost -} - -// computeImageInputCost calculates input cost: tokens first, then per-pixel, then per-image count fallback. -func computeImageInputCost(pricing *configstoreTables.TableModelPricing, imageUsage *schemas.ImageUsage, totalTokens int, pixels int, tier serviceTier) float64 { - // Try token-based pricing first - var inputTextTokens, inputImageTokens int - if imageUsage.InputTokensDetails != nil { - inputImageTokens = imageUsage.InputTokensDetails.ImageTokens - inputTextTokens = imageUsage.InputTokensDetails.TextTokens - } else { - inputTextTokens = imageUsage.InputTokens - } - - if inputTextTokens > 0 || inputImageTokens > 0 { - return float64(inputTextTokens)*tieredInputRate(pricing, totalTokens, tier) + - float64(inputImageTokens)*tieredImageInputRate(pricing, totalTokens, tier) - } - - // Per-pixel pricing fallback - if pricing.InputCostPerPixel != nil && pixels > 0 && imageUsage.NumInputImages > 0 { - return float64(pixels*imageUsage.NumInputImages) * *pricing.InputCostPerPixel - } - - // Fall back to per-image count pricing - if pricing.InputCostPerImage != nil && imageUsage.NumInputImages > 0 { - return float64(imageUsage.NumInputImages) * *pricing.InputCostPerImage - } - - return 0 -} - -// computeImageOutputCost calculates output cost: tokens first, then per-pixel, then per-image count fallback. -// imageQuality: "low", "medium", "high", "auto" use quality-specific rates when available; other values use base/size-tier rates. -func computeImageOutputCost(pricing *configstoreTables.TableModelPricing, imageUsage *schemas.ImageUsage, totalTokens int, pixels int, imageQuality string, tier serviceTier) float64 { - // Try token-based pricing first - var outputTextTokens, outputImageTokens int - if imageUsage.OutputTokensDetails != nil { - outputImageTokens = imageUsage.OutputTokensDetails.ImageTokens - outputTextTokens = imageUsage.OutputTokensDetails.TextTokens - } else { - outputImageTokens = imageUsage.OutputTokens - } - - if outputTextTokens > 0 || outputImageTokens > 0 { - return float64(outputTextTokens)*tieredOutputRate(pricing, totalTokens, tier) + - float64(outputImageTokens)*tieredImageOutputRate(pricing, totalTokens, tier) - } - - // Per-pixel pricing fallback - if pricing.OutputCostPerPixel != nil && pixels > 0 { - numOutputImages := 1 - if imageUsage.OutputTokensDetails != nil && imageUsage.OutputTokensDetails.NImages > 0 { - numOutputImages = imageUsage.OutputTokensDetails.NImages - } - return float64(pixels*numOutputImages) * *pricing.OutputCostPerPixel - } - - // Fall back to per-image count pricing with size-tier selection - // TODO: handle premium image flag when it becomes available in imageUsage - numOutputImages := 1 - if imageUsage.OutputTokensDetails != nil && imageUsage.OutputTokensDetails.NImages > 0 { - numOutputImages = imageUsage.OutputTokensDetails.NImages - } - var perImageRate *float64 - q := imageQuality - if q == "" { - q = "auto" - } - switch q { - case "low": - if pricing.OutputCostPerImageLowQuality != nil { - perImageRate = pricing.OutputCostPerImageLowQuality - } - case "medium": - if pricing.OutputCostPerImageMediumQuality != nil { - perImageRate = pricing.OutputCostPerImageMediumQuality - } - case "high": - if pricing.OutputCostPerImageHighQuality != nil { - perImageRate = pricing.OutputCostPerImageHighQuality - } - case "auto": - if pricing.OutputCostPerImageAutoQuality != nil { - perImageRate = pricing.OutputCostPerImageAutoQuality - } - } - if perImageRate == nil { - const pixels512x512 = 512 * 512 - const pixels1024x1024 = 1024 * 1024 - const pixels2048x2048 = 2048 * 2048 - const pixels4096x4096 = 4096 * 4096 - switch { - case pixels >= pixels4096x4096 && pricing.OutputCostPerImageAbove4096x4096Pixels != nil: - perImageRate = pricing.OutputCostPerImageAbove4096x4096Pixels - case pixels >= pixels2048x2048 && pricing.OutputCostPerImageAbove2048x2048Pixels != nil: - perImageRate = pricing.OutputCostPerImageAbove2048x2048Pixels - case pixels >= pixels1024x1024 && pricing.OutputCostPerImageAbove1024x1024Pixels != nil: - perImageRate = pricing.OutputCostPerImageAbove1024x1024Pixels - case pixels >= pixels512x512 && pricing.OutputCostPerImageAbove512x512Pixels != nil: - perImageRate = pricing.OutputCostPerImageAbove512x512Pixels - default: - perImageRate = pricing.OutputCostPerImage - } - } - if perImageRate != nil { - return float64(numOutputImages) * *perImageRate - } - - return 0 -} - -// computeVideoCost handles video generation requests. -// Input and output are calculated independently — tokens first, then per-second fallback. -func computeVideoCost(pricing *configstoreTables.TableModelPricing, usage *schemas.BifrostLLMUsage, videoSeconds *int, tier serviceTier) float64 { - totalTokens := safeTotalTokens(usage) - - // Input: text prompt tokens first, then per-second fallback - inputCost := 0.0 - if usage != nil && usage.PromptTokens > 0 { - inputCost = float64(usage.PromptTokens) * tieredInputRate(pricing, totalTokens, tier) - } else if videoSeconds != nil && *videoSeconds > 0 { - if rate := tieredVideoInputPerSecondRate(pricing, totalTokens); rate > 0 { - inputCost = float64(*videoSeconds) * rate - } - } - - // Output: completion tokens first, then per-second fallback - outputCost := 0.0 - if usage != nil && usage.CompletionTokens > 0 { - outputCost = float64(usage.CompletionTokens) * tieredOutputRate(pricing, totalTokens, tier) - } else if videoSeconds != nil && *videoSeconds > 0 { - if pricing.OutputCostPerVideoPerSecond != nil { - outputCost = float64(*videoSeconds) * *pricing.OutputCostPerVideoPerSecond - } else if pricing.OutputCostPerSecond != nil { - outputCost = float64(*videoSeconds) * *pricing.OutputCostPerSecond - } - } - - return inputCost + outputCost -} - -// computeOCRCost handles OCR requests, billing per page processed. -// ocr_cost_per_page covers base processing; annotation_cost_per_page is added when set. -func computeOCRCost(pricing *configstoreTables.TableModelPricing, ocrProcessedPages *int, ocrIsAnnotated *bool) float64 { - if ocrProcessedPages == nil { - return 0 - } - pages := float64(*ocrProcessedPages) - cost := 0.0 - if pricing.OCRCostPerPage != nil { - cost += pages * *pricing.OCRCostPerPage - } - if ocrIsAnnotated != nil && *ocrIsAnnotated && pricing.AnnotationCostPerPage != nil { - cost += pages * *pricing.AnnotationCostPerPage - } - return cost -} - -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -// tierFromString constructs a serviceTier from an OpenAI service_tier response value. -func tierFromString(s *schemas.BifrostServiceTier) serviceTier { - if s == nil { - return serviceTier{} - } - switch *s { - case schemas.BifrostServiceTierPriority: - return serviceTier{isPriority: true} - case schemas.BifrostServiceTierFlex: - return serviceTier{isFlex: true} - default: - return serviceTier{} - } -} - -// tieredInputRate returns the effective per-token input rate based on total token count. -// Flex applies a flat rate. Priority-specific tier rates are preferred where available. -func tieredInputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { - if tier.isFlex && pricing.InputCostPerTokenFlex != nil { - return *pricing.InputCostPerTokenFlex - } - if totalTokens > TokenTierAbove272K { - if tier.isPriority && pricing.InputCostPerTokenAbove272kTokensPriority != nil { - return *pricing.InputCostPerTokenAbove272kTokensPriority - } - if pricing.InputCostPerTokenAbove272kTokens != nil { - return *pricing.InputCostPerTokenAbove272kTokens - } - } - if totalTokens > TokenTierAbove200K { - if tier.isPriority && pricing.InputCostPerTokenAbove200kTokensPriority != nil { - return *pricing.InputCostPerTokenAbove200kTokensPriority - } - if pricing.InputCostPerTokenAbove200kTokens != nil { - return *pricing.InputCostPerTokenAbove200kTokens - } - } - if totalTokens > TokenTierAbove128K && pricing.InputCostPerTokenAbove128kTokens != nil { - return *pricing.InputCostPerTokenAbove128kTokens - } - if tier.isPriority && pricing.InputCostPerTokenPriority != nil { - return *pricing.InputCostPerTokenPriority - } - if pricing.InputCostPerToken != nil { - return *pricing.InputCostPerToken - } - return 0 -} - -// tieredOutputRate returns the effective per-token output rate based on total token count. -// Flex applies a flat rate. Priority-specific tier rates are preferred where available. -func tieredOutputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { - if tier.isFlex && pricing.OutputCostPerTokenFlex != nil { - return *pricing.OutputCostPerTokenFlex - } - if totalTokens > TokenTierAbove272K { - if tier.isPriority && pricing.OutputCostPerTokenAbove272kTokensPriority != nil { - return *pricing.OutputCostPerTokenAbove272kTokensPriority - } - if pricing.OutputCostPerTokenAbove272kTokens != nil { - return *pricing.OutputCostPerTokenAbove272kTokens - } - } - if totalTokens > TokenTierAbove200K { - if tier.isPriority && pricing.OutputCostPerTokenAbove200kTokensPriority != nil { - return *pricing.OutputCostPerTokenAbove200kTokensPriority - } - if pricing.OutputCostPerTokenAbove200kTokens != nil { - return *pricing.OutputCostPerTokenAbove200kTokens - } - } - if totalTokens > TokenTierAbove128K && pricing.OutputCostPerTokenAbove128kTokens != nil { - return *pricing.OutputCostPerTokenAbove128kTokens - } - - if tier.isPriority && pricing.OutputCostPerTokenPriority != nil { - return *pricing.OutputCostPerTokenPriority - } - - if pricing.OutputCostPerToken != nil { - return *pricing.OutputCostPerToken - } - - return 0 -} - -// tieredImageInputRate returns the effective rate for image tokens on the input side. -// Falls back to the general tieredInputRate when no image-specific rate is configured. -func tieredImageInputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { - if totalTokens > TokenTierAbove128K && pricing.InputCostPerImageAbove128kTokens != nil { - return *pricing.InputCostPerImageAbove128kTokens - } - if pricing.InputCostPerImageToken != nil { - return *pricing.InputCostPerImageToken - } - return tieredInputRate(pricing, totalTokens, tier) -} - -// tieredImageOutputRate returns the effective rate for image tokens on the output side. -// Falls back to the general tieredOutputRate when no image-specific rate is configured. -func tieredImageOutputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { - if pricing.OutputCostPerImageToken != nil { - return *pricing.OutputCostPerImageToken - } - return tieredOutputRate(pricing, totalTokens, tier) -} - -// tieredAudioInputPerSecondRate returns the effective per-second rate for audio input. -func tieredAudioInputPerSecondRate(pricing *configstoreTables.TableModelPricing, totalTokens int) float64 { - if totalTokens > TokenTierAbove128K && pricing.InputCostPerAudioPerSecondAbove128kTokens != nil { - return *pricing.InputCostPerAudioPerSecondAbove128kTokens - } - if pricing.InputCostPerAudioPerSecond != nil { - return *pricing.InputCostPerAudioPerSecond - } - if pricing.InputCostPerSecond != nil { - return *pricing.InputCostPerSecond - } - return 0 -} - -// tieredVideoInputPerSecondRate returns the effective per-second rate for video input. -func tieredVideoInputPerSecondRate(pricing *configstoreTables.TableModelPricing, totalTokens int) float64 { - if totalTokens > TokenTierAbove128K && pricing.InputCostPerVideoPerSecondAbove128kTokens != nil { - return *pricing.InputCostPerVideoPerSecondAbove128kTokens - } - if pricing.InputCostPerVideoPerSecond != nil { - return *pricing.InputCostPerVideoPerSecond - } - return 0 -} - -// tieredAudioTokenInputRate returns the effective per-token rate for audio input tokens. -// Falls back to the general tieredInputRate when no audio-specific rate is configured. -func tieredAudioTokenInputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { - if pricing.InputCostPerAudioToken != nil { - return *pricing.InputCostPerAudioToken - } - return tieredInputRate(pricing, totalTokens, tier) -} - -// tieredAudioTokenOutputRate returns the effective per-token rate for audio output tokens. -// Falls back to the general tieredOutputRate when no audio-specific rate is configured. -func tieredAudioTokenOutputRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { - if pricing.OutputCostPerAudioToken != nil { - return *pricing.OutputCostPerAudioToken - } - return tieredOutputRate(pricing, totalTokens, tier) -} - -func tieredCacheReadInputTokenRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { - if tier.isFlex && pricing.CacheReadInputTokenCostFlex != nil { - return *pricing.CacheReadInputTokenCostFlex - } - if totalTokens > TokenTierAbove272K { - if tier.isPriority && pricing.CacheReadInputTokenCostAbove272kTokensPriority != nil { - return *pricing.CacheReadInputTokenCostAbove272kTokensPriority - } - if pricing.CacheReadInputTokenCostAbove272kTokens != nil { - return *pricing.CacheReadInputTokenCostAbove272kTokens - } - } - if totalTokens > TokenTierAbove200K { - if tier.isPriority && pricing.CacheReadInputTokenCostAbove200kTokensPriority != nil { - return *pricing.CacheReadInputTokenCostAbove200kTokensPriority - } - if pricing.CacheReadInputTokenCostAbove200kTokens != nil { - return *pricing.CacheReadInputTokenCostAbove200kTokens - } - } - if tier.isPriority && pricing.CacheReadInputTokenCostPriority != nil { - return *pricing.CacheReadInputTokenCostPriority - } - if pricing.CacheReadInputTokenCost != nil { - return *pricing.CacheReadInputTokenCost - } - return tieredInputRate(pricing, totalTokens, tier) -} - -// Note: flex tier is not checked here because cache creation is not a concept in -// OpenAI's pricing model (the only provider that uses flex tier). Only cache read -// has a flex-specific rate. -func tieredCacheCreationInputTokenRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { - if totalTokens > TokenTierAbove200K && pricing.CacheCreationInputTokenCostAbove200kTokens != nil { - return *pricing.CacheCreationInputTokenCostAbove200kTokens - } - if pricing.CacheCreationInputTokenCost != nil { - return *pricing.CacheCreationInputTokenCost - } - return tieredInputRate(pricing, totalTokens, tier) -} - -func tieredCacheCreationInputAbove1hrTokenRate(pricing *configstoreTables.TableModelPricing, totalTokens int, tier serviceTier) float64 { - if totalTokens > TokenTierAbove200K && pricing.CacheCreationInputTokenCostAbove1hrAbove200kTokens != nil { - return *pricing.CacheCreationInputTokenCostAbove1hrAbove200kTokens - } - if pricing.CacheCreationInputTokenCostAbove1hr != nil { - return *pricing.CacheCreationInputTokenCostAbove1hr - } - return tieredCacheCreationInputTokenRate(pricing, totalTokens, tier) -} - -func safeTotalTokens(usage *schemas.BifrostLLMUsage) int { - if usage == nil { - return 0 - } - return usage.TotalTokens -} - -// parseImagePixels parses a size string like "1024x1024" into total pixel count. -// Returns 0 if the size string is empty or malformed. -func parseImagePixels(size string) int { - if size == "" { - return 0 - } - parts := strings.SplitN(size, "x", 2) - if len(parts) != 2 { - return 0 - } - w, err := strconv.Atoi(parts[0]) - if err != nil || w <= 0 { - return 0 - } - h, err := strconv.Atoi(parts[1]) - if err != nil || h <= 0 { - return 0 - } - return w * h -} - -// populateOutputImageCount sets the output image count on ImageUsage from len(Data) -// when OutputTokensDetails.NImages is not already populated. -func populateOutputImageCount(imageUsage *schemas.ImageUsage, dataLen int) { - if imageUsage == nil || dataLen == 0 { - return - } - if imageUsage.OutputTokensDetails == nil { - imageUsage.OutputTokensDetails = &schemas.ImageTokenDetails{} - } - if imageUsage.OutputTokensDetails.NImages == 0 { - imageUsage.OutputTokensDetails.NImages = dataLen - } -} - -// --------------------------------------------------------------------------- -// Pricing resolution -// --------------------------------------------------------------------------- - -// resolvePricing resolves the pricing entry for a request directly from the -// RoutingInfo populated on the response/error by core.bifrost at request time. -// -// Lookup precedence — AliasModelName → AliasModelID → ModelName. Each -// non-empty candidate is tried against the base catalog in order; the first -// hit wins. -// -// - AliasModelName (RoutingInfo.ResolvedKeyAlias.ModelName) is the canonical -// model name the admin tagged on the matched alias. Catches the -// opaque-deployment-ID case where the wire model wouldn't hit the catalog -// on its own. -// - AliasModelID (RoutingInfo.ResolvedKeyAlias.ModelID) is the wire model -// when an alias matched. nil/empty otherwise. -// - ModelName (RoutingInfo.Model) is the model string the caller sent — the -// alias key when an alias matched, or the raw user input when none did. -// -// Overrides are applied keyed by the wire model (AliasModelID when an alias -// matched, otherwise ModelName) so per-deployment override pricing stays -// addressable in either flow. -func (mc *ModelCatalog) resolvePricing(routingInfo schemas.RoutingInfo, requestType schemas.RequestType, scopes PricingLookupScopes) *configstoreTables.TableModelPricing { - provider := string(routingInfo.Provider) - var aliasModelID, aliasModelName string - if rka := routingInfo.ResolvedKeyAlias; rka != nil { - aliasModelID = rka.ModelID - if rka.ModelName != nil { - aliasModelName = *rka.ModelName - } - } - overrideKey := aliasModelID - if overrideKey == "" { - overrideKey = routingInfo.Model - } - mc.logger.Debug("looking up pricing for wire model %s and provider %s of request type %s", overrideKey, provider, normalizeRequestType(requestType)) - - if scopes.Provider == "" { - scopes.Provider = provider - } - - for _, candidate := range []string{aliasModelName, aliasModelID, routingInfo.Model} { - if candidate == "" { - continue - } - base, exists := mc.getBasePricing(candidate, provider, requestType) - if exists && base != nil { - result, _ := mc.applyPricingOverrides(overrideKey, requestType, *base, scopes) - return &result - } - mc.logger.Debug("pricing not found for %s, trying next candidate", candidate) - } - - // No base catalog entry found; still try overrides in case the user defined - // override-only pricing for a model not in the built-in catalog. - mc.logger.Debug("pricing not found for any candidate (provider %s), trying override-only pricing keyed by %s", provider, overrideKey) - result, applied := mc.applyPricingOverrides(overrideKey, requestType, configstoreTables.TableModelPricing{}, scopes) - if applied { - return &result - } - mc.logger.Debug("no pricing found for wire model %s and provider %s, skipping cost calculation", overrideKey, provider) - return nil -} - -// getBasePricing looks up catalog pricing for the given model, provider, and request type. -// It applies a provider-specific fallback chain when an exact match is not found: -// -// - Gemini: retries under the "vertex" provider, then falls back to chat mode for Responses requests. -// - Vertex: strips the "provider/model" prefix and retries, then falls back to chat mode for Responses requests. -// - Bedrock: prepends the "anthropic." namespace for Claude models, then falls back to chat mode for Responses requests. -// - All providers: for Responses/ResponsesStream requests, retries the lookup in chat mode. -// - All providers: for ImageEdit/ImageVariation requests, retries the lookup in image-generation mode. -// -// The method acquires a read lock for the duration of the lookup. -// -// Input: model — exact model name to look up. -// -// provider — provider identifier (e.g. "openai", "anthropic"). -// requestType — the request type used to derive the pricing mode. -// -// Output: TableModelPricing — the matched pricing row (zero value when not found). -// -// bool — true when a pricing entry was found, false otherwise. -func (mc *ModelCatalog) getBasePricing(model, provider string, requestType schemas.RequestType) (*configstoreTables.TableModelPricing, bool) { - mc.mu.RLock() - defer mc.mu.RUnlock() - - mode := normalizeRequestType(requestType) - - pricing, ok := mc.pricingData[makeKey(model, provider, mode)] - if ok { - return &pricing, true - } - - // Lookup in vertex if gemini not found - if provider == string(schemas.Gemini) { - mc.logger.Debug("primary lookup failed, trying vertex provider for the same model") - pricing, ok = mc.pricingData[makeKey(model, "vertex", mode)] - if ok { - return &pricing, true - } - - // Lookup in chat if responses not found - if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest || requestType == schemas.CompactionRequest { - mc.logger.Debug("secondary lookup failed, trying vertex provider for the same model in chat completion") - pricing, ok = mc.pricingData[makeKey(model, "vertex", normalizeRequestType(schemas.ChatCompletionRequest))] - if ok { - return &pricing, true - } - } - } - - if provider == string(schemas.Vertex) { - // Vertex models can be of the form "provider/model", so try to lookup the model without the provider prefix and keep the original provider - if strings.Contains(model, "/") { - modelWithoutProvider := strings.SplitN(model, "/", 2)[1] - mc.logger.Debug("primary lookup failed, trying vertex provider for the same model with provider/model format %s", modelWithoutProvider) - pricing, ok = mc.pricingData[makeKey(modelWithoutProvider, "vertex", mode)] - if ok { - return &pricing, true - } - - // Lookup in chat if responses not found - if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest || requestType == schemas.CompactionRequest { - mc.logger.Debug("secondary lookup failed, trying vertex provider for the same model in chat completion") - pricing, ok = mc.pricingData[makeKey(modelWithoutProvider, "vertex", normalizeRequestType(schemas.ChatCompletionRequest))] - if ok { - return &pricing, true - } - } - } - } - - if provider == string(schemas.Bedrock) { - // If model is claude without "anthropic." prefix, try with "anthropic." prefix - if !strings.Contains(model, "anthropic.") && schemas.IsAnthropicModel(model) { - mc.logger.Debug("primary lookup failed, trying with anthropic. prefix for the same model") - pricing, ok = mc.pricingData[makeKey("anthropic."+model, provider, mode)] - if ok { - return &pricing, true - } - - // Lookup in chat if responses not found - if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest || requestType == schemas.CompactionRequest { - mc.logger.Debug("secondary lookup failed, trying chat provider for the same model in chat completion") - pricing, ok = mc.pricingData[makeKey("anthropic."+model, provider, normalizeRequestType(schemas.ChatCompletionRequest))] - if ok { - return &pricing, true - } - } - } - } - - // Lookup in chat if responses/compaction not found - if requestType == schemas.ResponsesRequest || requestType == schemas.ResponsesStreamRequest || requestType == schemas.WebSocketResponsesRequest || requestType == schemas.RealtimeRequest || requestType == schemas.CompactionRequest { - mc.logger.Debug("primary lookup failed, trying chat provider for the same model in chat completion") - pricing, ok = mc.pricingData[makeKey(model, provider, normalizeRequestType(schemas.ChatCompletionRequest))] - if ok { - return &pricing, true - } - } - - // Lookup in image generation if image edit not found - if requestType == schemas.ImageEditRequest || - requestType == schemas.ImageEditStreamRequest || - requestType == schemas.ImageVariationRequest { - mc.logger.Debug("primary lookup failed, trying image generation provider for the same model") - pricing, ok = mc.pricingData[makeKey(model, provider, normalizeRequestType(schemas.ImageGenerationRequest))] - if ok { - return &pricing, true - } - } - - // Lookup fallback chain for container_create: - // 1. Try chat mode for the same model (e.g. "container-1g" in chat mode) - // 2. Try the base "container" model in chat mode (default rate when no memory-specific entry exists) - if requestType == schemas.ContainerCreateRequest { - mc.logger.Debug("primary lookup failed, trying chat mode for container create pricing") - pricing, ok = mc.pricingData[makeKey(model, provider, normalizeRequestType(schemas.ChatCompletionRequest))] - if ok { - return &pricing, true - } - if model != "container" { - mc.logger.Debug("memory-specific container pricing not found, falling back to base container entry") - pricing, ok = mc.pricingData[makeKey("container", provider, normalizeRequestType(schemas.ChatCompletionRequest))] - if ok { - return &pricing, true - } - } - } - - return nil, false -} - -// UpsertModelPricingAttributes writes the additional_attributes column for -// every pricing row that matches (model, provider), then reloads the pricing -// cache so the new values are immediately visible to list-models. Returns -// the number of rows updated (0 = no such pricing row, which callers must -// surface as a validation error). An empty/nil attrs map clears the column. +// UpsertModelPricingAttributes writes additional_attributes for every row +// matching (model, provider) and reloads the pricing cache. func (mc *ModelCatalog) UpsertModelPricingAttributes(ctx context.Context, model string, provider schemas.ModelProvider, attrs map[string]string) (int64, error) { - if mc.configStore == nil { - return 0, fmt.Errorf("model catalog requires a config store") - } - rows, err := mc.configStore.UpsertModelPricingAttributes(ctx, model, string(provider), attrs) - if err != nil { - return 0, err - } - if rows == 0 { - return 0, nil - } - if err := mc.loadPricingFromDatabase(ctx); err != nil { - return rows, fmt.Errorf("failed to reload pricing cache after attribute write: %w", err) - } - return rows, nil + return mc.datasheet.UpsertModelPricingAttributes(ctx, model, provider, attrs) } -// --------------------------------------------------------------------------- -// Passthrough pricing helpers -// --------------------------------------------------------------------------- - -// detectPassthroughRequestType maps a provider + stripped path to a RequestType. -func detectPassthroughRequestType(provider schemas.ModelProvider, path string) schemas.RequestType { - if idx := strings.IndexByte(path, '?'); idx >= 0 { - path = path[:idx] - } - path = strings.TrimRight(path, "/") - switch provider { - case schemas.OpenAI, schemas.Azure: - switch { - case strings.HasSuffix(path, "/chat/completions"): - return schemas.ChatCompletionRequest - case strings.HasSuffix(path, "/completions"): - return schemas.TextCompletionRequest - case strings.HasSuffix(path, "/embeddings"): - return schemas.EmbeddingRequest - case strings.HasSuffix(path, "/responses/compact"): - return schemas.CompactionRequest - case strings.HasSuffix(path, "/responses"): - return schemas.ResponsesRequest - case strings.HasSuffix(path, "/images/generations"): - return schemas.ImageGenerationRequest - case strings.HasSuffix(path, "/images/edits"): - return schemas.ImageEditRequest - case strings.HasSuffix(path, "/images/variations"): - return schemas.ImageVariationRequest - case strings.HasSuffix(path, "/audio/speech"): - return schemas.SpeechRequest - case strings.HasSuffix(path, "/audio/transcriptions"), - strings.HasSuffix(path, "/audio/translations"): - return schemas.TranscriptionRequest - case strings.HasSuffix(path, "/containers"): - return schemas.ContainerCreateRequest - case strings.Contains(path, "/video"): - return schemas.VideoGenerationRequest - default: - return schemas.ChatCompletionRequest - } - case schemas.Gemini, schemas.Vertex: - // Interactions API paths carry no colon action suffix. - if strings.Contains(path, "/interactions") { - return schemas.ResponsesRequest - } - colonIdx := strings.LastIndexByte(path, ':') - if colonIdx < 0 { - return schemas.ChatCompletionRequest - } - switch path[colonIdx+1:] { - case "generateContent", "streamGenerateContent": - return schemas.ResponsesRequest - case "embedContent", "batchEmbedContents": - return schemas.EmbeddingRequest - case "generateImages": - return schemas.ImageGenerationRequest - case "predict": - return schemas.EmbeddingRequest - case "predictLongRunning": - return schemas.VideoGenerationRequest - default: - return schemas.ChatCompletionRequest - } - case schemas.Anthropic: - switch { - case strings.HasSuffix(path, "/messages"): - return schemas.ResponsesRequest - case strings.HasSuffix(path, "/complete"): - return schemas.TextCompletionRequest - default: - return schemas.ResponsesRequest - } - default: - return schemas.ChatCompletionRequest - } +func (mc *ModelCatalog) SetPricingOverrides(rows []configstoreTables.TablePricingOverride) error { + return mc.datasheet.SetOverrides(rows) } -// inferPassthroughRequestType determines the request type from usage fields (primary) -// and falls back to path detection for text/embedding/responses where LLMUsage is ambiguous. -func inferPassthroughRequestType(provider schemas.ModelProvider, path string, su *schemas.BifrostPassthroughUsage) schemas.RequestType { - if su != nil { - if su.ContainerIdentifier != "" { - return schemas.ContainerCreateRequest - } - if su.ImageUsage != nil { - return schemas.ImageGenerationRequest - } - if su.AudioInputChars > 0 { - return schemas.SpeechRequest - } - if su.AudioTokenDetails != nil || su.AudioSeconds != nil { - return schemas.TranscriptionRequest - } - if su.VideoSeconds != nil { - return schemas.VideoGenerationRequest - } - } - return detectPassthroughRequestType(provider, path) +func (mc *ModelCatalog) UpsertPricingOverrides(rows ...*configstoreTables.TablePricingOverride) error { + return mc.datasheet.UpsertOverrides(rows...) } -// passthroughUsageToCostInput converts BifrostPassthroughUsage into costInput. -func passthroughUsageToCostInput(su *schemas.BifrostPassthroughUsage) costInput { - var input costInput - if su.LLMUsage != nil { - input.usage = su.LLMUsage - } - if su.ServiceTier != nil { - input.tier = tierFromString(su.ServiceTier) - } - if su.ImageUsage != nil { - input.imageUsage = su.ImageUsage - input.imageSize = su.ImageSize - input.imageQuality = su.ImageQuality - } - if su.AudioInputChars > 0 { - input.audioTextInputChars = su.AudioInputChars - } - if su.AudioSeconds != nil { - input.audioSeconds = su.AudioSeconds - } - if su.AudioTokenDetails != nil { - input.audioTokenDetails = su.AudioTokenDetails - } - if su.VideoSeconds != nil { - input.videoSeconds = su.VideoSeconds - } - if su.ContainerIdentifier != "" { - input.containerIdentifierString = su.ContainerIdentifier - } - return input +func (mc *ModelCatalog) DeletePricingOverride(id string) { + mc.datasheet.DeleteOverride(id) } diff --git a/framework/modelcatalog/pricing_overrides.go b/framework/modelcatalog/pricing_overrides.go deleted file mode 100644 index baecf513475..00000000000 --- a/framework/modelcatalog/pricing_overrides.go +++ /dev/null @@ -1,470 +0,0 @@ -package modelcatalog - -import ( - "context" - "fmt" - "sort" - "strings" - - "github.com/maximhq/bifrost/core/schemas" - "github.com/maximhq/bifrost/framework/configstore" - configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" -) - -// PricingLookupScopes carries the runtime identifiers used to resolve scoped -// pricing overrides during cost calculation. -type PricingLookupScopes struct { - VirtualKeyID string - SelectedKeyID string - Provider string -} - -// PricingLookupScopesFromContext builds a PricingLookupScopes from a BifrostContext. -// It reads the governance virtual key ID (not the raw VK token) and the selected key ID. -// provider should be the provider name string (e.g. "openai"), pass "" if unavailable. -// Returns nil only when ctx is nil. An empty scopes value is still returned when all fields -// are empty so that global-scope overrides are always evaluated. -// DO NOT USE THIS FUNCTION IN A GO ROUTINE. This is because it reads from ctx which is cancelled when the request ends. -// Better to call it in PostHooks synchronously and then pass the scopes object to the pricing manager. -// Only use this in go routines when you know for sure that the request will not end before the go routine completes. -func PricingLookupScopesFromContext(ctx *schemas.BifrostContext, provider string) *PricingLookupScopes { - if ctx == nil { - return nil - } - virtualKeyID, _ := ctx.Value(schemas.BifrostContextKeyGovernanceVirtualKeyID).(string) - selectedKeyID, _ := ctx.Value(schemas.BifrostContextKeySelectedKeyID).(string) - return &PricingLookupScopes{ - VirtualKeyID: virtualKeyID, - SelectedKeyID: selectedKeyID, - Provider: provider, - } -} - -// ScopeKind identifies which governance scope an override applies to. -type ScopeKind string - -const ( - ScopeKindGlobal ScopeKind = "global" - ScopeKindProvider ScopeKind = "provider" - ScopeKindProviderKey ScopeKind = "provider_key" - ScopeKindVirtualKey ScopeKind = "virtual_key" - ScopeKindVirtualKeyProvider ScopeKind = "virtual_key_provider" - ScopeKindVirtualKeyProviderKey ScopeKind = "virtual_key_provider_key" -) - -// MatchType controls how an override pattern is matched against model names. -type MatchType string - -const ( - MatchTypeExact MatchType = "exact" - MatchTypeWildcard MatchType = "wildcard" -) - -// PricingOverride describes a scoped pricing override shared across config storage, -// model catalog compilation, and governance APIs. -type PricingOverride struct { - ID string `json:"id"` - Name string `json:"name"` - ScopeKind ScopeKind `json:"scope_kind"` - VirtualKeyID *string `json:"virtual_key_id,omitempty"` - ProviderID *string `json:"provider_id,omitempty"` - ProviderKeyID *string `json:"provider_key_id,omitempty"` - MatchType MatchType `json:"match_type"` - Pattern string `json:"pattern"` - RequestTypes []schemas.RequestType `json:"request_types,omitempty"` - Options PricingOptions `json:"options"` -} - -// customPricingEntry is a single flattened override ready for lookup. -type customPricingEntry struct { - id string - scopeKind ScopeKind - virtualKeyID string - providerID string - providerKeyID string - pattern string // exact model name, or wildcard prefix (trailing * stripped) - wildcard bool - requestModes map[string]struct{} // always non-nil for valid overrides - options PricingOptions -} - -// customPricingData is the in-memory lookup structure for pricing overrides. -// Exact matches are indexed by model name; wildcards are a flat slice. -type customPricingData struct { - exact map[string][]customPricingEntry - wildcard []customPricingEntry -} - -// IsValid validates the shared pricing override contract before persistence or runtime use. -// -// Input: override — the PricingOverride to validate (receiver). -// Output: error — non-nil if any scope, pattern, or request-type constraint is violated. -func (override *PricingOverride) IsValid() error { - if err := override.validateScopeKind(); err != nil { - return err - } - if err := override.validatePattern(); err != nil { - return err - } - return override.validateRequestTypes() -} - -// validateScopeKind validates the scope identifiers required by override.ScopeKind. -// -// Input: override — receiver; ScopeKind and the three optional ID fields are inspected. -// Output: error — non-nil when required identifiers are absent or forbidden ones are present. -func (override *PricingOverride) validateScopeKind() error { - switch override.ScopeKind { - case ScopeKindGlobal: - if override.VirtualKeyID != nil || override.ProviderID != nil || override.ProviderKeyID != nil { - return fmt.Errorf("global scope_kind must not include scope identifiers") - } - case ScopeKindProvider: - if override.ProviderID == nil { - return fmt.Errorf("provider_id is required for provider scope_kind") - } - if override.VirtualKeyID != nil || override.ProviderKeyID != nil { - return fmt.Errorf("provider scope_kind only supports provider_id") - } - case ScopeKindProviderKey: - if override.ProviderKeyID == nil { - return fmt.Errorf("provider_key_id is required for provider_key scope_kind") - } - if override.VirtualKeyID != nil || override.ProviderID != nil { - return fmt.Errorf("provider_key scope_kind only supports provider_key_id") - } - case ScopeKindVirtualKey: - if override.VirtualKeyID == nil { - return fmt.Errorf("virtual_key_id is required for virtual_key scope_kind") - } - if override.ProviderID != nil || override.ProviderKeyID != nil { - return fmt.Errorf("virtual_key scope_kind only supports virtual_key_id") - } - case ScopeKindVirtualKeyProvider: - if override.VirtualKeyID == nil || override.ProviderID == nil { - return fmt.Errorf("virtual_key_id and provider_id are required for virtual_key_provider scope_kind") - } - if override.ProviderKeyID != nil { - return fmt.Errorf("virtual_key_provider scope_kind does not support provider_key_id") - } - case ScopeKindVirtualKeyProviderKey: - if override.VirtualKeyID == nil || override.ProviderID == nil || override.ProviderKeyID == nil { - return fmt.Errorf("virtual_key_id, provider_id, and provider_key_id are required for virtual_key_provider_key scope_kind") - } - default: - return fmt.Errorf("unsupported scope_kind %q", override.ScopeKind) - } - return nil -} - -// validatePattern checks that Pattern is non-empty and consistent with MatchType. -// -// Input: override — receiver; Pattern and MatchType are inspected. -// Output: error — non-nil when the pattern is empty, contains a wildcard for exact mode, -// -// or does not end with a single trailing "*" for wildcard mode. -func (override *PricingOverride) validatePattern() error { - pattern := strings.TrimSpace(override.Pattern) - if pattern == "" { - return fmt.Errorf("pattern is required") - } - switch override.MatchType { - case MatchTypeExact: - if strings.Contains(pattern, "*") { - return fmt.Errorf("exact match pattern must not contain wildcards") - } - case MatchTypeWildcard: - if !strings.HasSuffix(pattern, "*") { - return fmt.Errorf("wildcard pattern must end with *") - } - if strings.Count(pattern, "*") != 1 { - return fmt.Errorf("wildcard pattern must contain exactly one trailing *") - } - default: - return fmt.Errorf("unsupported match_type %q", override.MatchType) - } - return nil -} - -// validateRequestTypes checks that RequestTypes is non-empty and that every entry is a -// supported base request type. Stream variants (e.g. chat_completion_stream) are rejected — -// the base type (chat_completion) already covers both streaming and non-streaming requests. -// -// Input: override — receiver; RequestTypes slice is inspected. -// Output: error — non-nil if RequestTypes is empty, or contains an unsupported or stream variant. -func (override *PricingOverride) validateRequestTypes() error { - if len(override.RequestTypes) == 0 { - return fmt.Errorf("request_types is required and must contain at least one value") - } - for _, rt := range override.RequestTypes { - if normalizeStreamRequestType(rt) != rt { - return fmt.Errorf("unsupported request_type %q: use the base type (e.g. %q covers both streaming and non-streaming)", rt, normalizeStreamRequestType(rt)) - } - if normalizeRequestType(rt) == "unknown" { - return fmt.Errorf("unsupported request_type %q", rt) - } - } - return nil -} - -// matchesScope reports whether the entry's governance scope matches the runtime identifiers. -// -// Input: scopes — runtime VirtualKeyID, SelectedKeyID, and Provider to match against. -// Output: bool — true when the entry's scope kind and stored IDs align with scopes. -func (e *customPricingEntry) matchesScope(scopes PricingLookupScopes) bool { - switch e.scopeKind { - case ScopeKindGlobal: - return true - case ScopeKindProvider: - return e.providerID == scopes.Provider - case ScopeKindProviderKey: - return e.providerKeyID == scopes.SelectedKeyID - case ScopeKindVirtualKey: - return e.virtualKeyID == scopes.VirtualKeyID - case ScopeKindVirtualKeyProvider: - return e.virtualKeyID == scopes.VirtualKeyID && e.providerID == scopes.Provider - case ScopeKindVirtualKeyProviderKey: - return e.virtualKeyID == scopes.VirtualKeyID && e.providerID == scopes.Provider && e.providerKeyID == scopes.SelectedKeyID - } - return false -} - -// matchesMode reports whether the entry applies to the given normalized request mode. -// -// Input: mode — normalized request type string (e.g. "chat", "embedding"). -// Output: bool — true when requestModes contains mode. -func (e *customPricingEntry) matchesMode(mode string) bool { - _, ok := e.requestModes[mode] - return ok -} - -// resolve walks the 6-scope priority hierarchy and returns the first matching -// pricing patch for the given model, request mode, and runtime scopes. -// -// Input: model — exact model name being priced. -// -// mode — normalized request type string (e.g. "chat", "embedding"). -// scopes — runtime governance identifiers used to narrow the scope search. -// -// Output: *PricingOptions — pointer to the first matching override's options, or nil if none match. -func (c *customPricingData) resolve(model, mode string, scopes PricingLookupScopes) *PricingOptions { - for _, scopeKind := range scopePriorityOrder(scopes) { - for i := range c.exact[model] { - e := &c.exact[model][i] - if e.scopeKind == scopeKind && e.matchesScope(scopes) && e.matchesMode(mode) { - return &e.options - } - } - for i := range c.wildcard { - e := &c.wildcard[i] - if e.scopeKind == scopeKind && e.matchesScope(scopes) && strings.HasPrefix(model, e.pattern) && e.matchesMode(mode) { - return &e.options - } - } - } - return nil -} - -// scopePriorityOrder returns scope kinds in most-specific-first order, -// skipping scopes that can't match given the available runtime identifiers. -// -// Input: scopes — runtime governance identifiers; empty fields cause the corresponding scope kinds to be omitted. -// Output: []ScopeKind — ordered list from most-specific (VirtualKeyProviderKey) to least-specific (Global). -func scopePriorityOrder(scopes PricingLookupScopes) []ScopeKind { - order := make([]ScopeKind, 0, 6) - if scopes.VirtualKeyID != "" && scopes.Provider != "" && scopes.SelectedKeyID != "" { - order = append(order, ScopeKindVirtualKeyProviderKey) - } - if scopes.VirtualKeyID != "" && scopes.Provider != "" { - order = append(order, ScopeKindVirtualKeyProvider) - } - if scopes.VirtualKeyID != "" { - order = append(order, ScopeKindVirtualKey) - } - if scopes.SelectedKeyID != "" { - order = append(order, ScopeKindProviderKey) - } - if scopes.Provider != "" { - order = append(order, ScopeKindProvider) - } - order = append(order, ScopeKindGlobal) - return order -} - -// buildCustomPricingData constructs a customPricingData lookup structure from a raw override slice. -// -// Input: overrides — slice of validated PricingOverride records loaded from the config store. -// Output: *customPricingData — ready-to-query structure with exact and wildcard indexes populated. -func buildCustomPricingData(overrides []PricingOverride) *customPricingData { - data := &customPricingData{ - exact: make(map[string][]customPricingEntry, len(overrides)), - } - for _, o := range overrides { - entry := customPricingEntry{ - id: o.ID, - scopeKind: o.ScopeKind, - options: o.Options, - } - if o.VirtualKeyID != nil { - entry.virtualKeyID = *o.VirtualKeyID - } - if o.ProviderID != nil { - entry.providerID = *o.ProviderID - } - if o.ProviderKeyID != nil { - entry.providerKeyID = *o.ProviderKeyID - } - entry.requestModes = make(map[string]struct{}, len(o.RequestTypes)) - for _, rt := range o.RequestTypes { - entry.requestModes[normalizeRequestType(rt)] = struct{}{} - } - pattern := strings.TrimSpace(o.Pattern) - switch o.MatchType { - case MatchTypeExact: - entry.pattern = pattern - data.exact[pattern] = append(data.exact[pattern], entry) - case MatchTypeWildcard: - entry.pattern = strings.TrimSuffix(pattern, "*") - entry.wildcard = true - data.wildcard = append(data.wildcard, entry) - } - } - // Sort wildcards by descending prefix length so more-specific patterns (e.g. "gpt-4*") - // are checked before broader ones (e.g. "gpt-*"), making precedence deterministic. - sort.Slice(data.wildcard, func(i, j int) bool { - return len(data.wildcard[i].pattern) > len(data.wildcard[j].pattern) - }) - return data -} - -// applyPricingOverrides resolves any active scoped pricing override for the given model -// and request type, then patches the catalog base pricing with the override values. -// It returns the original pricing unchanged when no custom pricing tree is loaded or -// when the request type cannot be mapped to a known pricing mode. -// -// Input: model — exact model name being priced. -// -// requestType — the request type used to derive the pricing mode. -// pricing — base pricing row from the catalog to patch. -// scopes — runtime governance identifiers used to narrow the override scope. -// -// Output: TableModelPricing — patched pricing row, or pricing unchanged if no override matches. -// bool — true when an override was applied, false otherwise. -func (mc *ModelCatalog) applyPricingOverrides(model string, requestType schemas.RequestType, pricing configstoreTables.TableModelPricing, scopes PricingLookupScopes) (configstoreTables.TableModelPricing, bool) { - mc.overridesMu.RLock() - custom := mc.customPricing - mc.overridesMu.RUnlock() - - if custom == nil { - return pricing, false - } - - mode := normalizeRequestType(requestType) - if mode == "unknown" { - return pricing, false - } - - if patch := custom.resolve(model, mode, scopes); patch != nil { - return patchPricing(pricing, *patch), true - } - return pricing, false -} - -// patchPricing applies override values onto a copy of the base pricing row. -// For all fields, a non-nil override pointer replaces the corresponding destination value; -// a nil override leaves the base value intact. -// The original pricing row is never modified; a patched copy is always returned. -// -// Input: pricing — base pricing row from the catalog. -// -// override — pricing options sourced from the matched override entry. -// -// Output: TableModelPricing — shallow copy of pricing with override fields applied. -func patchPricing(pricing configstoreTables.TableModelPricing, override PricingOptions) configstoreTables.TableModelPricing { - patched := pricing - - for _, field := range []struct { - dst **float64 - src *float64 - }{ - {dst: &patched.InputCostPerToken, src: override.InputCostPerToken}, - {dst: &patched.OutputCostPerToken, src: override.OutputCostPerToken}, - {dst: &patched.InputCostPerTokenPriority, src: override.InputCostPerTokenPriority}, - {dst: &patched.OutputCostPerTokenPriority, src: override.OutputCostPerTokenPriority}, - {dst: &patched.InputCostPerTokenFlex, src: override.InputCostPerTokenFlex}, - {dst: &patched.OutputCostPerTokenFlex, src: override.OutputCostPerTokenFlex}, - {dst: &patched.InputCostPerVideoPerSecond, src: override.InputCostPerVideoPerSecond}, - {dst: &patched.OutputCostPerVideoPerSecond, src: override.OutputCostPerVideoPerSecond}, - {dst: &patched.OutputCostPerSecond, src: override.OutputCostPerSecond}, - {dst: &patched.InputCostPerAudioPerSecond, src: override.InputCostPerAudioPerSecond}, - {dst: &patched.InputCostPerSecond, src: override.InputCostPerSecond}, - {dst: &patched.InputCostPerAudioToken, src: override.InputCostPerAudioToken}, - {dst: &patched.OutputCostPerAudioToken, src: override.OutputCostPerAudioToken}, - {dst: &patched.InputCostPerCharacter, src: override.InputCostPerCharacter}, - {dst: &patched.InputCostPerTokenAbove128kTokens, src: override.InputCostPerTokenAbove128kTokens}, - {dst: &patched.InputCostPerImageAbove128kTokens, src: override.InputCostPerImageAbove128kTokens}, - {dst: &patched.InputCostPerVideoPerSecondAbove128kTokens, src: override.InputCostPerVideoPerSecondAbove128kTokens}, - {dst: &patched.InputCostPerAudioPerSecondAbove128kTokens, src: override.InputCostPerAudioPerSecondAbove128kTokens}, - {dst: &patched.OutputCostPerTokenAbove128kTokens, src: override.OutputCostPerTokenAbove128kTokens}, - {dst: &patched.InputCostPerTokenAbove200kTokens, src: override.InputCostPerTokenAbove200kTokens}, - {dst: &patched.InputCostPerTokenAbove200kTokensPriority, src: override.InputCostPerTokenAbove200kTokensPriority}, - {dst: &patched.OutputCostPerTokenAbove200kTokens, src: override.OutputCostPerTokenAbove200kTokens}, - {dst: &patched.OutputCostPerTokenAbove200kTokensPriority, src: override.OutputCostPerTokenAbove200kTokensPriority}, - {dst: &patched.InputCostPerTokenAbove272kTokens, src: override.InputCostPerTokenAbove272kTokens}, - {dst: &patched.InputCostPerTokenAbove272kTokensPriority, src: override.InputCostPerTokenAbove272kTokensPriority}, - {dst: &patched.OutputCostPerTokenAbove272kTokens, src: override.OutputCostPerTokenAbove272kTokens}, - {dst: &patched.OutputCostPerTokenAbove272kTokensPriority, src: override.OutputCostPerTokenAbove272kTokensPriority}, - {dst: &patched.CacheCreationInputTokenCostAbove200kTokens, src: override.CacheCreationInputTokenCostAbove200kTokens}, - {dst: &patched.CacheReadInputTokenCostAbove200kTokens, src: override.CacheReadInputTokenCostAbove200kTokens}, - {dst: &patched.CacheReadInputTokenCost, src: override.CacheReadInputTokenCost}, - {dst: &patched.CacheCreationInputTokenCost, src: override.CacheCreationInputTokenCost}, - {dst: &patched.CacheCreationInputTokenCostAbove1hr, src: override.CacheCreationInputTokenCostAbove1hr}, - {dst: &patched.CacheCreationInputTokenCostAbove1hrAbove200kTokens, src: override.CacheCreationInputTokenCostAbove1hrAbove200kTokens}, - {dst: &patched.CacheCreationInputAudioTokenCost, src: override.CacheCreationInputAudioTokenCost}, - {dst: &patched.CacheReadInputTokenCostPriority, src: override.CacheReadInputTokenCostPriority}, - {dst: &patched.CacheReadInputTokenCostFlex, src: override.CacheReadInputTokenCostFlex}, - {dst: &patched.CacheReadInputTokenCostAbove200kTokensPriority, src: override.CacheReadInputTokenCostAbove200kTokensPriority}, - {dst: &patched.CacheReadInputTokenCostAbove272kTokens, src: override.CacheReadInputTokenCostAbove272kTokens}, - {dst: &patched.CacheReadInputTokenCostAbove272kTokensPriority, src: override.CacheReadInputTokenCostAbove272kTokensPriority}, - {dst: &patched.InputCostPerTokenBatches, src: override.InputCostPerTokenBatches}, - {dst: &patched.OutputCostPerTokenBatches, src: override.OutputCostPerTokenBatches}, - {dst: &patched.InputCostPerImageToken, src: override.InputCostPerImageToken}, - {dst: &patched.OutputCostPerImageToken, src: override.OutputCostPerImageToken}, - {dst: &patched.InputCostPerImage, src: override.InputCostPerImage}, - {dst: &patched.OutputCostPerImage, src: override.OutputCostPerImage}, - {dst: &patched.InputCostPerPixel, src: override.InputCostPerPixel}, - {dst: &patched.OutputCostPerPixel, src: override.OutputCostPerPixel}, - {dst: &patched.OutputCostPerImagePremiumImage, src: override.OutputCostPerImagePremiumImage}, - {dst: &patched.OutputCostPerImageAbove512x512Pixels, src: override.OutputCostPerImageAbove512x512Pixels}, - {dst: &patched.OutputCostPerImageAbove512x512PixelsPremium, src: override.OutputCostPerImageAbove512x512PixelsPremium}, - {dst: &patched.OutputCostPerImageAbove1024x1024Pixels, src: override.OutputCostPerImageAbove1024x1024Pixels}, - {dst: &patched.OutputCostPerImageAbove1024x1024PixelsPremium, src: override.OutputCostPerImageAbove1024x1024PixelsPremium}, - {dst: &patched.OutputCostPerImageAbove2048x2048Pixels, src: override.OutputCostPerImageAbove2048x2048Pixels}, - {dst: &patched.OutputCostPerImageAbove4096x4096Pixels, src: override.OutputCostPerImageAbove4096x4096Pixels}, - {dst: &patched.CacheReadInputImageTokenCost, src: override.CacheReadInputImageTokenCost}, - {dst: &patched.SearchContextCostPerQuery, src: override.SearchContextCostPerQuery}, - {dst: &patched.CodeInterpreterCostPerSession, src: override.CodeInterpreterCostPerSession}, - {dst: &patched.OutputCostPerImageLowQuality, src: override.OutputCostPerImageLowQuality}, - {dst: &patched.OutputCostPerImageMediumQuality, src: override.OutputCostPerImageMediumQuality}, - {dst: &patched.OutputCostPerImageHighQuality, src: override.OutputCostPerImageHighQuality}, - {dst: &patched.OutputCostPerImageAutoQuality, src: override.OutputCostPerImageAutoQuality}, - {dst: &patched.OCRCostPerPage, src: override.OCRCostPerPage}, - {dst: &patched.AnnotationCostPerPage, src: override.AnnotationCostPerPage}, - } { - if field.src != nil { - *field.dst = field.src - } - } - return patched -} - -func (mc *ModelCatalog) loadPricingOverridesFromStore(ctx context.Context) error { - if mc.configStore == nil { - return nil - } - rows, err := mc.configStore.GetPricingOverrides(ctx, configstore.PricingOverrideFilters{}) - if err != nil { - return err - } - return mc.SetPricingOverrides(rows) -} diff --git a/framework/modelcatalog/pricing_overrides_test.go b/framework/modelcatalog/pricing_overrides_test.go deleted file mode 100644 index 1f2dd9e5969..00000000000 --- a/framework/modelcatalog/pricing_overrides_test.go +++ /dev/null @@ -1,507 +0,0 @@ -package modelcatalog - -import ( - "testing" - - bifrost "github.com/maximhq/bifrost/core" - "github.com/maximhq/bifrost/core/schemas" - configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -type noOpLogger struct{} - -func (noOpLogger) Debug(string, ...any) {} -func (noOpLogger) Info(string, ...any) {} -func (noOpLogger) Warn(string, ...any) {} -func (noOpLogger) Error(string, ...any) {} -func (noOpLogger) Fatal(string, ...any) {} -func (noOpLogger) SetLevel(schemas.LogLevel) {} -func (noOpLogger) SetOutputType(schemas.LoggerOutputType) {} -func (noOpLogger) LogHTTPRequest(schemas.LogLevel, string) schemas.LogEventBuilder { - return schemas.NoopLogEvent -} - -func TestGetPricing_OverridePrecedenceExactWildcard(t *testing.T) { - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.pricingData[makeKey("gpt-4o", "openai", "chat")] = configstoreTables.TableModelPricing{ - Model: "gpt-4o", - Provider: "openai", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - } - - providerID := "openai" - require.NoError(t, mc.SetPricingOverrides([]configstoreTables.TablePricingOverride{ - { - ID: "openai-override-0", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeWildcard), - Pattern: "gpt-*", - RequestTypes: []schemas.RequestType{schemas.ChatCompletionRequest}, - PricingPatchJSON: `{"input_cost_per_token":10}`, - }, - { - ID: "openai-override-1", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeExact), - Pattern: "gpt-4o", - RequestTypes: []schemas.RequestType{schemas.ChatCompletionRequest}, - PricingPatchJSON: `{"input_cost_per_token":20}`, - }, - })) - - pricing := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "openai"}) - require.NotNil(t, pricing) - require.NotNil(t, pricing.InputCostPerToken) - assert.Equal(t, 20.0, *pricing.InputCostPerToken) -} - -func TestGetPricing_RequestTypeSpecificOverrideBeatsGeneric(t *testing.T) { - t.Skip() - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.pricingData[makeKey("gpt-4o", "openai", "responses")] = configstoreTables.TableModelPricing{ - Model: "gpt-4o", - Provider: "openai", - Mode: "responses", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - } - - providerID := "openai" - require.NoError(t, mc.SetPricingOverrides([]configstoreTables.TablePricingOverride{ - { - ID: "openai-generic", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeExact), - Pattern: "gpt-4o", - PricingPatchJSON: `{"input_cost_per_token":9}`, - }, - { - ID: "openai-specific", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeExact), - Pattern: "gpt-4o", - RequestTypes: []schemas.RequestType{schemas.ResponsesRequest}, - PricingPatchJSON: `{"input_cost_per_token":15}`, - }, - })) - - pricing := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o"}, schemas.ResponsesRequest, PricingLookupScopes{Provider: "openai"}) - require.NotNil(t, pricing) - assert.Equal(t, 15.0, pricing.InputCostPerToken) -} - -func TestGetPricing_AppliesOverrideAfterFallbackResolution(t *testing.T) { - t.Skip() - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.pricingData[makeKey("gpt-4o", "vertex", "chat")] = configstoreTables.TableModelPricing{ - Model: "gpt-4o", - Provider: "vertex", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - } - - geminiProviderID := "gemini" - require.NoError(t, mc.SetPricingOverrides([]configstoreTables.TablePricingOverride{ - { - ID: "gemini-override", - ScopeKind: string(ScopeKindProvider), - ProviderID: &geminiProviderID, - MatchType: string(MatchTypeExact), - Pattern: "gpt-4o", - PricingPatchJSON: `{"input_cost_per_token":7}`, - }, - })) - - pricing := mc.resolvePricing(schemas.RoutingInfo{Provider: "gemini", Model: "gpt-4o"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "gemini"}) - require.NotNil(t, pricing) - assert.Equal(t, 7.0, pricing.InputCostPerToken) -} - -func TestGetPricing_DeploymentLookupUsesResolvedModelForOverrideMatching(t *testing.T) { - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.pricingData[makeKey("dep-gpt4o", "openai", "chat")] = configstoreTables.TableModelPricing{ - Model: "dep-gpt4o", - Provider: "openai", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - } - - providerID := "openai" - require.NoError(t, mc.SetPricingOverrides([]configstoreTables.TablePricingOverride{ - { - ID: "resolved-model-override", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeExact), - Pattern: "dep-gpt4o", - RequestTypes: []schemas.RequestType{schemas.ChatCompletionRequest}, - PricingPatchJSON: `{"input_cost_per_token":7}`, - }, - })) - - // Override pattern matches the resolved model name ("dep-gpt4o"), not the - // originally requested name ("gpt-4o"), because resolved model has priority. - pricing := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o", ResolvedKeyAlias: &schemas.ResolvedKeyAlias{ModelID: "dep-gpt4o"}}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "openai"}) - require.NotNil(t, pricing) - require.NotNil(t, pricing.InputCostPerToken) - assert.Equal(t, 7.0, *pricing.InputCostPerToken) -} - -func TestGetPricing_FallbackUsesRequestedProviderForScopeMatching(t *testing.T) { - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.pricingData[makeKey("gpt-4o", "vertex", "chat")] = configstoreTables.TableModelPricing{ - Model: "gpt-4o", - Provider: "vertex", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - } - - geminiProviderID := "gemini" - vertexProviderID := "vertex" - require.NoError(t, mc.SetPricingOverrides([]configstoreTables.TablePricingOverride{ - { - ID: "gemini-provider-override", - ScopeKind: string(ScopeKindProvider), - ProviderID: &geminiProviderID, - MatchType: string(MatchTypeExact), - Pattern: "gpt-4o", - RequestTypes: []schemas.RequestType{schemas.ChatCompletionRequest}, - PricingPatchJSON: `{"input_cost_per_token":5}`, - }, - { - ID: "vertex-provider-override", - ScopeKind: string(ScopeKindProvider), - ProviderID: &vertexProviderID, - MatchType: string(MatchTypeExact), - Pattern: "gpt-4o", - RequestTypes: []schemas.RequestType{schemas.ChatCompletionRequest}, - PricingPatchJSON: `{"input_cost_per_token":9}`, - }, - })) - - pricing := mc.resolvePricing(schemas.RoutingInfo{Provider: "gemini", Model: "gpt-4o"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "gemini"}) - require.NotNil(t, pricing) - require.NotNil(t, pricing.InputCostPerToken) - assert.Equal(t, 5.0, *pricing.InputCostPerToken) -} - -func TestGetPricing_ExactOverrideDoesNotMatchProviderPrefixedModel(t *testing.T) { - t.Skip() - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.pricingData[makeKey("openai/gpt-4o", "openai", "chat")] = configstoreTables.TableModelPricing{ - Model: "openai/gpt-4o", - Provider: "openai", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - } - - providerID := "openai" - require.NoError(t, mc.SetPricingOverrides([]configstoreTables.TablePricingOverride{ - { - ID: "openai-override-0", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeExact), - Pattern: "gpt-4o", - PricingPatchJSON: `{"input_cost_per_token":19}`, - }, - })) - - pricing := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "openai/gpt-4o"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "openai"}) - require.NotNil(t, pricing) - assert.Equal(t, 1.0, pricing.InputCostPerToken) -} - -func TestGetPricing_NoMatchingOverrideLeavesPricingUnchanged(t *testing.T) { - t.Skip() - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - baseCacheRead := 0.4 - mc.pricingData[makeKey("gpt-4o", "openai", "chat")] = configstoreTables.TableModelPricing{ - Model: "gpt-4o", - Provider: "openai", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - CacheReadInputTokenCost: &baseCacheRead, - } - - providerID := "openai" - require.NoError(t, mc.SetPricingOverrides([]configstoreTables.TablePricingOverride{ - { - ID: "openai-override-0", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeWildcard), - Pattern: "claude-*", - PricingPatchJSON: `{"input_cost_per_token":9}`, - }, - })) - - pricing := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "openai"}) - require.NotNil(t, pricing) - assert.Equal(t, 1.0, pricing.InputCostPerToken) - assert.Equal(t, 2.0, pricing.OutputCostPerToken) - require.NotNil(t, pricing.CacheReadInputTokenCost) - assert.Equal(t, 0.4, *pricing.CacheReadInputTokenCost) -} - -func TestDeleteProviderPricingOverrides_StopsApplying(t *testing.T) { - t.Skip() - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.pricingData[makeKey("gpt-4o", "openai", "chat")] = configstoreTables.TableModelPricing{ - Model: "gpt-4o", - Provider: "openai", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - } - - providerID := "openai" - require.NoError(t, mc.SetPricingOverrides([]configstoreTables.TablePricingOverride{ - { - ID: "openai-override-0", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeExact), - Pattern: "gpt-4o", - PricingPatchJSON: `{"input_cost_per_token":11}`, - }, - })) - - pricing := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "openai"}) - require.NotNil(t, pricing) - assert.Equal(t, 11.0, pricing.InputCostPerToken) - - require.NoError(t, mc.SetPricingOverrides(nil)) - - pricing = mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "openai"}) - require.NotNil(t, pricing) - assert.Equal(t, 1.0, pricing.InputCostPerToken) -} - -func TestGetPricing_WildcardSpecificityLongerLiteralWins(t *testing.T) { - t.Skip() - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.pricingData[makeKey("gpt-4o-mini", "openai", "chat")] = configstoreTables.TableModelPricing{ - Model: "gpt-4o-mini", - Provider: "openai", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - } - - providerID := "openai" - require.NoError(t, mc.SetPricingOverrides([]configstoreTables.TablePricingOverride{ - { - ID: "openai-override-0", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeWildcard), - Pattern: "gpt-*", - PricingPatchJSON: `{"input_cost_per_token":5}`, - }, - { - ID: "openai-override-1", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeWildcard), - Pattern: "gpt-4o*", - PricingPatchJSON: `{"input_cost_per_token":6}`, - }, - })) - - pricing := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o-mini"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "openai"}) - require.NotNil(t, pricing) - assert.Equal(t, 6.0, pricing.InputCostPerToken) -} - -// TestGetPricing_FirstInsertionWinsOnTie verifies that when multiple wildcard overrides -// match the same model and scope, the first one inserted takes precedence. -func TestGetPricing_FirstInsertionWinsOnTie(t *testing.T) { - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.pricingData[makeKey("gpt-4o-mini", "openai", "chat")] = configstoreTables.TableModelPricing{ - Model: "gpt-4o-mini", - Provider: "openai", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - } - - providerID := "openai" - require.NoError(t, mc.SetPricingOverrides([]configstoreTables.TablePricingOverride{ - { - ID: "a-override", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeWildcard), - Pattern: "gpt-4o*", - RequestTypes: []schemas.RequestType{schemas.ChatCompletionRequest}, - PricingPatchJSON: `{"input_cost_per_token":8}`, - }, - { - ID: "b-override", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerID, - MatchType: string(MatchTypeWildcard), - Pattern: "gpt-4o*", - RequestTypes: []schemas.RequestType{schemas.ChatCompletionRequest}, - PricingPatchJSON: `{"input_cost_per_token":9}`, - }, - })) - - pricing := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o-mini"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "openai"}) - require.NotNil(t, pricing) - require.NotNil(t, pricing.InputCostPerToken) - assert.Equal(t, 8.0, *pricing.InputCostPerToken) -} - -func TestPatchPricing_PartialPatchOnlyChangesSpecifiedFields(t *testing.T) { - t.Skip() - baseCacheRead := 0.4 - baseInputImage := 0.7 - base := configstoreTables.TableModelPricing{ - Model: "gpt-4o", - Provider: "openai", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - CacheReadInputTokenCost: &baseCacheRead, - InputCostPerImage: &baseInputImage, - } - - cacheRead := 0.9 - patched := patchPricing(base, PricingOptions{ - InputCostPerToken: bifrost.Ptr(3.0), - CacheReadInputTokenCost: &cacheRead, - }) - - assert.Equal(t, 3.0, patched.InputCostPerToken) - require.NotNil(t, patched.CacheReadInputTokenCost) - assert.Equal(t, 0.9, *patched.CacheReadInputTokenCost) - - assert.Equal(t, 2.0, patched.OutputCostPerToken) - require.NotNil(t, patched.InputCostPerImage) - assert.Equal(t, 0.7, *patched.InputCostPerImage) -} - -func TestApplyScopedPricingOverrides_ScopePrecedence(t *testing.T) { - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - - providerScopeID := "openai" - providerKeyScopeID := "provider-key-1" - virtualKeyScopeID := "virtual-key-1" - - require.NoError(t, mc.SetPricingOverrides([]configstoreTables.TablePricingOverride{ - { - ID: "global", - ScopeKind: string(ScopeKindGlobal), - MatchType: string(MatchTypeExact), - Pattern: "gpt-5-nano", - RequestTypes: []schemas.RequestType{schemas.ChatCompletionRequest}, - PricingPatchJSON: `{"input_cost_per_token":2}`, - }, - { - ID: "provider", - ScopeKind: string(ScopeKindProvider), - ProviderID: &providerScopeID, - MatchType: string(MatchTypeExact), - Pattern: "gpt-5-nano", - RequestTypes: []schemas.RequestType{schemas.ChatCompletionRequest}, - PricingPatchJSON: `{"input_cost_per_token":3}`, - }, - { - ID: "provider-key", - ScopeKind: string(ScopeKindProviderKey), - ProviderKeyID: &providerKeyScopeID, - MatchType: string(MatchTypeExact), - Pattern: "gpt-5-nano", - RequestTypes: []schemas.RequestType{schemas.ChatCompletionRequest}, - PricingPatchJSON: `{"input_cost_per_token":4}`, - }, - { - ID: "virtual-key", - ScopeKind: string(ScopeKindVirtualKey), - VirtualKeyID: &virtualKeyScopeID, - MatchType: string(MatchTypeExact), - Pattern: "gpt-5-nano", - RequestTypes: []schemas.RequestType{schemas.ChatCompletionRequest}, - PricingPatchJSON: `{"input_cost_per_token":5}`, - }, - })) - - base := configstoreTables.TableModelPricing{ - Model: "gpt-5-nano", - Provider: "openai", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(1.0), - OutputCostPerToken: bifrost.Ptr(2.0), - } - - tests := []struct { - name string - scopes PricingLookupScopes - expected float64 - }{ - { - name: "virtual key wins over provider key, provider and global", - scopes: PricingLookupScopes{ - VirtualKeyID: virtualKeyScopeID, - SelectedKeyID: providerKeyScopeID, - Provider: providerScopeID, - }, - expected: 5.0, - }, - { - name: "provider key wins over provider and global", - scopes: PricingLookupScopes{ - SelectedKeyID: providerKeyScopeID, - Provider: providerScopeID, - }, - expected: 4.0, - }, - { - name: "provider wins over global", - scopes: PricingLookupScopes{ - Provider: providerScopeID, - }, - expected: 3.0, - }, - { - name: "global applies when no narrower scope is provided", - scopes: PricingLookupScopes{}, - expected: 2.0, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - patched, applied := mc.applyPricingOverrides("gpt-5-nano", schemas.ChatCompletionRequest, base, tc.scopes) - require.True(t, applied) - require.NotNil(t, patched.InputCostPerToken) - assert.Equal(t, tc.expected, *patched.InputCostPerToken) - }) - } -} diff --git a/framework/modelcatalog/pricing_test.go b/framework/modelcatalog/pricing_test.go deleted file mode 100644 index 7faf7271c92..00000000000 --- a/framework/modelcatalog/pricing_test.go +++ /dev/null @@ -1,2791 +0,0 @@ -package modelcatalog - -import ( - "context" - "encoding/json" - "os" - "testing" - - bifrost "github.com/maximhq/bifrost/core" - "github.com/maximhq/bifrost/core/schemas" - configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// --------------------------------------------------------------------------- -// helpers -// --------------------------------------------------------------------------- - -// chatPricing returns a TableModelPricing with the given per-token rates. -func chatPricing(input, output float64) configstoreTables.TableModelPricing { - return configstoreTables.TableModelPricing{ - Model: "test-model", - Provider: "test-provider", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(input), - OutputCostPerToken: bifrost.Ptr(output), - } -} - -// testCatalogWithPricing creates a catalog pre-loaded with the given pricing entries. -func testCatalogWithPricing(entries map[string]configstoreTables.TableModelPricing) *ModelCatalog { - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - for k, v := range entries { - mc.pricingData[k] = v - } - return mc -} - -// routingInfoFor builds a minimal RoutingInfo populated by core.bifrost for a -// non-aliased request — the form pricing reads from. -func routingInfoFor(provider schemas.ModelProvider, model string) schemas.RoutingInfo { - return schemas.RoutingInfo{Provider: provider, Model: model} -} - -// makeChatResponse builds a minimal BifrostResponse for a chat completion. -func makeChatResponse(provider schemas.ModelProvider, model string, usage *schemas.BifrostLLMUsage) *schemas.BifrostResponse { - return &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - Usage: usage, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - RoutingInfo: routingInfoFor(provider, model), - }, - }, - } -} - -// makeEmbeddingResponse builds a minimal BifrostResponse for an embedding request. -func makeEmbeddingResponse(provider schemas.ModelProvider, model string, usage *schemas.BifrostLLMUsage) *schemas.BifrostResponse { - return &schemas.BifrostResponse{ - EmbeddingResponse: &schemas.BifrostEmbeddingResponse{ - Usage: usage, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.EmbeddingRequest, - RoutingInfo: routingInfoFor(provider, model), - }, - }, - } -} - -// makeRerankResponse builds a minimal BifrostResponse for a rerank request. -func makeRerankResponse(provider schemas.ModelProvider, model string, usage *schemas.BifrostLLMUsage) *schemas.BifrostResponse { - return &schemas.BifrostResponse{ - RerankResponse: &schemas.BifrostRerankResponse{ - Usage: usage, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.RerankRequest, - RoutingInfo: routingInfoFor(provider, model), - }, - }, - } -} - -// makeImageResponse builds a minimal BifrostResponse for an image generation request. -func makeImageResponse(provider schemas.ModelProvider, model string, usage *schemas.ImageUsage) *schemas.BifrostResponse { - return &schemas.BifrostResponse{ - ImageGenerationResponse: &schemas.BifrostImageGenerationResponse{ - Usage: usage, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ImageGenerationRequest, - RoutingInfo: routingInfoFor(provider, model), - }, - }, - } -} - -func derefF(f *float64) float64 { - if f == nil { - return 0 - } - return *f -} - -// ========================================================================= -// 1. computeTextCost — unit tests (pure function, no catalog) -// ========================================================================= - -func TestComputeTextCost_BasicInputOutput(t *testing.T) { - // GPT-4o: $5/M input, $15/M output - p := chatPricing(0.000005, 0.000015) - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - } - cost := computeTextCost(&p, usage, serviceTier{}) - // 1000 * 0.000005 + 500 * 0.000015 = 0.005 + 0.0075 = 0.0125 - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestComputeTextCost_NilUsage(t *testing.T) { - p := chatPricing(0.000005, 0.000015) - assert.Equal(t, 0.0, computeTextCost(&p, nil, serviceTier{})) -} - -func TestComputeTextCost_ZeroTokens(t *testing.T) { - p := chatPricing(0.000005, 0.000015) - usage := &schemas.BifrostLLMUsage{} - assert.Equal(t, 0.0, computeTextCost(&p, usage, serviceTier{})) -} - -func TestComputeTextCost_WithCachedPromptTokens(t *testing.T) { - // Claude 3.5 Sonnet (Bedrock): input=$3/M, output=$15/M, cache_read=$0.3/M, cache_creation=$3.75/M - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCost = bifrost.Ptr(0.0000003) - p.CacheCreationInputTokenCost = bifrost.Ptr(0.00000375) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 2000, - CompletionTokens: 500, - TotalTokens: 2500, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedReadTokens: 1500, // 1500 read from cache - CachedWriteTokens: 200, // 200 cache creation tokens - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Both cached read and write tokens are input-side deductions from promptTokens. - // Input: (2000-1500-200)*0.000003 + 1500*0.0000003 + 200*0.00000375 = 0.0009 + 0.00045 + 0.00075 = 0.0021 - // Output: 500*0.000015 = 0.0075 - // Total: 0.0021 + 0.0075 = 0.0096 - assert.InDelta(t, 0.0096, cost, 1e-12) -} - -func TestComputeTextCost_With1hrCacheCreationTokens(t *testing.T) { - // claude-3-5-sonnet-20241022-v2:0 on Bedrock: - // input=$3/M, output=$15/M, cache_creation=$3.75/M, cache_creation_1hr=$7.50/M, cache_read=$0.3/M - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCost = bifrost.Ptr(3e-7) - p.CacheCreationInputTokenCost = bifrost.Ptr(0.00000375) - p.CacheCreationInputTokenCostAbove1hr = bifrost.Ptr(0.0000075) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 2000, - CompletionTokens: 500, - TotalTokens: 2500, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedWriteTokens: 1000, - CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{ - CachedWriteTokens1h: 600, // 600 at 1hr rate, 400 at standard rate - }, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Input (non-cached): (2000-1000)*0.000003 = 0.003 - // Cache creation (1hr): 600*0.0000075 = 0.0045 - // Cache creation (standard): 400*0.00000375 = 0.0015 - // Output: 500*0.000015 = 0.0075 - // Total: 0.003 + 0.0045 + 0.0015 + 0.0075 = 0.0165 - assert.InDelta(t, 0.0165, cost, 1e-12) -} - -func TestComputeTextCost_StandardCacheCreationPricingLesserThan1hr(t *testing.T) { - // Standard (5-min TTL) cache creation is cheaper than 1hr TTL cache creation. - // 1hr rate ($7.50/M) is 2x the standard rate ($3.75/M). - p := chatPricing(0.000003, 0.000015) - p.CacheCreationInputTokenCost = bifrost.Ptr(0.00000375) - p.CacheCreationInputTokenCostAbove1hr = bifrost.Ptr(0.0000075) - - base := &schemas.BifrostLLMUsage{ - PromptTokens: 2000, - CompletionTokens: 500, - TotalTokens: 2500, - } - - usageStandard := *base - usageStandard.PromptTokensDetails = &schemas.ChatPromptTokensDetails{ - CachedWriteTokens: 1000, - } - - usage1hr := *base - usage1hr.PromptTokensDetails = &schemas.ChatPromptTokensDetails{ - CachedWriteTokens: 1000, - CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{ - CachedWriteTokens1h: 1000, // all 1000 tokens at 1hr rate - }, - } - - costStandard := computeTextCost(&p, &usageStandard, serviceTier{}) - cost1hr := computeTextCost(&p, &usage1hr, serviceTier{}) - - assert.Less(t, costStandard, cost1hr, "standard cache creation should cost less than 1hr cache creation") -} - -func TestComputeTextCost_1hrCacheCreationFallsBackToStandardWhenAbove1hrRateAbsent(t *testing.T) { - // claude-3-5-haiku on Bedrock has no cache_creation_input_token_cost_above_1hr entry. - // Tokens marked as 1hr cache writes must fall back to the standard cache creation rate. - p := chatPricing(8e-7, 0.000004) - p.CacheReadInputTokenCost = bifrost.Ptr(8e-8) - p.CacheCreationInputTokenCost = bifrost.Ptr(0.000001) - // CacheCreationInputTokenCostAbove1hr intentionally left nil - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 2000, - CompletionTokens: 500, - TotalTokens: 2500, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedWriteTokens: 1000, - CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{ - CachedWriteTokens1h: 1000, // all 1hr tokens, but no above_1hr rate configured - }, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Input (non-cached): (2000-1000)*8e-7 = 0.0008 - // Cache creation (1hr fallback → standard 0.000001): 1000*0.000001 = 0.001 - // Output: 500*0.000004 = 0.002 - // Total: 0.0008 + 0.001 + 0.002 = 0.0038 - assert.InDelta(t, 0.0038, cost, 1e-12) -} - -func TestComputeTextCost_CacheWriteTokenDetailsNil_FallsBackToStandardCreationRate(t *testing.T) { - // CachedWriteTokens is set but CachedWriteTokenDetails is nil. - // All write tokens must use the standard cache creation rate even though above_1hr is configured. - p := chatPricing(0.000003, 0.000015) - p.CacheCreationInputTokenCost = bifrost.Ptr(0.00000375) - p.CacheCreationInputTokenCostAbove1hr = bifrost.Ptr(0.0000075) // present but must not be used - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 2000, - CompletionTokens: 500, - TotalTokens: 2500, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedWriteTokens: 1000, - CachedWriteTokenDetails: nil, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Input (non-cached): (2000-1000)*0.000003 = 0.003 - // Cache creation (standard, no 1hr details): 1000*0.00000375 = 0.00375 - // Output: 500*0.000015 = 0.0075 - // Total: 0.003 + 0.00375 + 0.0075 = 0.01425 - assert.InDelta(t, 0.01425, cost, 1e-12) -} - -func TestComputeTextCost_CacheWriteTokenDetails1hZero_FallsBackToStandardCreationRate(t *testing.T) { - // CachedWriteTokenDetails is present but CachedWriteTokens1h is 0 (e.g. all tokens - // used 5-min TTL). All write tokens must use the standard cache creation rate. - p := chatPricing(0.000003, 0.000015) - p.CacheCreationInputTokenCost = bifrost.Ptr(0.00000375) - p.CacheCreationInputTokenCostAbove1hr = bifrost.Ptr(0.0000075) // present but must not be used - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 2000, - CompletionTokens: 500, - TotalTokens: 2500, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedWriteTokens: 1000, - CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{ - CachedWriteTokens1h: 0, - }, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Input (non-cached): (2000-1000)*0.000003 = 0.003 - // Cache creation (standard, 1h count is 0): 1000*0.00000375 = 0.00375 - // Output: 500*0.000015 = 0.0075 - // Total: 0.003 + 0.00375 + 0.0075 = 0.01425 - assert.InDelta(t, 0.01425, cost, 1e-12) -} - -func TestComputeTextCost_1hrCacheCreationAbove200k_UsesAbove1hrAbove200kRate(t *testing.T) { - // claude-3-5-sonnet-20241022-v2:0 on Bedrock has all four cache creation tiers. - // When totalTokens > 200k and CachedWriteTokens1h > 0, the above_1hr_above_200k rate - // ($15/M) must be used — the most specific tier wins. - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenAbove200kTokens = bifrost.Ptr(0.000006) - p.OutputCostPerTokenAbove200kTokens = bifrost.Ptr(0.00003) - p.CacheCreationInputTokenCost = bifrost.Ptr(0.00000375) - p.CacheCreationInputTokenCostAbove200kTokens = bifrost.Ptr(0.0000075) - p.CacheCreationInputTokenCostAbove1hr = bifrost.Ptr(0.0000075) - p.CacheCreationInputTokenCostAbove1hrAbove200kTokens = bifrost.Ptr(0.000015) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 180000, - CompletionTokens: 25000, - TotalTokens: 205000, // above 200k threshold - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedWriteTokens: 10000, - CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{ - CachedWriteTokens1h: 10000, - }, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Input rate (>200k): 0.000006; output rate (>200k): 0.00003 - // Input (non-cached): (180000-10000)*0.000006 = 170000*0.000006 = 1.02 - // Cache creation 1hr above 200k: 10000*0.000015 = 0.15 - // Output: 25000*0.00003 = 0.75 - // Total: 1.02 + 0.15 + 0.75 = 1.92 - assert.InDelta(t, 1.92, cost, 1e-9) -} - -func TestComputeTextCost_1hrCacheCreationAbove200k_FallsBackToAbove1hrWhenAbove200kRateAbsent(t *testing.T) { - // When CacheCreationInputTokenCostAbove1hrAbove200kTokens is absent but - // CacheCreationInputTokenCostAbove1hr is present, the 1hr rate must be used - // even for >200k requests. - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenAbove200kTokens = bifrost.Ptr(0.000006) - p.OutputCostPerTokenAbove200kTokens = bifrost.Ptr(0.00003) - p.CacheCreationInputTokenCost = bifrost.Ptr(0.00000375) - p.CacheCreationInputTokenCostAbove200kTokens = bifrost.Ptr(0.0000075) - p.CacheCreationInputTokenCostAbove1hr = bifrost.Ptr(0.000009) // distinct from above_200k to make fallback unambiguous - // CacheCreationInputTokenCostAbove1hrAbove200kTokens intentionally left nil - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 180000, - CompletionTokens: 25000, - TotalTokens: 205000, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedWriteTokens: 10000, - CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{ - CachedWriteTokens1h: 10000, - }, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Cache creation 1hr (no above_200k_1hr rate, uses above_1hr): 10000*0.000009 = 0.09 - // Input (non-cached): 170000*0.000006 = 1.02 - // Output: 25000*0.00003 = 0.75 - // Total: 1.02 + 0.09 + 0.75 = 1.86 - assert.InDelta(t, 1.86, cost, 1e-9) -} - -func TestComputeTextCost_1hrCacheCreationAbove200k_FallsBackToStandardAbove200kWhenNo1hrRates(t *testing.T) { - // When neither above_1hr field is present, 1hr tokens on a >200k request fall back - // to the standard above_200k cache creation rate. - p := chatPricing(8e-7, 0.000004) - p.InputCostPerTokenAbove200kTokens = bifrost.Ptr(0.0000016) - p.OutputCostPerTokenAbove200kTokens = bifrost.Ptr(0.000008) - p.CacheCreationInputTokenCost = bifrost.Ptr(0.000001) - p.CacheCreationInputTokenCostAbove200kTokens = bifrost.Ptr(0.000002) - // Neither CacheCreationInputTokenCostAbove1hr nor Above1hrAbove200k is set - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 180000, - CompletionTokens: 25000, - TotalTokens: 205000, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedWriteTokens: 10000, - CachedWriteTokenDetails: &schemas.ChatCachedWriteTokenDetails{ - CachedWriteTokens1h: 10000, - }, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Cache creation (1hr → no 1hr rates → standard above_200k): 10000*0.000002 = 0.02 - // Input (non-cached): 170000*0.0000016 = 0.272 - // Output: 25000*0.000008 = 0.2 - // Total: 0.272 + 0.02 + 0.2 = 0.492 - assert.InDelta(t, 0.492, cost, 1e-9) -} - -func TestComputeTextCost_Tiered200k(t *testing.T) { - // Claude 3.5 Sonnet Bedrock 200k tier: input=$6/M, output=$30/M - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenAbove200kTokens = bifrost.Ptr(0.000006) - p.OutputCostPerTokenAbove200kTokens = bifrost.Ptr(0.00003) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 180000, - CompletionTokens: 30000, - TotalTokens: 210000, // Above 200k threshold - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Uses tiered rate since total > 200k - // 180000 * 0.000006 + 30000 * 0.00003 = 1.08 + 0.90 = 1.98 - assert.InDelta(t, 1.98, cost, 1e-9) -} - -func TestComputeTextCost_Below200kUsesBaseRate(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenAbove200kTokens = bifrost.Ptr(0.000006) - p.OutputCostPerTokenAbove200kTokens = bifrost.Ptr(0.00003) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, // Below 200k - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Uses base rate since total < 200k - // 1000 * 0.000003 + 500 * 0.000015 = 0.003 + 0.0075 = 0.0105 - assert.InDelta(t, 0.0105, cost, 1e-12) -} - -func TestComputeTextCost_Tiered272k(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenAbove200kTokens = new(0.000006) - p.OutputCostPerTokenAbove200kTokens = new(0.00003) - p.InputCostPerTokenAbove272kTokens = new(0.000009) - p.OutputCostPerTokenAbove272kTokens = new(0.000045) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 250000, - CompletionTokens: 30000, - TotalTokens: 280000, // Above 272k threshold - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Uses 272k tiered rate since total > 272k - // 250000 * 0.000009 + 30000 * 0.000045 = 2.25 + 1.35 = 3.60 - assert.InDelta(t, 3.60, cost, 1e-9) -} - -func TestComputeTextCost_Between200kAnd272kUses200kRate(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenAbove200kTokens = new(0.000006) - p.OutputCostPerTokenAbove200kTokens = new(0.00003) - p.InputCostPerTokenAbove272kTokens = new(0.000009) - p.OutputCostPerTokenAbove272kTokens = new(0.000045) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 200000, - CompletionTokens: 30000, - TotalTokens: 230000, // Between 200k and 272k - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Uses 200k tiered rate since total > 200k but <= 272k - // 200000 * 0.000006 + 30000 * 0.00003 = 1.20 + 0.90 = 2.10 - assert.InDelta(t, 2.10, cost, 1e-9) -} - -func TestComputeTextCost_272kTierWithCacheRead(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenAbove272kTokens = new(0.000009) - p.OutputCostPerTokenAbove272kTokens = new(0.000045) - p.CacheReadInputTokenCost = new(0.0000003) - p.CacheReadInputTokenCostAbove272kTokens = new(0.0000009) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 250000, - CompletionTokens: 30000, - TotalTokens: 280000, // Above 272k - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedReadTokens: 50000, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Non-cached input: (250000-50000) * 0.000009 = 200000 * 0.000009 = 1.80 - // Cached read: 50000 * 0.0000009 = 0.045 - // Output: 30000 * 0.000045 = 1.35 - // Total: 1.80 + 0.045 + 1.35 = 3.195 - assert.InDelta(t, 3.195, cost, 1e-9) -} - -func TestComputeTextCost_SearchQueryCost(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.SearchContextCostPerQuery = bifrost.Ptr(0.01) // $0.01 per search query - - numQueries := 3 - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - CompletionTokensDetails: &schemas.ChatCompletionTokensDetails{ - NumSearchQueries: &numQueries, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // 1000*0.000003 + 500*0.000015 + 3*0.01 = 0.003 + 0.0075 + 0.03 = 0.0405 - assert.InDelta(t, 0.0405, cost, 1e-12) -} - -func TestComputeTextCost_NoCacheRateFallsBackToBaseInputRate(t *testing.T) { - // If cache rate fields are nil, tieredCacheReadInputTokenRate falls back to base InputCostPerToken - p := chatPricing(0.000005, 0.000015) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedReadTokens: 400, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Non-cached prompt: (1000-400)*0.000005 = 600*0.000005 = 0.003 - // Cached prompt: 400 tokens at base input rate (no cache rate set) = 400*0.000005 = 0.002 - // Output: 500*0.000015 = 0.0075 - // Total: 0.003 + 0.002 + 0.0075 = 0.0125 - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -// ========================================================================= -// 2. computeEmbeddingCost — unit tests -// ========================================================================= - -func TestComputeEmbeddingCost_Basic(t *testing.T) { - // Titan Embed Text v1: $0.1/M input - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.0000001), - OutputCostPerToken: bifrost.Ptr(0.0), - } - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 5000, - TotalTokens: 5000, - } - cost := computeEmbeddingCost(&p, usage, serviceTier{}) - // 5000 * 0.0000001 = 0.0005 - assert.InDelta(t, 0.0005, cost, 1e-12) -} - -func TestComputeEmbeddingCost_NilUsage(t *testing.T) { - p := configstoreTables.TableModelPricing{InputCostPerToken: new(0.0000001)} - assert.Equal(t, 0.0, computeEmbeddingCost(&p, nil, serviceTier{})) -} - -// ========================================================================= -// 3. computeRerankCost — unit tests -// ========================================================================= - -func TestComputeRerankCost_Basic(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000001), - OutputCostPerToken: bifrost.Ptr(0.000002), - } - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 2000, - CompletionTokens: 100, - TotalTokens: 2100, - } - cost := computeRerankCost(&p, usage, serviceTier{}) - // 2000*0.000001 + 100*0.000002 = 0.002 + 0.0002 = 0.0022 - assert.InDelta(t, 0.0022, cost, 1e-12) -} - -func TestComputeRerankCost_WithSearchCost(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.0), - OutputCostPerToken: bifrost.Ptr(0.0), - SearchContextCostPerQuery: bifrost.Ptr(0.001), - } - numQueries := 5 - usage := &schemas.BifrostLLMUsage{ - CompletionTokensDetails: &schemas.ChatCompletionTokensDetails{ - NumSearchQueries: &numQueries, - }, - } - cost := computeRerankCost(&p, usage, serviceTier{}) - assert.InDelta(t, 0.005, cost, 1e-12) -} - -func TestComputeRerankCost_NilUsage(t *testing.T) { - p := configstoreTables.TableModelPricing{InputCostPerToken: new(0.001)} - assert.Equal(t, 0.0, computeRerankCost(&p, nil, serviceTier{})) -} - -// ========================================================================= -// 4. computeSpeechCost — unit tests -// ========================================================================= - -func TestComputeSpeechCost_TokensPreferredOverDuration(t *testing.T) { - // TTS: input=text tokens, output=audio tokens (preferred over per-second) - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.0000025), - OutputCostPerToken: bifrost.Ptr(0.00001), - OutputCostPerSecond: bifrost.Ptr(0.00025), - } - seconds := 60 - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 100, - CompletionTokens: 200, - TotalTokens: 300, - } - cost := computeSpeechCost(&p, usage, &seconds, 0, serviceTier{}) - // Input: 100 text tokens * $0.0000025 = $0.00025 - // Output: 200 audio tokens present → uses token rate $0.00001, NOT per-second - // 200 * $0.00001 = $0.002 - // Total: $0.00225 - assert.InDelta(t, 0.00225, cost, 1e-12) -} - -func TestComputeSpeechCost_OutputFallsBackToPerSecond(t *testing.T) { - // TTS: no output tokens → falls back to per-second output pricing - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000001), - OutputCostPerToken: bifrost.Ptr(0.000002), - OutputCostPerSecond: bifrost.Ptr(0.0001), - } - seconds := 120 - usage := &schemas.BifrostLLMUsage{PromptTokens: 500} - cost := computeSpeechCost(&p, usage, &seconds, 0, serviceTier{}) - // Input: 500 * $0.000001 = $0.0005 - // Output: no CompletionTokens → falls back to 120 * $0.0001 = $0.012 - // Total: $0.0125 - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestComputeSpeechCost_OutputAudioTokenRate(t *testing.T) { - // TTS: output uses OutputCostPerAudioToken when available - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000001), - OutputCostPerToken: bifrost.Ptr(0.000002), - OutputCostPerAudioToken: bifrost.Ptr(0.00005), - } - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 200, - CompletionTokens: 100, - TotalTokens: 300, - } - cost := computeSpeechCost(&p, usage, nil, 0, serviceTier{}) - // Input: 200 * $0.000001 = $0.0002 - // Output: 100 * $0.00005 = $0.005 (OutputCostPerAudioToken preferred) - // Total: $0.0052 - assert.InDelta(t, 0.0052, cost, 1e-12) -} - -func TestComputeSpeechCost_TokenFallback(t *testing.T) { - p := chatPricing(0.000005, 0.000015) - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - } - cost := computeSpeechCost(&p, usage, nil, 0, serviceTier{}) // No audio seconds → token fallback - // 1000*0.000005 + 500*0.000015 = 0.005 + 0.0075 = 0.0125 - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestComputeSpeechCost_NilUsageNilSeconds(t *testing.T) { - p := chatPricing(0.000005, 0.000015) - assert.Equal(t, 0.0, computeSpeechCost(&p, nil, nil, 0, serviceTier{})) -} - -// ========================================================================= -// 5. computeTranscriptionCost — unit tests -// ========================================================================= - -func TestComputeTranscriptionCost_DurationBased(t *testing.T) { - // assemblyai/nano: input_cost_per_second=0.00010278 - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.0), - OutputCostPerToken: bifrost.Ptr(0.0), - InputCostPerSecond: bifrost.Ptr(0.00010278), - } - seconds := 300 // 5 minutes - cost := computeTranscriptionCost(&p, nil, &seconds, nil, serviceTier{}) - // 300 * 0.00010278 = 0.030834 - assert.InDelta(t, 0.030834, cost, 1e-9) -} - -func TestComputeTranscriptionCost_AudioTokenDetails(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000005), - OutputCostPerToken: bifrost.Ptr(0.000015), - InputCostPerAudioToken: bifrost.Ptr(0.00001), - } - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 2000, - CompletionTokens: 500, - TotalTokens: 2500, - } - audioDetails := &schemas.TranscriptionUsageInputTokenDetails{ - AudioTokens: 1500, - TextTokens: 500, - } - cost := computeTranscriptionCost(&p, usage, nil, audioDetails, serviceTier{}) - // Audio: 1500*0.00001 = 0.015 - // Text: 500*0.000005 = 0.0025 - // Output: 500*0.000015 = 0.0075 - // Total: 0.025 - assert.InDelta(t, 0.025, cost, 1e-12) -} - -func TestComputeTranscriptionCost_TokenFallback(t *testing.T) { - p := chatPricing(0.000005, 0.000015) - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 200, - TotalTokens: 1200, - } - cost := computeTranscriptionCost(&p, usage, nil, nil, serviceTier{}) - // 1000*0.000005 + 200*0.000015 = 0.005 + 0.003 = 0.008 - assert.InDelta(t, 0.008, cost, 1e-12) -} - -func TestComputeTranscriptionCost_TokenDetailsPreferredOverDuration(t *testing.T) { - // STT: audio token details present → uses tokens, not per-second - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000005), - OutputCostPerToken: bifrost.Ptr(0.0), - InputCostPerAudioPerSecond: bifrost.Ptr(0.0001), - InputCostPerAudioToken: bifrost.Ptr(0.00001), - } - seconds := 60 - audioDetails := &schemas.TranscriptionUsageInputTokenDetails{ - AudioTokens: 5000, - TextTokens: 1000, - } - cost := computeTranscriptionCost(&p, nil, &seconds, audioDetails, serviceTier{}) - // Input: audio token details present → tokens preferred over per-second - // 5000 audio * $0.00001 = $0.05 - // 1000 text * $0.000005 = $0.005 - // Output: nil usage → $0 - // Total: $0.055 - assert.InDelta(t, 0.055, cost, 1e-12) -} - -func TestComputeTranscriptionCost_DurationFallbackWhenNoTokens(t *testing.T) { - // STT: no audio token details, no prompt tokens → falls back to per-second - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000005), - OutputCostPerToken: bifrost.Ptr(0.000015), - InputCostPerAudioPerSecond: bifrost.Ptr(0.0001), - } - seconds := 60 - usage := &schemas.BifrostLLMUsage{ - CompletionTokens: 200, - TotalTokens: 200, - } - cost := computeTranscriptionCost(&p, usage, &seconds, nil, serviceTier{}) - // Input: no audio details, PromptTokens=0 → falls back to 60 * $0.0001 = $0.006 - // Output: 200 * $0.000015 = $0.003 - // Total: $0.009 - assert.InDelta(t, 0.009, cost, 1e-12) -} - -// ========================================================================= -// 6. computeImageCost — unit tests -// ========================================================================= - -func TestComputeImageCost_PerImage(t *testing.T) { - // dall-e-3 (aiml): output_cost_per_image=$0.052 - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.0), - OutputCostPerToken: bifrost.Ptr(0.0), - OutputCostPerImage: bifrost.Ptr(0.052), - } - usage := &schemas.ImageUsage{ - OutputTokensDetails: &schemas.ImageTokenDetails{ - NImages: 2, - }, - } - cost := computeImageCost(&p, usage, "", "", serviceTier{}) - // 2 * 0.052 = 0.104 - assert.InDelta(t, 0.104, cost, 1e-12) -} - -func TestComputeImageCost_PerImageDefaultsToOne(t *testing.T) { - p := configstoreTables.TableModelPricing{ - OutputCostPerImage: bifrost.Ptr(0.052), - } - usage := &schemas.ImageUsage{} // No token details → defaults to 1 image - cost := computeImageCost(&p, usage, "", "", serviceTier{}) - assert.InDelta(t, 0.052, cost, 1e-12) -} - -func TestComputeImageCost_TokenBased(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000005), - OutputCostPerToken: bifrost.Ptr(0.000015), - } - usage := &schemas.ImageUsage{ - InputTokens: 1000, - OutputTokens: 500, - TotalTokens: 1500, - } - cost := computeImageCost(&p, usage, "", "", serviceTier{}) - // 1000*0.000005 + 500*0.000015 = 0.005 + 0.0075 = 0.0125 - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestComputeImageCost_TokenBasedWithDetails(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000005), - OutputCostPerToken: bifrost.Ptr(0.000015), - } - usage := &schemas.ImageUsage{ - InputTokens: 2000, - OutputTokens: 1000, - TotalTokens: 3000, - InputTokensDetails: &schemas.ImageTokenDetails{ - TextTokens: 500, - ImageTokens: 1500, - }, - OutputTokensDetails: &schemas.ImageTokenDetails{ - TextTokens: 200, - ImageTokens: 800, - }, - } - cost := computeImageCost(&p, usage, "", "", serviceTier{}) - // Input: (500+1500)*0.000005 = 2000*0.000005 = 0.01 - // Output: (200+800)*0.000015 = 1000*0.000015 = 0.015 - // Total: 0.025 - assert.InDelta(t, 0.025, cost, 1e-12) -} - -func TestComputeImageCost_NilUsage(t *testing.T) { - p := configstoreTables.TableModelPricing{OutputCostPerImage: new(0.05)} - assert.Equal(t, 0.0, computeImageCost(&p, nil, "", "", serviceTier{})) -} - -func TestComputeImageCost_InputAndOutputPerImage(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerImage: bifrost.Ptr(0.01), - OutputCostPerImage: bifrost.Ptr(0.05), - } - usage := &schemas.ImageUsage{ - NumInputImages: 3, - OutputTokensDetails: &schemas.ImageTokenDetails{NImages: 2}, - } - cost := computeImageCost(&p, usage, "", "", serviceTier{}) - // 3 input * $0.01 + 2 output * $0.05 = $0.03 + $0.10 = $0.13 - assert.InDelta(t, 0.13, cost, 1e-12) -} - -func TestComputeImageCost_PerPixelOutput(t *testing.T) { - p := configstoreTables.TableModelPricing{ - OutputCostPerPixel: bifrost.Ptr(0.000000019), // ~$0.02 for 1024x1024 - } - usage := &schemas.ImageUsage{ - OutputTokensDetails: &schemas.ImageTokenDetails{NImages: 1}, - } - cost := computeImageCost(&p, usage, "1024x1024", "", serviceTier{}) - // 1024*1024 * 1 * 0.000000019 = 1048576 * 0.000000019 ≈ 0.01992 - assert.InDelta(t, 1048576*0.000000019, cost, 1e-12) -} - -func TestComputeImageCost_PerPixelInputAndOutput(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerPixel: bifrost.Ptr(0.00000001), - OutputCostPerPixel: bifrost.Ptr(0.00000002), - } - usage := &schemas.ImageUsage{ - NumInputImages: 2, - OutputTokensDetails: &schemas.ImageTokenDetails{NImages: 3}, - } - cost := computeImageCost(&p, usage, "512x512", "", serviceTier{}) - pixels := 512 * 512 // 262144 - // Input: 262144 * 2 * 0.00000001 = 0.00524288 - // Output: 262144 * 3 * 0.00000002 = 0.01572864 - expected := float64(pixels*2)*0.00000001 + float64(pixels*3)*0.00000002 - assert.InDelta(t, expected, cost, 1e-12) -} - -func TestComputeImageCost_TokensPreferredOverPixels(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000005), - OutputCostPerToken: bifrost.Ptr(0.000015), - InputCostPerPixel: bifrost.Ptr(0.00000001), - OutputCostPerPixel: bifrost.Ptr(0.00000002), - } - usage := &schemas.ImageUsage{ - InputTokens: 1000, - OutputTokens: 500, - TotalTokens: 1500, - } - cost := computeImageCost(&p, usage, "1024x1024", "", serviceTier{}) - // Tokens should win: 1000*0.000005 + 500*0.000015 = 0.0125 - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestComputeImageCost_PixelsPreferredOverPerImage(t *testing.T) { - p := configstoreTables.TableModelPricing{ - OutputCostPerPixel: bifrost.Ptr(0.00000002), - OutputCostPerImage: bifrost.Ptr(999.0), // should not be used - } - usage := &schemas.ImageUsage{ - OutputTokensDetails: &schemas.ImageTokenDetails{NImages: 1}, - } - cost := computeImageCost(&p, usage, "256x256", "", serviceTier{}) - // Per-pixel should win: 65536 * 1 * 0.00000002 = 0.00131072 - assert.InDelta(t, 65536*0.00000002, cost, 1e-12) -} - -func TestComputeImageCost_PerPixelFallsBackToPerImage_WhenNoSize(t *testing.T) { - p := configstoreTables.TableModelPricing{ - OutputCostPerPixel: bifrost.Ptr(0.00000002), - OutputCostPerImage: bifrost.Ptr(0.05), - } - usage := &schemas.ImageUsage{ - OutputTokensDetails: &schemas.ImageTokenDetails{NImages: 2}, - } - cost := computeImageCost(&p, usage, "", "", serviceTier{}) - // No size → pixels=0, falls through to per-image: 2 * $0.05 = $0.10 - assert.InDelta(t, 0.10, cost, 1e-12) -} - -func TestComputeImageCost_QualityBasedRates(t *testing.T) { - usage := &schemas.ImageUsage{ - OutputTokensDetails: &schemas.ImageTokenDetails{NImages: 1}, - } - // Quality-specific rates take precedence over base/size-tier - p := configstoreTables.TableModelPricing{ - OutputCostPerImage: bifrost.Ptr(0.01), - OutputCostPerImageLowQuality: bifrost.Ptr(0.02), - OutputCostPerImageMediumQuality: bifrost.Ptr(0.03), - OutputCostPerImageHighQuality: bifrost.Ptr(0.04), - OutputCostPerImageAutoQuality: bifrost.Ptr(0.05), - } - assert.InDelta(t, 0.02, computeImageCost(&p, usage, "", "low", serviceTier{}), 1e-12) - assert.InDelta(t, 0.03, computeImageCost(&p, usage, "", "medium", serviceTier{}), 1e-12) - assert.InDelta(t, 0.04, computeImageCost(&p, usage, "", "high", serviceTier{}), 1e-12) - assert.InDelta(t, 0.05, computeImageCost(&p, usage, "", "auto", serviceTier{}), 1e-12) - // "hd" does not match any quality case so perImageRate stays nil → size/base fallback. - assert.InDelta(t, 0.01, computeImageCost(&p, usage, "", "hd", serviceTier{}), 1e-12) - // Empty quality is treated as auto - assert.InDelta(t, 0.05, computeImageCost(&p, usage, "", "", serviceTier{}), 1e-12) -} - -func TestParseImagePixels(t *testing.T) { - assert.Equal(t, 1048576, parseImagePixels("1024x1024")) - assert.Equal(t, 262144, parseImagePixels("512x512")) - assert.Equal(t, 1835008, parseImagePixels("1792x1024")) - assert.Equal(t, 0, parseImagePixels("")) - assert.Equal(t, 0, parseImagePixels("invalid")) - assert.Equal(t, 0, parseImagePixels("1024")) - assert.Equal(t, 0, parseImagePixels("0x1024")) - assert.Equal(t, 0, parseImagePixels("-1x1024")) -} - -// ========================================================================= -// 7. computeVideoCost — unit tests -// ========================================================================= - -func TestComputeVideoCost_DurationBased(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000001), - OutputCostPerToken: bifrost.Ptr(0.0), - OutputCostPerVideoPerSecond: bifrost.Ptr(0.001), - } - seconds := 30 - usage := &schemas.BifrostLLMUsage{PromptTokens: 500, TotalTokens: 500} - cost := computeVideoCost(&p, usage, &seconds, serviceTier{}) - // Output: 30 * 0.001 = 0.03 - // Input: 500 * 0.000001 = 0.0005 - // Total: 0.0305 - assert.InDelta(t, 0.0305, cost, 1e-12) -} - -func TestComputeVideoCost_OutputCostPerSecondFallback(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.0), - OutputCostPerToken: bifrost.Ptr(0.0), - OutputCostPerSecond: bifrost.Ptr(0.002), - } - seconds := 10 - cost := computeVideoCost(&p, nil, &seconds, serviceTier{}) - assert.InDelta(t, 0.02, cost, 1e-12) -} - -func TestComputeVideoCost_NilSeconds(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000001), - OutputCostPerVideoPerSecond: bifrost.Ptr(0.001), - } - usage := &schemas.BifrostLLMUsage{PromptTokens: 1000} - cost := computeVideoCost(&p, usage, nil, serviceTier{}) - // Only input tokens: 1000 * 0.000001 = 0.001 - assert.InDelta(t, 0.001, cost, 1e-12) -} - -// ========================================================================= -// 8. tieredInputRate / tieredOutputRate -// ========================================================================= - -func TestTieredInputRate_BelowThreshold(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000003), - InputCostPerTokenAbove200kTokens: bifrost.Ptr(0.000006), - } - assert.Equal(t, 0.000003, tieredInputRate(&p, 100000, serviceTier{})) -} - -func TestTieredInputRate_AboveThreshold(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000003), - InputCostPerTokenAbove200kTokens: bifrost.Ptr(0.000006), - } - assert.Equal(t, 0.000006, tieredInputRate(&p, 210000, serviceTier{})) -} - -func TestTieredInputRate_AboveThresholdNoTieredRate(t *testing.T) { - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000003), - } - // Falls back to base rate when tiered field is nil - assert.Equal(t, 0.000003, tieredInputRate(&p, 300000, serviceTier{})) -} - -func TestTieredOutputRate_AboveThreshold(t *testing.T) { - p := configstoreTables.TableModelPricing{ - OutputCostPerToken: bifrost.Ptr(0.000015), - OutputCostPerTokenAbove200kTokens: bifrost.Ptr(0.00003), - } - assert.Equal(t, 0.00003, tieredOutputRate(&p, 250000, serviceTier{})) -} - -// ========================================================================= -// 9. extractCostInput — usage extraction -// ========================================================================= - -func TestExtractCostInput_ChatResponse(t *testing.T) { - usage := &schemas.BifrostLLMUsage{PromptTokens: 100, CompletionTokens: 50, TotalTokens: 150} - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{Usage: usage}, - } - input := extractCostInput(resp) - require.NotNil(t, input.usage) - assert.Equal(t, 100, input.usage.PromptTokens) - assert.Equal(t, 50, input.usage.CompletionTokens) -} - -func TestExtractCostInput_EmbeddingResponse(t *testing.T) { - usage := &schemas.BifrostLLMUsage{PromptTokens: 200, TotalTokens: 200} - resp := &schemas.BifrostResponse{ - EmbeddingResponse: &schemas.BifrostEmbeddingResponse{Usage: usage}, - } - input := extractCostInput(resp) - require.NotNil(t, input.usage) - assert.Equal(t, 200, input.usage.PromptTokens) -} - -func TestExtractCostInput_ImageResponse(t *testing.T) { - imgUsage := &schemas.ImageUsage{InputTokens: 100, OutputTokens: 200, TotalTokens: 300} - resp := &schemas.BifrostResponse{ - ImageGenerationResponse: &schemas.BifrostImageGenerationResponse{Usage: imgUsage}, - } - input := extractCostInput(resp) - assert.Nil(t, input.usage) - require.NotNil(t, input.imageUsage) - assert.Equal(t, 300, input.imageUsage.TotalTokens) -} - -func TestExtractCostInput_TranscriptionWithSeconds(t *testing.T) { - sec := 60 - resp := &schemas.BifrostResponse{ - TranscriptionResponse: &schemas.BifrostTranscriptionResponse{ - Usage: &schemas.TranscriptionUsage{ - Seconds: &sec, - InputTokens: bifrost.Ptr(1000), - OutputTokens: bifrost.Ptr(200), - TotalTokens: bifrost.Ptr(1200), - }, - }, - } - input := extractCostInput(resp) - require.NotNil(t, input.usage) - require.NotNil(t, input.audioSeconds) - assert.Equal(t, 60, *input.audioSeconds) - assert.Equal(t, 1000, input.usage.PromptTokens) -} - -func TestExtractCostInput_SpeechResponse(t *testing.T) { - resp := &schemas.BifrostResponse{ - SpeechResponse: &schemas.BifrostSpeechResponse{ - Usage: &schemas.SpeechUsage{ - InputTokens: 100, - OutputTokens: 500, - TotalTokens: 600, - }, - }, - } - input := extractCostInput(resp) - require.NotNil(t, input.usage) - assert.Equal(t, 100, input.usage.PromptTokens) - assert.Equal(t, 500, input.usage.CompletionTokens) - assert.Equal(t, 600, input.usage.TotalTokens) -} - -func TestExtractCostInput_VideoResponse(t *testing.T) { - sec := "15" - resp := &schemas.BifrostResponse{ - VideoGenerationResponse: &schemas.BifrostVideoGenerationResponse{ - Seconds: &sec, - }, - } - input := extractCostInput(resp) - require.NotNil(t, input.videoSeconds) - assert.Equal(t, 15, *input.videoSeconds) -} - -func TestExtractCostInput_VideoResponseInvalidSeconds(t *testing.T) { - sec := "not-a-number" - resp := &schemas.BifrostResponse{ - VideoGenerationResponse: &schemas.BifrostVideoGenerationResponse{ - Seconds: &sec, - }, - } - input := extractCostInput(resp) - assert.Nil(t, input.videoSeconds) -} - -// ========================================================================= -// 10. Semantic cache billing (calculateCostWithCache) -// ========================================================================= - -func TestCalculateCost_SemanticCacheDirectHit(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): { - Model: "gpt-4o", Provider: "openai", Mode: "chat", - InputCostPerToken: bifrost.Ptr(0.000005), OutputCostPerToken: bifrost.Ptr(0.000015), - }, - }) - - hitType := "direct" - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - Usage: &schemas.BifrostLLMUsage{PromptTokens: 100, CompletionTokens: 50, TotalTokens: 150}, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4o"), - CacheDebug: &schemas.BifrostCacheDebug{ - CacheHit: true, - HitType: &hitType, - }, - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - assert.Equal(t, 0.0, cost) -} - -func TestCalculateCost_SemanticCacheSemanticHit(t *testing.T) { - embProvider := "openai" - embModel := "text-embedding-3-small" - embTokens := 500 - - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): { - Model: "gpt-4o", Provider: "openai", Mode: "chat", - InputCostPerToken: bifrost.Ptr(0.000005), OutputCostPerToken: bifrost.Ptr(0.000015), - }, - makeKey("text-embedding-3-small", "openai", "embedding"): { - Model: "text-embedding-3-small", Provider: "openai", Mode: "embedding", - InputCostPerToken: bifrost.Ptr(0.00000002), - }, - }) - - hitType := "semantic" - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - Usage: &schemas.BifrostLLMUsage{PromptTokens: 100, CompletionTokens: 50, TotalTokens: 150}, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4o"), - CacheDebug: &schemas.BifrostCacheDebug{ - CacheHit: true, - HitType: &hitType, - ProviderUsed: &embProvider, - ModelUsed: &embModel, - InputTokens: &embTokens, - }, - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - // Only embedding cost: 500 * 0.00000002 = 0.00001 - assert.InDelta(t, 0.00001, cost, 1e-12) -} - -func TestCalculateCost_SemanticCacheMiss(t *testing.T) { - embProvider := "openai" - embModel := "text-embedding-3-small" - embTokens := 500 - - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): { - Model: "gpt-4o", Provider: "openai", Mode: "chat", - InputCostPerToken: bifrost.Ptr(0.000005), OutputCostPerToken: bifrost.Ptr(0.000015), - }, - makeKey("text-embedding-3-small", "openai", "embedding"): { - Model: "text-embedding-3-small", Provider: "openai", Mode: "embedding", - InputCostPerToken: bifrost.Ptr(0.00000002), - }, - }) - - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - Usage: &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500}, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4o"), - CacheDebug: &schemas.BifrostCacheDebug{ - CacheHit: false, - ProviderUsed: &embProvider, - ModelUsed: &embModel, - InputTokens: &embTokens, - }, - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - // Base cost: 1000*0.000005 + 500*0.000015 = 0.005 + 0.0075 = 0.0125 - // Embedding cost: 500 * 0.00000002 = 0.00001 - // Total: 0.01251 - assert.InDelta(t, 0.01251, cost, 1e-12) -} - -func TestCalculateCost_SemanticCacheHitNoEmbeddingInfo(t *testing.T) { - mc := testCatalogWithPricing(nil) - - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - ExtraFields: schemas.BifrostResponseExtraFields{ - CacheDebug: &schemas.BifrostCacheDebug{ - CacheHit: true, - // No ProviderUsed, ModelUsed, InputTokens - }, - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - assert.Equal(t, 0.0, cost) -} - -// ========================================================================= -// 11. CalculateCost integration — end-to-end -// ========================================================================= - -func TestCalculateCost_NilResponse(t *testing.T) { - mc := testCatalogWithPricing(nil) - assert.Equal(t, 0.0, mc.CalculateCost(nil, nil)) -} - -func TestCalculateCost_ProviderComputedCostPassthrough(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - - resp := makeChatResponse(schemas.OpenAI, "gpt-4o", &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - Cost: &schemas.BifrostCost{ - TotalCost: 0.99, // Provider already calculated - }, - }) - - cost := mc.CalculateCost(resp, nil) - assert.Equal(t, 0.99, cost) -} - -func TestCalculateCost_NoUsageData(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - - resp := makeChatResponse(schemas.OpenAI, "gpt-4o", nil) - cost := mc.CalculateCost(resp, nil) - assert.Equal(t, 0.0, cost) -} - -func TestCalculateCost_ChatCompletion_GPT4o(t *testing.T) { - // GPT-4o: $5/M input, $15/M output, cache_read=$0.5/M - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): { - Model: "gpt-4o", Provider: "openai", Mode: "chat", - InputCostPerToken: bifrost.Ptr(0.000005), - OutputCostPerToken: bifrost.Ptr(0.000015), - CacheReadInputTokenCost: bifrost.Ptr(0.0000005), - }, - }) - - resp := makeChatResponse(schemas.OpenAI, "gpt-4o", &schemas.BifrostLLMUsage{ - PromptTokens: 10000, - CompletionTokens: 2000, - TotalTokens: 12000, - }) - - cost := mc.CalculateCost(resp, nil) - // 10000*0.000005 + 2000*0.000015 = 0.05 + 0.03 = 0.08 - assert.InDelta(t, 0.08, cost, 1e-12) -} - -func TestCalculateCost_ChatCompletion_Claude35Sonnet_WithCache(t *testing.T) { - // Claude 3.5 Sonnet (Bedrock): $3/M input, $15/M output, cache_read=$0.3/M, cache_creation=$3.75/M - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("anthropic.claude-3-5-sonnet-20241022-v2:0", "bedrock", "chat"): { - Model: "anthropic.claude-3-5-sonnet-20241022-v2:0", Provider: "bedrock", Mode: "chat", - InputCostPerToken: bifrost.Ptr(0.000003), - OutputCostPerToken: bifrost.Ptr(0.000015), - CacheReadInputTokenCost: bifrost.Ptr(0.0000003), - CacheCreationInputTokenCost: bifrost.Ptr(0.00000375), - InputCostPerTokenAbove200kTokens: bifrost.Ptr(0.000006), - OutputCostPerTokenAbove200kTokens: bifrost.Ptr(0.00003), - }, - }) - - resp := makeChatResponse(schemas.Bedrock, "anthropic.claude-3-5-sonnet-20241022-v2:0", &schemas.BifrostLLMUsage{ - PromptTokens: 5000, - CompletionTokens: 1000, - TotalTokens: 6000, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedReadTokens: 3000, // 3000 cache read tokens - CachedWriteTokens: 500, // 500 cache creation tokens - }, - }) - - cost := mc.CalculateCost(resp, nil) - // Both cached read and write tokens are input-side deductions from promptTokens. - // Input: (5000-3000-500)*0.000003 + 3000*0.0000003 + 500*0.00000375 = 0.0045 + 0.0009 + 0.001875 = 0.007275 - // Output: 1000*0.000015 = 0.015 - // Total: 0.007275 + 0.015 = 0.022275 - assert.InDelta(t, 0.022275, cost, 1e-12) -} - -func TestCalculateCost_Embedding(t *testing.T) { - // Titan Embed Text v1: $0.1/M input - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("amazon.titan-embed-text-v1", "bedrock", "embedding"): { - Model: "amazon.titan-embed-text-v1", Provider: "bedrock", Mode: "embedding", - InputCostPerToken: bifrost.Ptr(0.0000001), - OutputCostPerToken: bifrost.Ptr(0.0), - }, - }) - - resp := makeEmbeddingResponse(schemas.Bedrock, "amazon.titan-embed-text-v1", &schemas.BifrostLLMUsage{ - PromptTokens: 10000, - TotalTokens: 10000, - }) - - cost := mc.CalculateCost(resp, nil) - // 10000 * 0.0000001 = 0.001 - assert.InDelta(t, 0.001, cost, 1e-12) -} - -func TestCalculateCost_Rerank(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("amazon.rerank-v1:0", "bedrock", "rerank"): { - Model: "amazon.rerank-v1:0", Provider: "bedrock", Mode: "rerank", - InputCostPerToken: bifrost.Ptr(0.0), - OutputCostPerToken: bifrost.Ptr(0.0), - }, - }) - - resp := makeRerankResponse(schemas.Bedrock, "amazon.rerank-v1:0", &schemas.BifrostLLMUsage{ - PromptTokens: 500, - TotalTokens: 500, - }) - - cost := mc.CalculateCost(resp, nil) - assert.Equal(t, 0.0, cost) -} - -func TestCalculateCost_ImageGeneration(t *testing.T) { - // dall-e-3 via aiml: output_cost_per_image=$0.052 - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("dall-e-3", "aiml", "image_generation"): { - Model: "dall-e-3", Provider: "aiml", Mode: "image_generation", - OutputCostPerImage: bifrost.Ptr(0.052), - }, - }) - - resp := makeImageResponse("aiml", "dall-e-3", &schemas.ImageUsage{ - OutputTokensDetails: &schemas.ImageTokenDetails{NImages: 3}, - }) - - cost := mc.CalculateCost(resp, nil) - // 3 * 0.052 = 0.156 - assert.InDelta(t, 0.156, cost, 1e-12) -} - -func TestCalculateCost_StreamRequestTypeNormalized(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - - // Stream request type should be normalized to base type - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - Usage: &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500}, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionStreamRequest, - RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4o"), - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestCalculateCost_WebSocketResponsesFallsBackToChatPricing(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - - resp := &schemas.BifrostResponse{ - ResponsesStreamResponse: &schemas.BifrostResponsesStreamResponse{ - Response: &schemas.BifrostResponsesResponse{ - Usage: &schemas.ResponsesResponseUsage{InputTokens: 1000, OutputTokens: 500, TotalTokens: 1500}, - }, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.WebSocketResponsesRequest, - RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4o"), - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestCalculateCost_NoPricingData(t *testing.T) { - mc := testCatalogWithPricing(nil) - resp := makeChatResponse(schemas.OpenAI, "unknown-model", &schemas.BifrostLLMUsage{ - PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500, - }) - cost := mc.CalculateCost(resp, nil) - assert.Equal(t, 0.0, cost) -} - -// ========================================================================= -// 12. Pricing resolution — getPricing fallback logic -// ========================================================================= - -func TestGetPricing_DirectLookup(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "openai"}) - assert.Equal(t, 0.000005, derefF(p.InputCostPerToken)) -} - -func TestGetPricing_GeminiFallsBackToVertex(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gemini-2.0-flash", "vertex", "chat"): { - Model: "gemini-2.0-flash", Provider: "vertex", Mode: "chat", - InputCostPerToken: bifrost.Ptr(0.0000001), OutputCostPerToken: bifrost.Ptr(0.0000004), - }, - }) - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "gemini", Model: "gemini-2.0-flash"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "gemini"}) - assert.Equal(t, 0.0000001, derefF(p.InputCostPerToken)) -} - -func TestGetPricing_VertexStripsProviderPrefix(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gemini-2.0-flash", "vertex", "chat"): chatPricing(0.0000001, 0.0000004), - }) - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "vertex", Model: "google/gemini-2.0-flash"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "vertex"}) - assert.Equal(t, 0.0000001, derefF(p.InputCostPerToken)) -} - -func TestGetPricing_BedrockAddsAnthropicPrefix(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("anthropic.claude-3-5-sonnet-20241022-v2:0", "bedrock", "chat"): chatPricing(0.000003, 0.000015), - }) - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "bedrock", Model: "claude-3-5-sonnet-20241022-v2:0"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "bedrock"}) - assert.Equal(t, 0.000003, derefF(p.InputCostPerToken)) -} - -func TestGetPricing_ResponsesFallsBackToChat(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o"}, schemas.ResponsesRequest, PricingLookupScopes{Provider: "openai"}) - assert.Equal(t, 0.000005, derefF(p.InputCostPerToken)) -} - -func TestGetPricing_ResponsesStreamFallsBackToChat(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o"}, schemas.ResponsesStreamRequest, PricingLookupScopes{Provider: "openai"}) - assert.Equal(t, 0.000005, derefF(p.InputCostPerToken)) -} - -func TestGetPricing_RealtimeFallsBackToChat(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o"}, schemas.RealtimeRequest, PricingLookupScopes{Provider: "openai"}) - assert.Equal(t, 0.000005, derefF(p.InputCostPerToken)) -} - -func TestGetPricing_GeminiResponsesFallsBackToVertexChat(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gemini-2.0-flash", "vertex", "chat"): chatPricing(0.0000001, 0.0000004), - }) - // gemini provider + responses request → try vertex + responses → try vertex + chat - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "gemini", Model: "gemini-2.0-flash"}, schemas.ResponsesRequest, PricingLookupScopes{Provider: "gemini"}) - assert.Equal(t, 0.0000001, derefF(p.InputCostPerToken)) -} - -func TestGetPricing_NotFound(t *testing.T) { - mc := testCatalogWithPricing(nil) - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "nonexistent"}, schemas.ChatCompletionRequest, PricingLookupScopes{Provider: "openai"}) - assert.Nil(t, p) -} - -// ========================================================================= -// 13. resolvePricing — deployment fallback -// ========================================================================= - -func TestResolvePricing_DeploymentFallback(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("my-deployment", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - - // Model not found directly, but deployment matches - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o-custom", ResolvedKeyAlias: &schemas.ResolvedKeyAlias{ModelID: "my-deployment"}}, schemas.ChatCompletionRequest, PricingLookupScopes{}) - require.NotNil(t, p) - assert.Equal(t, 0.000005, derefF(p.InputCostPerToken)) -} - -func TestResolvePricing_ResolvedModelHasPriority(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - makeKey("my-deployment", "openai", "chat"): chatPricing(0.000001, 0.000002), - }) - - // Resolved model ("my-deployment") is looked up first and has priority - // over the originally requested model ("gpt-4o"). - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "gpt-4o", ResolvedKeyAlias: &schemas.ResolvedKeyAlias{ModelID: "my-deployment"}}, schemas.ChatCompletionRequest, PricingLookupScopes{}) - require.NotNil(t, p) - assert.Equal(t, 0.000001, derefF(p.InputCostPerToken)) -} - -func TestResolvePricing_NothingFound(t *testing.T) { - mc := testCatalogWithPricing(nil) - p := mc.resolvePricing(schemas.RoutingInfo{Provider: "openai", Model: "unknown"}, schemas.ChatCompletionRequest, PricingLookupScopes{}) - assert.Nil(t, p) -} - -// ========================================================================= -// 14. normalizeStreamRequestType -// ========================================================================= - -func TestNormalizeStreamRequestType(t *testing.T) { - tests := []struct { - input schemas.RequestType - expected schemas.RequestType - }{ - {schemas.ChatCompletionStreamRequest, schemas.ChatCompletionRequest}, - {schemas.TextCompletionStreamRequest, schemas.TextCompletionRequest}, - {schemas.ResponsesStreamRequest, schemas.ResponsesRequest}, - {schemas.SpeechStreamRequest, schemas.SpeechRequest}, - {schemas.TranscriptionStreamRequest, schemas.TranscriptionRequest}, - {schemas.ImageGenerationStreamRequest, schemas.ImageGenerationRequest}, - {schemas.ImageEditStreamRequest, schemas.ImageEditRequest}, - {schemas.RealtimeRequest, schemas.RealtimeRequest}, // realtime is its own base type - {schemas.ChatCompletionRequest, schemas.ChatCompletionRequest}, // non-stream unchanged - {schemas.EmbeddingRequest, schemas.EmbeddingRequest}, // non-stream unchanged - } - - for _, tt := range tests { - assert.Equal(t, tt.expected, normalizeStreamRequestType(tt.input), "for input %s", tt.input) - } -} - -// ========================================================================= -// 15. responsesUsageToBifrostUsage -// ========================================================================= - -func TestResponsesUsageToBifrostUsage_Basic(t *testing.T) { - u := &schemas.ResponsesResponseUsage{ - InputTokens: 100, - OutputTokens: 50, - TotalTokens: 150, - } - result := responsesUsageToBifrostUsage(u) - assert.Equal(t, 100, result.PromptTokens) - assert.Equal(t, 50, result.CompletionTokens) - assert.Equal(t, 150, result.TotalTokens) - assert.Nil(t, result.PromptTokensDetails) - assert.Nil(t, result.CompletionTokensDetails) -} - -func TestResponsesUsageToBifrostUsage_WithTokenDetails(t *testing.T) { - numQueries := 2 - u := &schemas.ResponsesResponseUsage{ - InputTokens: 1000, - OutputTokens: 500, - TotalTokens: 1500, - InputTokensDetails: &schemas.ResponsesResponseInputTokens{ - CachedReadTokens: 300, - CachedWriteTokens: 50, - TextTokens: 600, - AudioTokens: 50, - ImageTokens: 50, - }, - OutputTokensDetails: &schemas.ResponsesResponseOutputTokens{ - ReasoningTokens: 100, - NumSearchQueries: &numQueries, - }, - } - result := responsesUsageToBifrostUsage(u) - - require.NotNil(t, result.PromptTokensDetails) - assert.Equal(t, 300, result.PromptTokensDetails.CachedReadTokens) - assert.Equal(t, 50, result.PromptTokensDetails.CachedWriteTokens) - assert.Equal(t, 600, result.PromptTokensDetails.TextTokens) - assert.Equal(t, 50, result.PromptTokensDetails.AudioTokens) - assert.Equal(t, 50, result.PromptTokensDetails.ImageTokens) - - require.NotNil(t, result.CompletionTokensDetails) - assert.Equal(t, 100, result.CompletionTokensDetails.ReasoningTokens) - require.NotNil(t, result.CompletionTokensDetails.NumSearchQueries) - assert.Equal(t, 2, *result.CompletionTokensDetails.NumSearchQueries) -} - -// ========================================================================= -// 16. Edge cases -// ========================================================================= - -func TestCalculateCost_200kTier_EndToEnd(t *testing.T) { - // Claude 3.5 Sonnet Bedrock with 200k tier pricing - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("anthropic.claude-3-5-sonnet-20240620-v1:0", "bedrock", "chat"): { - Model: "anthropic.claude-3-5-sonnet-20240620-v1:0", Provider: "bedrock", Mode: "chat", - InputCostPerToken: bifrost.Ptr(0.000003), - OutputCostPerToken: bifrost.Ptr(0.000015), - InputCostPerTokenAbove200kTokens: bifrost.Ptr(0.000006), - OutputCostPerTokenAbove200kTokens: bifrost.Ptr(0.00003), - CacheReadInputTokenCost: bifrost.Ptr(0.0000003), - CacheCreationInputTokenCost: bifrost.Ptr(0.00000375), - CacheReadInputTokenCostAbove200kTokens: bifrost.Ptr(0.0000006), - CacheCreationInputTokenCostAbove200kTokens: bifrost.Ptr(0.0000075), - }, - }) - - resp := makeChatResponse(schemas.Bedrock, "anthropic.claude-3-5-sonnet-20240620-v1:0", &schemas.BifrostLLMUsage{ - PromptTokens: 190000, - CompletionTokens: 20000, - TotalTokens: 210000, // Above 200k - }) - - cost := mc.CalculateCost(resp, nil) - // Tiered rate: input=0.000006, output=0.00003 - // 190000*0.000006 + 20000*0.00003 = 1.14 + 0.6 = 1.74 - assert.InDelta(t, 1.74, cost, 1e-9) -} - -func TestCalculateCost_272kTier_EndToEnd(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("claude-3-7-sonnet", "anthropic", "chat"): { - Model: "claude-3-7-sonnet", - Provider: "anthropic", - Mode: "chat", - InputCostPerToken: new(0.000003), - OutputCostPerToken: new(0.000015), - InputCostPerTokenAbove200kTokens: new(0.000006), - OutputCostPerTokenAbove200kTokens: new(0.00003), - InputCostPerTokenAbove272kTokens: new(0.000009), - OutputCostPerTokenAbove272kTokens: new(0.000045), - CacheReadInputTokenCost: new(0.0000003), - CacheReadInputTokenCostAbove200kTokens: new(0.0000006), - CacheReadInputTokenCostAbove272kTokens: new(0.0000009), - }, - }) - - resp := makeChatResponse(schemas.Anthropic, "claude-3-7-sonnet", &schemas.BifrostLLMUsage{ - PromptTokens: 250000, - CompletionTokens: 30000, - TotalTokens: 280000, // Above 272k - }) - - cost := mc.CalculateCost(resp, nil) - // Tiered rate: input=0.000009, output=0.000045 - // 250000*0.000009 + 30000*0.000045 = 2.25 + 1.35 = 3.60 - assert.InDelta(t, 3.60, cost, 1e-9) -} - -func TestCalculateCost_272kTier_CacheReadFallbackChain(t *testing.T) { - // Verifies the 272k cache read rate takes precedence over 200k and base rates - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("claude-3-7-sonnet", "anthropic", "chat"): { - Model: "claude-3-7-sonnet", - Provider: "anthropic", - Mode: "chat", - InputCostPerToken: new(0.000003), - OutputCostPerToken: new(0.000015), - InputCostPerTokenAbove272kTokens: new(0.000009), - OutputCostPerTokenAbove272kTokens: new(0.000045), - CacheReadInputTokenCost: new(0.0000003), - CacheReadInputTokenCostAbove200kTokens: new(0.0000006), - CacheReadInputTokenCostAbove272kTokens: new(0.0000009), - }, - }) - - resp := makeChatResponse(schemas.Anthropic, "claude-3-7-sonnet", &schemas.BifrostLLMUsage{ - PromptTokens: 250000, - CompletionTokens: 30000, - TotalTokens: 280000, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedReadTokens: 50000, - }, - }) - - cost := mc.CalculateCost(resp, nil) - // Non-cached input: (250000-50000) * 0.000009 = 200000 * 0.000009 = 1.80 - // Cached read (272k rate): 50000 * 0.0000009 = 0.045 - // Output: 30000 * 0.000045 = 1.35 - // Total: 1.80 + 0.045 + 1.35 = 3.195 - assert.InDelta(t, 3.195, cost, 1e-9) -} - -// ========================================================================= -// Priority tier tests -// ========================================================================= - -func TestComputeTextCost_PriorityUsesInputOutputPriorityRate(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenPriority = new(0.000006) - p.OutputCostPerTokenPriority = new(0.00003) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - } - - cost := computeTextCost(&p, usage, serviceTier{isPriority: true}) - - // Uses priority rates: 1000*0.000006 + 500*0.00003 = 0.006 + 0.015 = 0.021 - assert.InDelta(t, 0.021, cost, 1e-12) -} - -func TestComputeTextCost_NonPriorityIgnoresPriorityRate(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenPriority = new(0.000006) - p.OutputCostPerTokenPriority = new(0.00003) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Uses base rates, ignores priority fields: 1000*0.000003 + 500*0.000015 = 0.003 + 0.0075 = 0.0105 - assert.InDelta(t, 0.0105, cost, 1e-12) -} - -func TestComputeTextCost_Priority272kTier(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenPriority = new(0.000006) - p.OutputCostPerTokenPriority = new(0.00003) - p.InputCostPerTokenAbove272kTokens = new(0.000009) - p.InputCostPerTokenAbove272kTokensPriority = new(0.000012) - p.OutputCostPerTokenAbove272kTokens = new(0.000045) - p.OutputCostPerTokenAbove272kTokensPriority = new(0.00006) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 250000, - CompletionTokens: 30000, - TotalTokens: 280000, - } - - cost := computeTextCost(&p, usage, serviceTier{isPriority: true}) - - // Uses 272k priority rates: 250000*0.000012 + 30000*0.00006 = 3.00 + 1.80 = 4.80 - assert.InDelta(t, 4.80, cost, 1e-9) -} - -func TestComputeTextCost_Priority272kTierFallsBackToNonPriority272k(t *testing.T) { - // Priority flag set but no priority-specific 272k rate — fall back to non-priority 272k - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenAbove272kTokens = new(0.000009) - p.OutputCostPerTokenAbove272kTokens = new(0.000045) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 250000, - CompletionTokens: 30000, - TotalTokens: 280000, - } - - cost := computeTextCost(&p, usage, serviceTier{isPriority: true}) - - // Falls back to non-priority 272k rate: 250000*0.000009 + 30000*0.000045 = 2.25 + 1.35 = 3.60 - assert.InDelta(t, 3.60, cost, 1e-9) -} - -func TestComputeTextCost_PriorityCacheReadRate(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenPriority = new(0.000006) - p.OutputCostPerTokenPriority = new(0.00003) - p.CacheReadInputTokenCost = new(0.0000003) - p.CacheReadInputTokenCostPriority = new(0.0000006) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedReadTokens: 400, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{isPriority: true}) - - // Non-cached input: (1000-400)*0.000006 = 600*0.000006 = 0.0036 - // Cached read (priority rate): 400*0.0000006 = 0.00024 - // Output: 500*0.00003 = 0.015 - // Total: 0.0036 + 0.00024 + 0.015 = 0.01884 - assert.InDelta(t, 0.01884, cost, 1e-12) -} - -func TestCalculateCost_PriorityTier_EndToEnd(t *testing.T) { - tier := schemas.BifrostServiceTierPriority - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): { - Model: "gpt-4o", - Provider: "openai", - Mode: "chat", - InputCostPerToken: new(0.000005), - OutputCostPerToken: new(0.000015), - InputCostPerTokenPriority: new(0.000010), - OutputCostPerTokenPriority: new(0.000030), - }, - }) - - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - ServiceTier: &tier, - Usage: &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - }, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4o"), - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - // Priority rates: 1000*0.000010 + 500*0.000030 = 0.010 + 0.015 = 0.025 - assert.InDelta(t, 0.025, cost, 1e-12) -} - -func TestCalculateCost_NonPriorityServiceTier_UsesBaseRate(t *testing.T) { - tier := schemas.BifrostServiceTierAuto - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): { - Model: "gpt-4o", - Provider: "openai", - Mode: "chat", - InputCostPerToken: new(0.000005), - OutputCostPerToken: new(0.000015), - InputCostPerTokenPriority: new(0.000010), - OutputCostPerTokenPriority: new(0.000030), - }, - }) - - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - ServiceTier: &tier, - Usage: &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - }, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4o"), - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - // Base rates (not priority): 1000*0.000005 + 500*0.000015 = 0.005 + 0.0075 = 0.0125 - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestTieredCacheReadRate_FallbackOrder(t *testing.T) { - // 272k rate takes precedence over 200k, 200k over base, base over input rate - t.Run("uses_272k_when_above_272k", func(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCost = new(0.0000003) - p.CacheReadInputTokenCostAbove200kTokens = new(0.0000006) - p.CacheReadInputTokenCostAbove272kTokens = new(0.0000009) - assert.Equal(t, 0.0000009, tieredCacheReadInputTokenRate(&p, 280000, serviceTier{})) - }) - t.Run("uses_200k_when_between_200k_and_272k", func(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCost = new(0.0000003) - p.CacheReadInputTokenCostAbove200kTokens = new(0.0000006) - p.CacheReadInputTokenCostAbove272kTokens = new(0.0000009) - assert.Equal(t, 0.0000006, tieredCacheReadInputTokenRate(&p, 230000, serviceTier{})) - }) - t.Run("uses_base_cache_rate_when_below_200k", func(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCost = new(0.0000003) - p.CacheReadInputTokenCostAbove200kTokens = new(0.0000006) - p.CacheReadInputTokenCostAbove272kTokens = new(0.0000009) - assert.Equal(t, 0.0000003, tieredCacheReadInputTokenRate(&p, 1500, serviceTier{})) - }) - t.Run("falls_back_to_input_rate_when_no_cache_rate_set", func(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - // No cache rates set at all - assert.Equal(t, 0.000003, tieredCacheReadInputTokenRate(&p, 280000, serviceTier{})) - }) - t.Run("priority_uses_272k_priority_rate", func(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCost = new(0.0000003) - p.CacheReadInputTokenCostPriority = new(0.0000006) - p.CacheReadInputTokenCostAbove272kTokens = new(0.0000009) - p.CacheReadInputTokenCostAbove272kTokensPriority = new(0.0000012) - assert.Equal(t, 0.0000012, tieredCacheReadInputTokenRate(&p, 280000, serviceTier{isPriority: true})) - }) - t.Run("priority_falls_back_to_272k_non_priority_when_priority_rate_missing", func(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCostAbove272kTokens = new(0.0000009) - assert.Equal(t, 0.0000009, tieredCacheReadInputTokenRate(&p, 280000, serviceTier{isPriority: true})) - }) - t.Run("priority_uses_priority_base_cache_rate_below_tiers", func(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCost = new(0.0000003) - p.CacheReadInputTokenCostPriority = new(0.0000006) - assert.Equal(t, 0.0000006, tieredCacheReadInputTokenRate(&p, 1500, serviceTier{isPriority: true})) - }) - t.Run("flex_uses_flex_cache_rate", func(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCost = new(0.0000003) - p.CacheReadInputTokenCostFlex = new(0.0000005) - assert.Equal(t, 0.0000005, tieredCacheReadInputTokenRate(&p, 1500, serviceTier{isFlex: true})) - }) - t.Run("flex_uses_flex_cache_rate_regardless_of_token_count", func(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCost = new(0.0000003) - p.CacheReadInputTokenCostFlex = new(0.0000005) - p.CacheReadInputTokenCostAbove272kTokens = new(0.0000009) - // Even above 272k, flex flat rate takes precedence - assert.Equal(t, 0.0000005, tieredCacheReadInputTokenRate(&p, 280000, serviceTier{isFlex: true})) - }) - t.Run("flex_falls_back_to_base_cache_rate_when_no_flex_cache_rate", func(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCost = new(0.0000003) - // No flex cache rate — falls back to base cache rate - assert.Equal(t, 0.0000003, tieredCacheReadInputTokenRate(&p, 1500, serviceTier{isFlex: true})) - }) - t.Run("flex_wins_over_272k_priority_and_priority_base_when_all_present", func(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.CacheReadInputTokenCostAbove272kTokens = new(5e-7) - p.CacheReadInputTokenCostFlex = new(1.3e-7) - p.CacheReadInputTokenCostPriority = new(5e-7) - p.CacheReadInputTokenCostAbove272kTokensPriority = new(0.000001) - // token count exceeds 272k — but flex flat rate should still win - assert.Equal(t, 1.3e-7, tieredCacheReadInputTokenRate(&p, 280000, serviceTier{isFlex: true})) - }) -} - -// ========================================================================= -// tierFromString tests -// ========================================================================= - -func TestTierFromString_Priority(t *testing.T) { - s := schemas.BifrostServiceTierPriority - tier := tierFromString(&s) - assert.True(t, tier.isPriority) - assert.False(t, tier.isFlex) -} - -func TestTierFromString_Flex(t *testing.T) { - s := schemas.BifrostServiceTierFlex - tier := tierFromString(&s) - assert.False(t, tier.isPriority) - assert.True(t, tier.isFlex) -} - -func TestTierFromString_Default(t *testing.T) { - for _, s := range []schemas.BifrostServiceTier{schemas.BifrostServiceTierAuto, schemas.BifrostServiceTierDefault, ""} { - tier := tierFromString(&s) - assert.False(t, tier.isPriority, "expected no priority for %q", s) - assert.False(t, tier.isFlex, "expected no flex for %q", s) - } -} - -func TestTierFromString_Nil(t *testing.T) { - tier := tierFromString(nil) - assert.False(t, tier.isPriority) - assert.False(t, tier.isFlex) -} - -// ========================================================================= -// Flex tier tests -// ========================================================================= - -func TestComputeTextCost_FlexUsesFlexRate(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenFlex = new(0.0000015) - p.OutputCostPerTokenFlex = new(0.0000075) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - } - - cost := computeTextCost(&p, usage, serviceTier{isFlex: true}) - - // Flex rates: 1000*0.0000015 + 500*0.0000075 = 0.0015 + 0.00375 = 0.00525 - assert.InDelta(t, 0.00525, cost, 1e-12) -} - -func TestComputeTextCost_NonFlexIgnoresFlexRate(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenFlex = new(0.0000015) - p.OutputCostPerTokenFlex = new(0.0000075) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - - // Base rates, flex fields ignored: 1000*0.000003 + 500*0.000015 = 0.003 + 0.0075 = 0.0105 - assert.InDelta(t, 0.0105, cost, 1e-12) -} - -func TestComputeTextCost_FlexIgnoresTokenTiers(t *testing.T) { - // Flex is a flat rate — token-count tiers (272k, 200k, 128k) do not apply. - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenFlex = new(0.0000015) - p.OutputCostPerTokenFlex = new(0.0000075) - p.InputCostPerTokenAbove272kTokens = new(0.000009) - p.OutputCostPerTokenAbove272kTokens = new(0.000045) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 250000, - CompletionTokens: 30000, - TotalTokens: 280000, - } - - cost := computeTextCost(&p, usage, serviceTier{isFlex: true}) - - // Flex flat rate overrides 272k tier: 250000*0.0000015 + 30000*0.0000075 = 0.375 + 0.225 = 0.60 - assert.InDelta(t, 0.60, cost, 1e-9) -} - -func TestComputeTextCost_FlexCacheReadRate(t *testing.T) { - p := chatPricing(0.000003, 0.000015) - p.InputCostPerTokenFlex = new(0.0000015) - p.OutputCostPerTokenFlex = new(0.0000075) - p.CacheReadInputTokenCost = new(0.0000003) - p.CacheReadInputTokenCostFlex = new(0.0000006) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedReadTokens: 400, - }, - } - - cost := computeTextCost(&p, usage, serviceTier{isFlex: true}) - - // Non-cached input: (1000-400)*0.0000015 = 600*0.0000015 = 0.0009 - // Cached read (flex rate): 400*0.0000006 = 0.00024 - // Output: 500*0.0000075 = 0.00375 - // Total: 0.0009 + 0.00024 + 0.00375 = 0.00489 - assert.InDelta(t, 0.00489, cost, 1e-12) -} - -func TestComputeTextCost_FlexFallsBackToBaseWhenNoFlexRate(t *testing.T) { - // isFlex set but no flex fields configured — falls back to base rates. - p := chatPricing(0.000003, 0.000015) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - } - - cost := computeTextCost(&p, usage, serviceTier{isFlex: true}) - - // Base rates used as fallback: 1000*0.000003 + 500*0.000015 = 0.003 + 0.0075 = 0.0105 - assert.InDelta(t, 0.0105, cost, 1e-12) -} - -func TestCalculateCost_FlexTier_EndToEnd(t *testing.T) { - tier := schemas.BifrostServiceTierFlex - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): { - Model: "gpt-4o", - Provider: "openai", - Mode: "chat", - InputCostPerToken: new(0.000005), - OutputCostPerToken: new(0.000015), - InputCostPerTokenFlex: new(0.0000025), - OutputCostPerTokenFlex: new(0.0000075), - }, - }) - - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - ServiceTier: &tier, - Usage: &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - }, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4o"), - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - // Flex rates: 1000*0.0000025 + 500*0.0000075 = 0.0025 + 0.00375 = 0.00625 - assert.InDelta(t, 0.00625, cost, 1e-12) -} - -func TestCalculateCost_FlexTier_FallsBackToBaseWhenNoFlexRate(t *testing.T) { - tier := schemas.BifrostServiceTierFlex - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - ServiceTier: &tier, - Usage: &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - }, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4o"), - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - // No flex rates configured — falls back to base: 1000*0.000005 + 500*0.000015 = 0.005 + 0.0075 = 0.0125 - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestCalculateCost_ProviderCostZeroTotalStillCalculates(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - - // Provider cost present but TotalCost is 0 → our calculation runs - resp := makeChatResponse(schemas.OpenAI, "gpt-4o", &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 500, - TotalTokens: 1500, - Cost: &schemas.BifrostCost{ - TotalCost: 0, - }, - }) - - cost := mc.CalculateCost(resp, nil) - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestCalculateCost_AllCachedTokens(t *testing.T) { - // All prompt tokens are from cache - p := chatPricing(0.000005, 0.000015) - p.CacheReadInputTokenCost = bifrost.Ptr(0.0000005) - - usage := &schemas.BifrostLLMUsage{ - PromptTokens: 1000, - CompletionTokens: 0, - TotalTokens: 1000, - PromptTokensDetails: &schemas.ChatPromptTokensDetails{ - CachedReadTokens: 1000, // All cached - }, - } - - cost := computeTextCost(&p, usage, serviceTier{}) - // Non-cached: 0, cached: 1000*0.0000005 = 0.0005 - assert.InDelta(t, 0.0005, cost, 1e-12) -} - -// ========================================================================= -// Nil usage fallbacks — per-unit pricing when no token data is reported -// ========================================================================= - -func TestCalculateCost_ImageGeneration_NilUsage_PerImagePricing(t *testing.T) { - // Image response exists but Usage is nil — should default to 1 image with per-image pricing - pricing := configstoreTables.TableModelPricing{ - Model: "dall-e-3", - Provider: "openai", - Mode: "image_generation", - InputCostPerToken: bifrost.Ptr(0.0), - OutputCostPerImage: bifrost.Ptr(0.04), - } - - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("dall-e-3", "openai", "image_generation"): pricing, - }) - - resp := makeImageResponse("openai", "dall-e-3", nil) - cost := mc.CalculateCost(resp, nil) - // 1 image * $0.04 = $0.04 - assert.InDelta(t, 0.04, cost, 1e-12) -} - -func TestCalculateCost_ImageGeneration_NilUsage_InputAndOutputPerImage(t *testing.T) { - // Both input and output per-image pricing, but no NumInputImages set - pricing := configstoreTables.TableModelPricing{ - Model: "test-image-model", - Provider: "test", - Mode: "image_generation", - InputCostPerImage: bifrost.Ptr(0.01), - OutputCostPerImage: bifrost.Ptr(0.04), - } - - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("test-image-model", "test", "image_generation"): pricing, - }) - - resp := makeImageResponse("test", "test-image-model", nil) - cost := mc.CalculateCost(resp, nil) - // NumInputImages is 0 (not populated from request), so only output pricing applies - // 1 output image * $0.04 = $0.04 - assert.InDelta(t, 0.04, cost, 1e-12) -} - -func TestCalculateCost_ImageGeneration_WithInputImages(t *testing.T) { - // Input + output per-image pricing with NumInputImages populated from request - pricing := configstoreTables.TableModelPricing{ - Model: "gpt-image-1", - Provider: "openai", - Mode: "image_generation", - InputCostPerImage: bifrost.Ptr(0.01), - OutputCostPerImage: bifrost.Ptr(0.04), - } - - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-image-1", "openai", "image_generation"): pricing, - }) - - resp := makeImageResponse("openai", "gpt-image-1", &schemas.ImageUsage{ - NumInputImages: 2, - }) - cost := mc.CalculateCost(resp, nil) - // 2 input images * $0.01 + 1 output image * $0.04 = $0.06 - assert.InDelta(t, 0.06, cost, 1e-12) -} - -func TestCalculateCost_ImageGeneration_OutputCountFromData(t *testing.T) { - // Output image count derived from len(Data) via populateOutputImageCount - pricing := configstoreTables.TableModelPricing{ - Model: "dall-e-3", - Provider: "openai", - Mode: "image_generation", - OutputCostPerImage: bifrost.Ptr(0.04), - } - - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("dall-e-3", "openai", "image_generation"): pricing, - }) - - resp := &schemas.BifrostResponse{ - ImageGenerationResponse: &schemas.BifrostImageGenerationResponse{ - Data: []schemas.ImageData{ - {URL: "https://example.com/img1.png", Index: 0}, - {URL: "https://example.com/img2.png", Index: 1}, - {URL: "https://example.com/img3.png", Index: 2}, - }, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ImageGenerationRequest, - RoutingInfo: routingInfoFor("openai", "dall-e-3"), - }, - }, - } - cost := mc.CalculateCost(resp, nil) - // 3 output images * $0.04 = $0.12 - assert.InDelta(t, 0.12, cost, 1e-12) -} - -func TestCalculateCost_ImageGeneration_NilUsage_NoPerImagePricing(t *testing.T) { - // No per-image pricing and no tokens — should return 0 - pricing := configstoreTables.TableModelPricing{ - Model: "token-only-model", - Provider: "test", - Mode: "image_generation", - InputCostPerToken: bifrost.Ptr(0.000001), - OutputCostPerToken: bifrost.Ptr(0.000002), - } - - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("token-only-model", "test", "image_generation"): pricing, - }) - - resp := makeImageResponse("test", "token-only-model", nil) - cost := mc.CalculateCost(resp, nil) - // No per-image pricing and all tokens are zero → 0 - assert.InDelta(t, 0.0, cost, 1e-12) -} - -func TestCalculateCost_ImageGeneration_EmptyUsage_PerImagePricing(t *testing.T) { - // Usage exists but all fields are zero — same as nil usage, should use per-image pricing - pricing := configstoreTables.TableModelPricing{ - Model: "dall-e-3", - Provider: "openai", - Mode: "image_generation", - OutputCostPerImage: bifrost.Ptr(0.04), - } - - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("dall-e-3", "openai", "image_generation"): pricing, - }) - - resp := makeImageResponse("openai", "dall-e-3", &schemas.ImageUsage{}) - cost := mc.CalculateCost(resp, nil) - assert.InDelta(t, 0.04, cost, 1e-12) -} - -func TestComputeImageCost_MixedInputTokensOutputPerImage(t *testing.T) { - // Input has tokens (text prompt), output has no tokens but per-image pricing - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000005), - OutputCostPerToken: bifrost.Ptr(0.000015), - OutputCostPerImage: bifrost.Ptr(0.04), - } - usage := &schemas.ImageUsage{ - InputTokens: 500, - OutputTokensDetails: &schemas.ImageTokenDetails{NImages: 2}, - } - cost := computeImageCost(&p, usage, "", "", serviceTier{}) - // Input: 500 tokens * $0.000005 = $0.0025 - // Output: no output tokens → falls back to 2 images * $0.04 = $0.08 - assert.InDelta(t, 0.0825, cost, 1e-12) -} - -func TestComputeImageCost_MixedInputPerImageOutputTokens(t *testing.T) { - // Input has no tokens but per-image count, output has tokens - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000005), - OutputCostPerToken: bifrost.Ptr(0.000015), - InputCostPerImage: bifrost.Ptr(0.01), - } - usage := &schemas.ImageUsage{ - NumInputImages: 3, - OutputTokens: 1000, - } - cost := computeImageCost(&p, usage, "", "", serviceTier{}) - // Input: no input tokens → falls back to 3 images * $0.01 = $0.03 - // Output: 1000 tokens * $0.000015 = $0.015 - assert.InDelta(t, 0.045, cost, 1e-12) -} - -func TestComputeImageCost_BothHaveTokens_IgnoresPerImage(t *testing.T) { - // Both sides have tokens — per-image pricing is ignored - p := configstoreTables.TableModelPricing{ - InputCostPerToken: bifrost.Ptr(0.000005), - OutputCostPerToken: bifrost.Ptr(0.000015), - InputCostPerImage: bifrost.Ptr(0.01), - OutputCostPerImage: bifrost.Ptr(0.04), - } - usage := &schemas.ImageUsage{ - InputTokens: 200, - OutputTokens: 800, - TotalTokens: 1000, - NumInputImages: 3, - } - cost := computeImageCost(&p, usage, "", "", serviceTier{}) - // Input: 200 * $0.000005 = $0.001 (tokens present, per-image ignored) - // Output: 800 * $0.000015 = $0.012 (tokens present, per-image ignored) - assert.InDelta(t, 0.013, cost, 1e-12) -} - -func TestCalculateCost_ResponsesWithCodeInterpreter(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4.1", "openai", "chat"): chatPricing(0.000002, 0.000008), - }) - - ciType := schemas.ResponsesMessageTypeCodeInterpreterCall - resp := &schemas.BifrostResponse{ - ResponsesResponse: &schemas.BifrostResponsesResponse{ - Usage: &schemas.ResponsesResponseUsage{ - InputTokens: 579, - OutputTokens: 334, - TotalTokens: 913, - }, - Output: []schemas.ResponsesMessage{ - {Type: &ciType}, - {Type: &ciType}, - }, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ResponsesRequest, - RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4.1"), - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - // Token cost only: 579*0.000002 + 334*0.000008 = 0.001158 + 0.002672 = 0.003830 - // Session cost is now tracked via ContainerCreateRequest, not per-response - assert.InDelta(t, 0.003830, cost, 1e-6) -} - -// --------------------------------------------------------------------------- -// computeContainerCreationCost -// --------------------------------------------------------------------------- - -func TestComputeContainerCreationCost_Basic(t *testing.T) { - p := configstoreTables.TableModelPricing{ - Model: "container", - Provider: "openai", - Mode: "chat", - CodeInterpreterCostPerSession: bifrost.Ptr(0.03), - } - assert.InDelta(t, 0.03, computeContainerCreationCost(&p), 1e-12) -} - -func TestComputeContainerCreationCost_NilPricing(t *testing.T) { - assert.Equal(t, 0.0, computeContainerCreationCost(nil)) -} - -func TestComputeContainerCreationCost_NilRate(t *testing.T) { - p := configstoreTables.TableModelPricing{ - Model: "container", - Provider: "openai", - Mode: "chat", - } - assert.Equal(t, 0.0, computeContainerCreationCost(&p)) -} - -// --------------------------------------------------------------------------- -// ContainerCreateRequest end-to-end via CalculateCost -// --------------------------------------------------------------------------- - -func TestCalculateCost_ContainerCreate_NoMemoryLimit(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("container", "openai", "chat"): { - Model: "container", - Provider: "openai", - Mode: "chat", - CodeInterpreterCostPerSession: bifrost.Ptr(0.03), - }, - }) - - resp := &schemas.BifrostResponse{ - ContainerCreateResponse: &schemas.BifrostContainerCreateResponse{ - ID: "cntr_abc123", - Name: "test-container", - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ContainerCreateRequest, - RoutingInfo: schemas.RoutingInfo{Provider: schemas.OpenAI}, - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - assert.InDelta(t, 0.03, cost, 1e-12) -} - -func TestCalculateCost_ContainerCreate_MemorySpecificEntry(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("container", "openai", "chat"): { - Model: "container", - Provider: "openai", - Mode: "chat", - CodeInterpreterCostPerSession: bifrost.Ptr(0.03), - }, - makeKey("container-4g", "openai", "chat"): { - Model: "container-4g", - Provider: "openai", - Mode: "chat", - CodeInterpreterCostPerSession: bifrost.Ptr(0.12), - }, - }) - - resp := &schemas.BifrostResponse{ - ContainerCreateResponse: &schemas.BifrostContainerCreateResponse{ - ID: "cntr_abc123", - Name: "test-container", - MemoryLimit: "4g", - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ContainerCreateRequest, - RoutingInfo: schemas.RoutingInfo{Provider: schemas.OpenAI}, - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - assert.InDelta(t, 0.12, cost, 1e-12) -} - -func TestCalculateCost_ContainerCreate_FallsBackToBaseEntry(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("container", "openai", "chat"): { - Model: "container", - Provider: "openai", - Mode: "chat", - CodeInterpreterCostPerSession: bifrost.Ptr(0.03), - }, - }) - - resp := &schemas.BifrostResponse{ - ContainerCreateResponse: &schemas.BifrostContainerCreateResponse{ - ID: "cntr_abc123", - Name: "test-container", - MemoryLimit: "4g", - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ContainerCreateRequest, - RoutingInfo: schemas.RoutingInfo{Provider: schemas.OpenAI}, - }, - }, - } - - // No container-4g entry — should fall back to base "container" rate - cost := mc.CalculateCost(resp, nil) - assert.InDelta(t, 0.03, cost, 1e-12) -} - -func TestCalculateCost_ContainerCreate_NoPricingEntry(t *testing.T) { - mc := testCatalogWithPricing(nil) - - resp := &schemas.BifrostResponse{ - ContainerCreateResponse: &schemas.BifrostContainerCreateResponse{ - ID: "cntr_abc123", - Name: "test-container", - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ContainerCreateRequest, - RoutingInfo: schemas.RoutingInfo{Provider: schemas.OpenAI}, - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - assert.Equal(t, 0.0, cost) -} - -// --------------------------------------------------------------------------- -// Backward-compat: RoutingInfo missing → synthesize from deprecated triplet -// -// Covers callers stuck on the legacy ExtraFields shape: -// - LoggerPlugin.RecalculateCosts replaying logs written before RoutingInfo existed -// - Third-party plugins / SDK users that haven't migrated to RoutingInfo -// -// The fallback only fires when RoutingInfo is fully empty (zero Provider, -// zero Model, nil ResolvedKeyAlias). Any partial population is trusted. -// --------------------------------------------------------------------------- - -func TestCalculateCost_BackCompat_LegacyFieldsOnly_NoAlias(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - - // Caller populates only the deprecated triplet — no RoutingInfo. - // Pricing should fall back to Provider + OriginalModelRequested. - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - Usage: &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500}, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - Provider: schemas.OpenAI, - OriginalModelRequested: "gpt-4o", - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - // 1000 * 0.000005 + 500 * 0.000015 = 0.005 + 0.0075 = 0.0125 - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestCalculateCost_BackCompat_LegacyFieldsOnly_WithAlias(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("my-deployment", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - - // Caller populates only the deprecated triplet with a distinct - // ResolvedModelUsed (i.e. an alias was matched at original request time - // and the wire model differs from the caller-facing name). The fallback - // should route ResolvedModelUsed into ResolvedKeyAlias.ModelID so the - // catalog lookup hits the deployment-keyed entry. - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - Usage: &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500}, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - Provider: schemas.OpenAI, - OriginalModelRequested: "my-alias-name", - ResolvedModelUsed: "my-deployment", - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - // 1000 * 0.000005 + 500 * 0.000015 = 0.0125, charged via the deployment-keyed entry - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestCalculateCost_BackCompat_RoutingInfoWinsOverLegacyFields(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - makeKey("gemini-2.0-flash", "gemini", "chat"): { - Model: "gemini-2.0-flash", - Provider: "gemini", - Mode: "chat", - InputCostPerToken: bifrost.Ptr(0.0000001), - OutputCostPerToken: bifrost.Ptr(0.0000004), - }, - }) - - // Both populated. The modern fields (RoutingInfo) must win — the - // fallback only fires when RoutingInfo is fully unset. - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - Usage: &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500}, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - RoutingInfo: routingInfoFor(schemas.OpenAI, "gpt-4o"), - Provider: schemas.Gemini, - OriginalModelRequested: "gemini-2.0-flash", - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - // Priced via RoutingInfo → openai/gpt-4o → 0.0125 (not the gemini rate). - assert.InDelta(t, 0.0125, cost, 1e-12) -} - -func TestCalculateCost_BackCompat_BothEmpty_ReturnsZero(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - - // Neither RoutingInfo nor the deprecated triplet are populated. - // Pricing has no way to identify the model; cost is 0. - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - Usage: &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500}, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - assert.Equal(t, 0.0, cost) -} - -func TestCalculateCost_BackCompat_PartialRoutingInfo_NoFallback(t *testing.T) { - mc := testCatalogWithPricing(map[string]configstoreTables.TableModelPricing{ - makeKey("gpt-4o", "openai", "chat"): chatPricing(0.000005, 0.000015), - }) - - // RoutingInfo has Model but no Provider. The legacy Provider field is - // also set. The fallback MUST NOT fire — partial RoutingInfo means the - // caller intended to use RoutingInfo. With Provider unset on RoutingInfo, - // the catalog lookup fails and cost is 0. (This guards the trigger - // against false positives.) - resp := &schemas.BifrostResponse{ - ChatResponse: &schemas.BifrostChatResponse{ - Usage: &schemas.BifrostLLMUsage{PromptTokens: 1000, CompletionTokens: 500, TotalTokens: 1500}, - ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - RoutingInfo: schemas.RoutingInfo{Model: "gpt-4o"}, - Provider: schemas.OpenAI, - }, - }, - } - - cost := mc.CalculateCost(resp, nil) - assert.Equal(t, 0.0, cost) -} - -// --------------------------------------------------------------------------- -// file:// URL loading tests -// --------------------------------------------------------------------------- - -func TestLoadPricingFromURL_FileScheme(t *testing.T) { - pricingData := map[string]PricingEntry{ - "gpt-4o": { - Provider: "openai", - Mode: "chat", - }, - } - data, err := json.Marshal(pricingData) - require.NoError(t, err) - - f, err := os.CreateTemp(t.TempDir(), "pricing-*.json") - require.NoError(t, err) - _, err = f.Write(data) - require.NoError(t, err) - f.Close() - - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.pricingURL = "file://" + f.Name() - - result, err := mc.loadPricingFromURL(context.Background()) - require.NoError(t, err) - require.Len(t, result, 1) - assert.Equal(t, "openai", result["gpt-4o"].Provider) -} - -func TestLoadPricingFromURL_FileMissing(t *testing.T) { - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.pricingURL = "file:///nonexistent/path/pricing.json" - - _, err := mc.loadPricingFromURL(context.Background()) - require.Error(t, err) -} - -func TestLoadModelParametersFromURL_FileScheme(t *testing.T) { - paramsData := map[string]json.RawMessage{ - "gpt-4o": json.RawMessage(`{"max_output_tokens":4096}`), - } - data, err := json.Marshal(paramsData) - require.NoError(t, err) - - f, err := os.CreateTemp(t.TempDir(), "model-parameters-*.json") - require.NoError(t, err) - _, err = f.Write(data) - require.NoError(t, err) - f.Close() - - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.modelParametersURL = "file://" + f.Name() - - result, err := mc.loadModelParametersFromURL(context.Background()) - require.NoError(t, err) - require.Len(t, result, 1) - assert.JSONEq(t, `{"max_output_tokens":4096}`, string(result["gpt-4o"])) -} - -func TestLoadModelParametersFromURL_FileMissing(t *testing.T) { - mc := newTestCatalog(nil, nil) - mc.logger = noOpLogger{} - mc.modelParametersURL = "file:///nonexistent/path/model-parameters.json" - - _, err := mc.loadModelParametersFromURL(context.Background()) - require.Error(t, err) -} diff --git a/framework/modelcatalog/refine_test.go b/framework/modelcatalog/refine_test.go deleted file mode 100644 index 297a055342f..00000000000 --- a/framework/modelcatalog/refine_test.go +++ /dev/null @@ -1,51 +0,0 @@ -package modelcatalog - -import ( - "testing" - - "github.com/maximhq/bifrost/core/schemas" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// TestRefineModelForProvider_ReplicateRefinesOpenAIModel verifies that -// Replicate can recover nested provider slugs for provider-pinned OpenAI-family models. -func TestRefineModelForProvider_ReplicateRefinesOpenAIModel(t *testing.T) { - mc := newTestCatalog(map[schemas.ModelProvider][]string{ - schemas.Replicate: {"openai/gpt-5-nano"}, - }, map[string]string{ - "openai/gpt-5-nano": "gpt-5-nano", - }) - - refined, err := mc.RefineModelForProvider(schemas.Replicate, "gpt-5-nano") - require.NoError(t, err) - assert.Equal(t, "openai/gpt-5-nano", refined) -} - -// TestRefineModelForProvider_ReplicatePreservesOwnerSlashModel verifies that -// standard Replicate owner/model slugs are not mistaken for nested provider slugs. -func TestRefineModelForProvider_ReplicatePreservesOwnerSlashModel(t *testing.T) { - mc := newTestCatalog(map[schemas.ModelProvider][]string{ - schemas.Replicate: {"meta/meta-llama-3-8b"}, - }, nil) - - refined, err := mc.RefineModelForProvider(schemas.Replicate, "meta/meta-llama-3-8b") - require.NoError(t, err) - assert.Equal(t, "meta/meta-llama-3-8b", refined) -} - -// TestRefineModelForProvider_ReplicateReturnsAmbiguousMatchError verifies that -// refinement fails fast when multiple nested provider slugs match the same base model. -func TestRefineModelForProvider_ReplicateReturnsAmbiguousMatchError(t *testing.T) { - mc := newTestCatalog(map[schemas.ModelProvider][]string{ - schemas.Replicate: { - "openai/gpt-5-nano", - "xai/gpt-5-nano", - }, - }, nil) - - refined, err := mc.RefineModelForProvider(schemas.Replicate, "gpt-5-nano") - require.Error(t, err) - assert.Empty(t, refined) - assert.Contains(t, err.Error(), "multiple compatible models found") -} diff --git a/framework/modelcatalog/shims.go b/framework/modelcatalog/shims.go new file mode 100644 index 00000000000..ef014f2fcc9 --- /dev/null +++ b/framework/modelcatalog/shims.go @@ -0,0 +1,43 @@ +// Compatibility shims preserved to keep server.go and the enterprise +// transport compiling without code changes in the same commit as the +// internal refactor. Each shim mirrors the pre-refactor API by aggregating +// into a single live entry keyed by "" (empty keyID). +// +// The follow-up PR replaces the call sites with per-key fanout via +// BifrostListModelsRequest.KeyID and then DELETES THIS WHOLE FILE. +package modelcatalog + +import ( + "github.com/maximhq/bifrost/core/schemas" +) + +// UpsertModelDataForProvider stores the merged filtered response for the +// provider in a single aggregated live entry. modelsInKeys is retained as a +// fallback when modelData is empty (provider list-models failed or no keys +// configured) — matches the pre-refactor "trust the user-allowed list" path. +// +// Deprecated: shim. Use UpsertLive per key once the call site adopts +// BifrostListModelsRequest.KeyID. +func (mc *ModelCatalog) UpsertModelDataForProvider(provider schemas.ModelProvider, modelData *schemas.BifrostListModelsResponse, modelsInKeys []schemas.Model) { + models := extractModelIDs(modelData, provider) + if len(models) == 0 { + models = extractModelIDs(&schemas.BifrostListModelsResponse{Data: modelsInKeys}, provider) + } + mc.live.Upsert(provider, "", false, models) +} + +// UpsertUnfilteredModelDataForProvider stores the unfiltered provider +// response in a single aggregated entry. +// +// Deprecated: shim. See UpsertModelDataForProvider. +func (mc *ModelCatalog) UpsertUnfilteredModelDataForProvider(provider schemas.ModelProvider, modelData *schemas.BifrostListModelsResponse) { + models := extractModelIDs(modelData, provider) + mc.live.Upsert(provider, "", true, models) +} + +// DeleteModelDataForProvider drops every live entry for the provider. +// +// Deprecated: shim. Use InvalidateLiveProvider. +func (mc *ModelCatalog) DeleteModelDataForProvider(provider schemas.ModelProvider) { + mc.live.InvalidateProvider(provider) +} diff --git a/framework/modelcatalog/sync.go b/framework/modelcatalog/sync.go deleted file mode 100644 index bb7c74e1e82..00000000000 --- a/framework/modelcatalog/sync.go +++ /dev/null @@ -1,545 +0,0 @@ -package modelcatalog - -import ( - "context" - "encoding/json" - "fmt" - "io" - "net/http" - "net/url" - "os" - "slices" - "sync" - "time" - - bifrost "github.com/maximhq/bifrost/core" - providerUtils "github.com/maximhq/bifrost/core/providers/utils" - "github.com/maximhq/bifrost/core/schemas" - configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" - "github.com/tidwall/gjson" - "gorm.io/gorm" -) - -const ( - urlFetchMaxRetries = 3 // retries after the first attempt (4 attempts total) - urlFetchMaxBackoff = 10 * time.Second // cap for exponential backoff (steps start at 1s) -) - -// syncPricing syncs pricing data from URL to database and updates cache -func (mc *ModelCatalog) syncPricing(ctx context.Context) error { - if mc.shouldSyncGate != nil { - if !mc.shouldSyncGate(ctx) { - return nil - } - } - // Load pricing data from URL - pricingData, err := WithRetries(ctx, urlFetchMaxRetries, urlFetchMaxBackoff, func() (map[string]PricingEntry, error) { - return mc.loadPricingFromURL(ctx) - }) - if err != nil { - // Check if we have existing data in database - pricingRecords, pricingErr := mc.configStore.GetModelPrices(ctx) - if pricingErr != nil { - return fmt.Errorf("failed to get pricing records: %w", pricingErr) - } - if len(pricingRecords) > 0 { - mc.logger.Warn("failed to fetch pricing from URL, falling back to existing database records: %v", err) - return nil - } else { - return fmt.Errorf("failed to load pricing data from URL and no existing data in database: %w", err) - } - } - - // Update database in transaction - err = mc.configStore.ExecuteTransaction(ctx, func(tx *gorm.DB) error { - // Deduplicate and insert new pricing data - seen := make(map[string]bool) - for modelKey, entry := range pricingData { - pricing := convertPricingDataToTableModelPricing(modelKey, entry) - // Create composite key for deduplication - key := makeKey(pricing.Model, pricing.Provider, pricing.Mode) - // Skip if already seen - if exists, ok := seen[key]; ok && exists { - continue - } - // Mark as seen - seen[key] = true - if err := mc.configStore.UpsertModelPrices(ctx, &pricing, tx); err != nil { - return fmt.Errorf("failed to create pricing record for model %s: %w", pricing.Model, err) - } - } - - // Clear seen map - seen = nil - - return nil - }) - if err != nil { - return fmt.Errorf("failed to sync pricing data to database: %w", err) - } - - // Reload cache from database - if err := mc.loadPricingFromDatabase(ctx); err != nil { - return fmt.Errorf("failed to reload pricing cache: %w", err) - } - - // Populate model params cache from pricing datasheet max_output_tokens - mc.populateModelParamsFromPricing(pricingData) - - mc.logger.Debug("successfully synced %d pricing records", len(pricingData)) - return nil -} - -// populateModelParamsFromPricing extracts max_output_tokens from pricing entries -// and populates the model params cache so that providers can look up max output -// tokens without a separate model-parameters sync. -func (mc *ModelCatalog) populateModelParamsFromPricing(pricingData map[string]PricingEntry) { - modelParamsEntries := make(map[string]providerUtils.ModelParams) - for modelKey, entry := range pricingData { - if entry.MaxOutputTokens != nil { - modelName := extractModelName(modelKey) - params := providerUtils.ModelParams{ - MaxOutputTokens: entry.MaxOutputTokens, - } - modelParamsEntries[modelName] = params - } - } - if len(modelParamsEntries) > 0 { - providerUtils.BulkSetModelParams(modelParamsEntries) - mc.logger.Debug("populated %d model params entries from pricing datasheet", len(modelParamsEntries)) - } -} - -// loadPricingFromURL loads pricing data from the configured URL (supports file:// and http(s)://) -func (mc *ModelCatalog) loadPricingFromURL(ctx context.Context) (map[string]PricingEntry, error) { - rawURL := mc.getPricingURL() - - parsed, err := url.Parse(rawURL) - if err != nil { - return nil, fmt.Errorf("failed to parse pricing URL: %w", err) - } - - var data []byte - - if parsed.Scheme == "file" { - data, err = os.ReadFile(parsed.Path) - if err != nil { - return nil, fmt.Errorf("failed to read pricing file: %w", err) - } - } else { - if err := bifrost.ValidateExternalURL(rawURL, true); err != nil { - return nil, fmt.Errorf("pricing URL validation failed: %w", err) - } - client := &http.Client{Timeout: DefaultPricingTimeout} - req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil) - if err != nil { - return nil, fmt.Errorf("failed to create HTTP request: %w", err) - } - resp, err := client.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to download pricing data: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("failed to download pricing data: HTTP %d", resp.StatusCode) - } - - data, err = io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read pricing data response: %w", err) - } - } - - var pricingData map[string]PricingEntry - if err := json.Unmarshal(data, &pricingData); err != nil { - return nil, fmt.Errorf("failed to unmarshal pricing data: %w", err) - } - - mc.logger.Debug("successfully loaded and parsed %d pricing records", len(pricingData)) - return pricingData, nil -} - -// loadPricingIntoMemoryFromURL loads pricing data from URL into memory cache (when config store is not available) -func (mc *ModelCatalog) loadPricingIntoMemoryFromURL(ctx context.Context) error { - pricingData, err := WithRetries(ctx, urlFetchMaxRetries, urlFetchMaxBackoff, func() (map[string]PricingEntry, error) { - return mc.loadPricingFromURL(ctx) - }) - if err != nil { - return fmt.Errorf("failed to load pricing data from URL: %w", err) - } - - mc.mu.Lock() - defer mc.mu.Unlock() - - // Clear and rebuild the pricing map - mc.pricingData = make(map[string]configstoreTables.TableModelPricing, len(pricingData)) - for modelKey, entry := range pricingData { - pricing := convertPricingDataToTableModelPricing(modelKey, entry) - key := makeKey(pricing.Model, pricing.Provider, pricing.Mode) - mc.pricingData[key] = pricing - } - - // Populate model params cache from pricing datasheet max_output_tokens - mc.populateModelParamsFromPricing(pricingData) - - return nil -} - -// ReloadPricing re-reads the pricing table into the in-memory cache. The -// management API uses this after a batched write so the new attributes are -// observable immediately. The existing 24-hour sync owns refreshing pricing -// fields from the upstream datasheet; this method just refreshes the cache. -func (mc *ModelCatalog) ReloadPricing(ctx context.Context) error { - return mc.loadPricingFromDatabase(ctx) -} - -// loadPricingFromDatabase loads pricing data from database into memory cache -func (mc *ModelCatalog) loadPricingFromDatabase(ctx context.Context) error { - if mc.configStore == nil { - return nil - } - - pricingRecords, err := mc.configStore.GetModelPrices(ctx) - if err != nil { - return fmt.Errorf("failed to load pricing from database: %w", err) - } - - mc.mu.Lock() - defer mc.mu.Unlock() - - // Clear and rebuild the pricing map - mc.pricingData = make(map[string]configstoreTables.TableModelPricing, len(pricingRecords)) - for _, pricing := range pricingRecords { - key := makeKey(pricing.Model, pricing.Provider, pricing.Mode) - mc.pricingData[key] = pricing - } - - mc.logger.Debug("loaded %d pricing records from database into memory", len(mc.pricingData)) - return nil -} - -// loadModelParametersFromDatabase bulk-loads model parameters from the DB into the provider -// utils cache (startup / ReloadFromDB). The SetCacheMissHandler path still loads one row at -// a time on cache miss; both use the same table JSON shape. -// Returns the number of rows loaded so callers can decide whether to background-sync from URL. -func (mc *ModelCatalog) loadModelParametersFromDatabase(ctx context.Context) (int, error) { - if mc.configStore == nil { - return 0, nil - } - - rows, err := mc.configStore.GetModelParameters(ctx) - if err != nil { - return 0, fmt.Errorf("failed to load model parameters from database: %w", err) - } - if len(rows) == 0 { - mc.logger.Debug("no model parameters rows in database") - return 0, nil - } - - paramsData := make(map[string]json.RawMessage, len(rows)) - for _, row := range rows { - paramsData[row.Model] = json.RawMessage(row.Data) - } - mc.applyModelParameters(paramsData) - mc.logger.Debug("loaded %d model parameters records from database into cache", len(rows)) - return len(rows), nil -} - -// startSyncWorker starts the background sync worker -func (mc *ModelCatalog) startSyncWorker(ctx context.Context) { - // IMPORTANT: scheduling model - // - // The sync worker wakes on a fixed ticker (syncWorkerTickerPeriod). On each - // wake it checks: - // - // time.Since(lastSyncTimestamp) >= pricingSyncInterval - // - // pricingSyncInterval defines the minimum elapsed time between syncs. The - // ticker period is the check granularity and must stay well below the - // minimum supported pricingSyncInterval, otherwise ticker drift (the few - // seconds a sync takes to complete) pushes the next check just under the - // threshold and the effective cadence doubles. - mc.syncTicker = time.NewTicker(syncWorkerTickerPeriod) - mc.wg.Add(1) - go mc.syncWorker(ctx) -} - -// withDistributedLock acquires a named distributed lock and executes fn under it. -// Pass retries=0 to block until acquired (Lock); pass retries>0 to use LockWithRetry. -func (mc *ModelCatalog) withDistributedLock(ctx context.Context, key string, retries int, fn func() error) error { - lock, err := mc.distributedLockManager.NewLock(key) - if err != nil { - return fmt.Errorf("failed to create lock %q: %w", key, err) - } - if retries > 0 { - if err := lock.LockWithRetry(ctx, retries); err != nil { - return fmt.Errorf("failed to acquire lock %q: %w", key, err) - } - } else { - if err := lock.Lock(ctx); err != nil { - return fmt.Errorf("failed to acquire lock %q: %w", key, err) - } - } - // Use a fresh context for unlock so that a cancelled or timed-out work context - // does not prevent the lock row from being deleted. If we reused ctx and it was - // already cancelled when the defer fires, ReleaseLock's DB call would fail - // silently and the lock would stay in the database until TTL expiry (30s), - // blocking every other node from acquiring it during that window. - defer func() { - if err := lock.Unlock(context.Background()); err != nil { - mc.logger.Warn("failed to release distributed lock %q: %v", key, err) - } - }() - return fn() -} - -// syncTick performs a single sync tick with proper lock management -// if the last sync was more than the sync interval ago, sync pricing and model parameters in parallel -func (mc *ModelCatalog) syncTick(ctx context.Context) { - mc.syncMu.RLock() - lastSync := mc.lastSyncedAt - interval := mc.syncInterval - mc.syncMu.RUnlock() - - if time.Since(lastSync) >= interval { - mc.logger.Debug("starting model catalog background sync") - if err := mc.withDistributedLock(ctx, "model_catalog_pricing_sync", 10, func() error { - // Sync pricing and model parameters in parallel - var wg sync.WaitGroup - var pricingErr, paramsErr error - wg.Add(2) - go func() { - defer wg.Done() - if err := mc.syncPricing(ctx); err != nil { - mc.logger.Error("background pricing sync failed: %v", err) - pricingErr = err - } - }() - go func() { - defer wg.Done() - if err := mc.syncModelParameters(ctx); err != nil { - mc.logger.Error("background model parameters sync failed: %v", err) - paramsErr = err - } - }() - wg.Wait() - - if pricingErr == nil && paramsErr == nil { - if mc.afterSyncHook != nil { - mc.afterSyncHook(ctx) - } - mc.syncMu.Lock() - mc.lastSyncedAt = time.Now() - mc.syncMu.Unlock() - } - if pricingErr != nil { - return pricingErr - } - return paramsErr - }); err != nil { - mc.logger.Error("failed to run model catalog sync: %v", err) - } - mc.logger.Debug("model catalog background sync completed") - } -} - -// syncWorker runs the background sync check -func (mc *ModelCatalog) syncWorker(ctx context.Context) { - defer mc.wg.Done() - defer mc.syncTicker.Stop() - - for { - select { - case <-ctx.Done(): - return - case <-mc.syncTicker.C: - mc.syncTick(ctx) - case <-mc.done: - return - } - } -} - -// --- Model Parameters sync --- - -func (mc *ModelCatalog) applyModelParameters(paramsData map[string]json.RawMessage) { - modelParamsEntries := make(map[string]providerUtils.ModelParams, len(paramsData)) - newResponseTypes := make(map[string][]string, len(paramsData)) - newParamsIndex := make(map[string][]string, len(paramsData)) - - for model, rawData := range paramsData { - var parsed modelParametersParseResult - if err := json.Unmarshal(rawData, &parsed); err != nil { - mc.logger.Warn("model-parameters-sync: skipping malformed parameters for model %s: %v", model, err) - continue - } - - outputs := make([]string, 0, len(parsed.SupportedEndpoints)) - for _, endpoint := range parsed.SupportedEndpoints { - if normalized := normalizeEndpointToOutputType(endpoint); normalized != "" && !slices.Contains(outputs, normalized) { - outputs = append(outputs, normalized) - } - } - - if parsed.Mode != nil { - if normalized := normalizeModeToOutputType(*parsed.Mode); normalized != "" && !slices.Contains(outputs, normalized) { - outputs = append(outputs, normalized) - } - } - - if !slices.Contains(outputs, "text_completion") { - provider := gjson.GetBytes(rawData, "provider") - if provider.Exists() { - key := makeKey(model, normalizeProvider(provider.String()), normalizeRequestType(schemas.TextCompletionRequest)) - - mc.mu.RLock() - _, ok := mc.pricingData[key] - mc.mu.RUnlock() - if ok { - outputs = append(outputs, "text_completion") - } - } - } - - if len(outputs) > 0 { - newResponseTypes[model] = outputs - } - - supported := extractSupportedParams(&parsed) - if len(supported) > 0 { - newParamsIndex[model] = supported - } - - var p struct { - MaxOutputTokens *int `json:"max_output_tokens"` - } - if err := json.Unmarshal(rawData, &p); err == nil && (p.MaxOutputTokens != nil || parsed.VertexMultiRegionOnly != nil) { - modelParamsEntries[model] = providerUtils.ModelParams{ - MaxOutputTokens: p.MaxOutputTokens, - IsVertexMultiRegionOnly: parsed.VertexMultiRegionOnly, - } - } - } - - mc.mu.Lock() - mc.supportedResponseTypes = newResponseTypes - mc.supportedParams = newParamsIndex - mc.mu.Unlock() - - if len(modelParamsEntries) > 0 { - providerUtils.BulkSetModelParams(modelParamsEntries) - } -} - -// loadModelParametersIntoMemoryFromURL loads model parameters from the remote URL into the -// provider utils cache (when config store is not available). -func (mc *ModelCatalog) loadModelParametersIntoMemoryFromURL(ctx context.Context) error { - paramsData, err := WithRetries(ctx, urlFetchMaxRetries, urlFetchMaxBackoff, func() (map[string]json.RawMessage, error) { - return mc.loadModelParametersFromURL(ctx) - }) - if err != nil { - return fmt.Errorf("failed to load model parameters from URL: %w", err) - } - mc.applyModelParameters(paramsData) - return nil -} - -// syncModelParameters syncs model parameters data from URL into memory cache -func (mc *ModelCatalog) syncModelParameters(ctx context.Context) error { - if mc.shouldSyncGate != nil { - if !mc.shouldSyncGate(ctx) { - mc.logger.Debug("model parameters sync cancelled by custom gate") - return nil - } - } - mc.logger.Debug("starting model parameters synchronization") - - paramsData, err := WithRetries(ctx, urlFetchMaxRetries, urlFetchMaxBackoff, func() (map[string]json.RawMessage, error) { - return mc.loadModelParametersFromURL(ctx) - }) - if err != nil { - if mc.configStore != nil { - rows, dbErr := mc.configStore.GetModelParameters(ctx) - if dbErr == nil && len(rows) > 0 { - mc.logger.Error("failed to load model parameters from URL, falling back to existing database records: %v", err) - return nil - } - } - return fmt.Errorf("failed to load model parameters from URL and no existing data in database: %w", err) - } - - // Persist to database if config store is available - if mc.configStore != nil { - err = mc.configStore.ExecuteTransaction(ctx, func(tx *gorm.DB) error { - for model, data := range paramsData { - params := &configstoreTables.TableModelParameters{ - Model: model, - Data: string(data), - } - if err := mc.configStore.UpsertModelParameters(ctx, params, tx); err != nil { - return fmt.Errorf("failed to upsert model parameters for model %s: %w", model, err) - } - } - return nil - }) - if err != nil { - return fmt.Errorf("failed to sync model parameters to database: %w", err) - } - } - - mc.applyModelParameters(paramsData) - - mc.logger.Info("successfully synced %d model parameters records", len(paramsData)) - return nil -} - -// loadModelParametersFromURL loads model parameters data from the configured URL (supports file:// and http(s)://) -func (mc *ModelCatalog) loadModelParametersFromURL(ctx context.Context) (map[string]json.RawMessage, error) { - rawURL := mc.getModelParametersURL() - - parsed, err := url.Parse(rawURL) - if err != nil { - return nil, fmt.Errorf("failed to parse model parameters URL: %w", err) - } - - var data []byte - - if parsed.Scheme == "file" { - data, err = os.ReadFile(parsed.Path) - if err != nil { - return nil, fmt.Errorf("failed to read model parameters file: %w", err) - } - } else { - if err := bifrost.ValidateExternalURL(rawURL, true); err != nil { - return nil, fmt.Errorf("model parameters URL validation failed: %w", err) - } - client := &http.Client{Timeout: DefaultModelParametersTimeout} - req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil) - if err != nil { - return nil, fmt.Errorf("failed to create HTTP request: %w", err) - } - resp, err := client.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to download model parameters data: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("failed to download model parameters data: HTTP %d", resp.StatusCode) - } - - data, err = io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("failed to read model parameters response: %w", err) - } - } - - var paramsData map[string]json.RawMessage - if err := json.Unmarshal(data, ¶msData); err != nil { - return nil, fmt.Errorf("failed to unmarshal model parameters data: %w", err) - } - - mc.logger.Debug("successfully loaded and parsed %d model parameters records", len(paramsData)) - return paramsData, nil -} diff --git a/framework/modelcatalog/utils.go b/framework/modelcatalog/utils.go deleted file mode 100644 index b26cfa15cda..00000000000 --- a/framework/modelcatalog/utils.go +++ /dev/null @@ -1,458 +0,0 @@ -package modelcatalog - -import ( - "context" - "slices" - "strings" - "time" - - "github.com/bytedance/sonic" - "github.com/maximhq/bifrost/core/schemas" - configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" -) - -const retryBackoffMin = time.Second - -// WithRetries runs op until it succeeds or maxRetries retries are exhausted -// (1 initial attempt + maxRetries retries). After each failure it waits with -// exponential backoff starting at 1 second (retryBackoffMin), capped at maxBackoff -// when maxBackoff > 0. If maxBackoff is zero, there is no upper cap on the delay. -func WithRetries[T any](ctx context.Context, maxRetries int, maxBackoff time.Duration, op func() (T, error)) (T, error) { - var zero T - if maxRetries < 0 { - maxRetries = 0 - } - var lastErr error - for attempt := 0; attempt <= maxRetries; attempt++ { - select { - case <-ctx.Done(): - return zero, ctx.Err() - default: - } - - if attempt > 0 { - backoff := retryBackoffMin * time.Duration(1< 0 && backoff > maxBackoff { - backoff = maxBackoff - } - select { - case <-ctx.Done(): - return zero, ctx.Err() - case <-time.After(backoff): - } - } - v, err := op() - if err == nil { - return v, nil - } - lastErr = err - } - return zero, lastErr -} - -// makeKey creates a unique key for a model, provider, and mode for pricingData map -func makeKey(model, provider, mode string) string { return model + "|" + provider + "|" + mode } - -// normalizeProvider normalizes the provider name to a consistent format -func normalizeProvider(p string) string { - if strings.Contains(p, "vertex_ai") || p == "google-vertex" { - return string(schemas.Vertex) - } else if strings.Contains(p, "bedrock") { - return string(schemas.Bedrock) - } else if strings.Contains(p, "cohere") { - return string(schemas.Cohere) - } else if strings.Contains(p, "runwayml") { - return string(schemas.Runway) - } else if strings.Contains(p, "fireworks_ai") { - return string(schemas.Fireworks) - } else { - return p - } -} - -// normalizeRequestType normalizes the request type to a consistent format -func normalizeRequestType(reqType schemas.RequestType) string { - baseType := "unknown" - - switch reqType { - case schemas.TextCompletionRequest, schemas.TextCompletionStreamRequest: - baseType = "completion" - case schemas.ChatCompletionRequest, schemas.ChatCompletionStreamRequest: - baseType = "chat" - case schemas.ResponsesRequest, schemas.ResponsesStreamRequest, schemas.WebSocketResponsesRequest, schemas.RealtimeRequest, schemas.CompactionRequest: - baseType = "responses" - case schemas.EmbeddingRequest: - baseType = "embedding" - case schemas.RerankRequest: - baseType = "rerank" - case schemas.SpeechRequest, schemas.SpeechStreamRequest: - baseType = "audio_speech" - case schemas.TranscriptionRequest, schemas.TranscriptionStreamRequest: - baseType = "audio_transcription" - case schemas.ImageGenerationRequest, schemas.ImageGenerationStreamRequest, schemas.ImageVariationRequest: - baseType = "image_generation" - case schemas.ImageEditRequest, schemas.ImageEditStreamRequest: - baseType = "image_edit" - case schemas.VideoGenerationRequest, schemas.VideoRemixRequest: - baseType = "video_generation" - case schemas.OCRRequest: - baseType = "ocr" - case schemas.ContainerCreateRequest: - baseType = "container_create" - } - - return baseType -} - -// normalizeStreamRequestType normalizes the stream request type to a consistent format -// It returns the base request type for the stream request type. -func normalizeStreamRequestType(rt schemas.RequestType) schemas.RequestType { - switch rt { - case schemas.TextCompletionStreamRequest: - return schemas.TextCompletionRequest - case schemas.ChatCompletionStreamRequest: - return schemas.ChatCompletionRequest - case schemas.ResponsesStreamRequest, schemas.WebSocketResponsesRequest: - return schemas.ResponsesRequest - case schemas.RealtimeRequest: - return schemas.RealtimeRequest - case schemas.SpeechStreamRequest: - return schemas.SpeechRequest - case schemas.TranscriptionStreamRequest: - return schemas.TranscriptionRequest - case schemas.ImageGenerationStreamRequest: - return schemas.ImageGenerationRequest - case schemas.ImageEditStreamRequest: - return schemas.ImageEditRequest - default: - return rt - } -} - -// extractModelName extracts the model name from a model key that may be in provider/model format -func extractModelName(modelKey string) string { - if strings.Contains(modelKey, "/") { - parts := strings.Split(modelKey, "/") - if len(parts) > 1 { - return strings.Join(parts[1:], "/") - } - } - return modelKey -} - -// convertPricingDataToTableModelPricing converts the pricing data to a TableModelPricing struct -func convertPricingDataToTableModelPricing(modelKey string, entry PricingEntry) configstoreTables.TableModelPricing { - provider := normalizeProvider(entry.Provider) - modelName := extractModelName(modelKey) - - return configstoreTables.TableModelPricing{ - Model: modelName, - BaseModel: entry.BaseModel, - Provider: provider, - Mode: entry.Mode, - ContextLength: entry.ContextLength, - MaxInputTokens: entry.MaxInputTokens, - MaxOutputTokens: entry.MaxOutputTokens, - Architecture: entry.Architecture, - - // Costs - Text - InputCostPerToken: entry.InputCostPerToken, - OutputCostPerToken: entry.OutputCostPerToken, - InputCostPerTokenBatches: entry.InputCostPerTokenBatches, - OutputCostPerTokenBatches: entry.OutputCostPerTokenBatches, - InputCostPerTokenPriority: entry.InputCostPerTokenPriority, - OutputCostPerTokenPriority: entry.OutputCostPerTokenPriority, - InputCostPerTokenFlex: entry.InputCostPerTokenFlex, - OutputCostPerTokenFlex: entry.OutputCostPerTokenFlex, - InputCostPerTokenAbove200kTokens: entry.InputCostPerTokenAbove200kTokens, - InputCostPerTokenAbove200kTokensPriority: entry.InputCostPerTokenAbove200kTokensPriority, - OutputCostPerTokenAbove200kTokens: entry.OutputCostPerTokenAbove200kTokens, - OutputCostPerTokenAbove200kTokensPriority: entry.OutputCostPerTokenAbove200kTokensPriority, - // Costs - 272k Tier - InputCostPerTokenAbove272kTokens: entry.InputCostPerTokenAbove272kTokens, - InputCostPerTokenAbove272kTokensPriority: entry.InputCostPerTokenAbove272kTokensPriority, - OutputCostPerTokenAbove272kTokens: entry.OutputCostPerTokenAbove272kTokens, - OutputCostPerTokenAbove272kTokensPriority: entry.OutputCostPerTokenAbove272kTokensPriority, - // Costs - Character - InputCostPerCharacter: entry.InputCostPerCharacter, - // Costs - 128k Tier - InputCostPerTokenAbove128kTokens: entry.InputCostPerTokenAbove128kTokens, - InputCostPerImageAbove128kTokens: entry.InputCostPerImageAbove128kTokens, - InputCostPerVideoPerSecondAbove128kTokens: entry.InputCostPerVideoPerSecondAbove128kTokens, - InputCostPerAudioPerSecondAbove128kTokens: entry.InputCostPerAudioPerSecondAbove128kTokens, - OutputCostPerTokenAbove128kTokens: entry.OutputCostPerTokenAbove128kTokens, - - // Costs - Cache - CacheCreationInputTokenCost: entry.CacheCreationInputTokenCost, - CacheReadInputTokenCost: entry.CacheReadInputTokenCost, - CacheCreationInputTokenCostAbove200kTokens: entry.CacheCreationInputTokenCostAbove200kTokens, - CacheReadInputTokenCostAbove200kTokens: entry.CacheReadInputTokenCostAbove200kTokens, - CacheReadInputTokenCostAbove200kTokensPriority: entry.CacheReadInputTokenCostAbove200kTokensPriority, - CacheCreationInputTokenCostAbove1hr: entry.CacheCreationInputTokenCostAbove1hr, - CacheCreationInputTokenCostAbove1hrAbove200kTokens: entry.CacheCreationInputTokenCostAbove1hrAbove200kTokens, - CacheCreationInputAudioTokenCost: entry.CacheCreationInputAudioTokenCost, - CacheReadInputTokenCostPriority: entry.CacheReadInputTokenCostPriority, - CacheReadInputTokenCostFlex: entry.CacheReadInputTokenCostFlex, - CacheReadInputImageTokenCost: entry.CacheReadInputImageTokenCost, - CacheReadInputTokenCostAbove272kTokens: entry.CacheReadInputTokenCostAbove272kTokens, - CacheReadInputTokenCostAbove272kTokensPriority: entry.CacheReadInputTokenCostAbove272kTokensPriority, - - // Costs - Image - InputCostPerImage: entry.InputCostPerImage, - InputCostPerPixel: entry.InputCostPerPixel, - OutputCostPerImage: entry.OutputCostPerImage, - OutputCostPerPixel: entry.OutputCostPerPixel, - OutputCostPerImagePremiumImage: entry.OutputCostPerImagePremiumImage, - OutputCostPerImageAbove512x512Pixels: entry.OutputCostPerImageAbove512x512Pixels, - OutputCostPerImageAbove512x512PixelsPremium: entry.OutputCostPerImageAbove512x512PixelsPremium, - OutputCostPerImageAbove1024x1024Pixels: entry.OutputCostPerImageAbove1024x1024Pixels, - OutputCostPerImageAbove1024x1024PixelsPremium: entry.OutputCostPerImageAbove1024x1024PixelsPremium, - OutputCostPerImageAbove2048x2048Pixels: entry.OutputCostPerImageAbove2048x2048Pixels, - OutputCostPerImageAbove4096x4096Pixels: entry.OutputCostPerImageAbove4096x4096Pixels, - OutputCostPerImageLowQuality: entry.OutputCostPerImageLowQuality, - OutputCostPerImageMediumQuality: entry.OutputCostPerImageMediumQuality, - OutputCostPerImageHighQuality: entry.OutputCostPerImageHighQuality, - OutputCostPerImageAutoQuality: entry.OutputCostPerImageAutoQuality, - // Costs - Image Token - InputCostPerImageToken: entry.InputCostPerImageToken, - OutputCostPerImageToken: entry.OutputCostPerImageToken, - - // Costs - Audio/Video - InputCostPerAudioToken: entry.InputCostPerAudioToken, - InputCostPerAudioPerSecond: entry.InputCostPerAudioPerSecond, - InputCostPerSecond: entry.InputCostPerSecond, - InputCostPerVideoPerSecond: entry.InputCostPerVideoPerSecond, - OutputCostPerAudioToken: entry.OutputCostPerAudioToken, - OutputCostPerVideoPerSecond: entry.OutputCostPerVideoPerSecond, - OutputCostPerSecond: entry.OutputCostPerSecond, - - // Costs - Other - SearchContextCostPerQuery: entry.SearchContextCostPerQuery, - CodeInterpreterCostPerSession: entry.CodeInterpreterCostPerSession, - - // Costs - OCR - OCRCostPerPage: entry.OCRCostPerPage, - AnnotationCostPerPage: entry.AnnotationCostPerPage, - } -} - -// convertTableModelPricingToPricingData converts the TableModelPricing struct to a PricingEntry struct -func convertTableModelPricingToPricingData(pricing *configstoreTables.TableModelPricing) *PricingEntry { - options := PricingOptions{ - // Costs - Text - InputCostPerToken: pricing.InputCostPerToken, - OutputCostPerToken: pricing.OutputCostPerToken, - InputCostPerTokenBatches: pricing.InputCostPerTokenBatches, - OutputCostPerTokenBatches: pricing.OutputCostPerTokenBatches, - InputCostPerTokenPriority: pricing.InputCostPerTokenPriority, - OutputCostPerTokenPriority: pricing.OutputCostPerTokenPriority, - InputCostPerTokenFlex: pricing.InputCostPerTokenFlex, - OutputCostPerTokenFlex: pricing.OutputCostPerTokenFlex, - InputCostPerTokenAbove200kTokens: pricing.InputCostPerTokenAbove200kTokens, - InputCostPerTokenAbove200kTokensPriority: pricing.InputCostPerTokenAbove200kTokensPriority, - OutputCostPerTokenAbove200kTokens: pricing.OutputCostPerTokenAbove200kTokens, - OutputCostPerTokenAbove200kTokensPriority: pricing.OutputCostPerTokenAbove200kTokensPriority, - // Costs - 272k Tier - InputCostPerTokenAbove272kTokens: pricing.InputCostPerTokenAbove272kTokens, - InputCostPerTokenAbove272kTokensPriority: pricing.InputCostPerTokenAbove272kTokensPriority, - OutputCostPerTokenAbove272kTokens: pricing.OutputCostPerTokenAbove272kTokens, - OutputCostPerTokenAbove272kTokensPriority: pricing.OutputCostPerTokenAbove272kTokensPriority, - // Costs - Character - InputCostPerCharacter: pricing.InputCostPerCharacter, - // Costs - 128k Tier - InputCostPerTokenAbove128kTokens: pricing.InputCostPerTokenAbove128kTokens, - InputCostPerImageAbove128kTokens: pricing.InputCostPerImageAbove128kTokens, - InputCostPerVideoPerSecondAbove128kTokens: pricing.InputCostPerVideoPerSecondAbove128kTokens, - InputCostPerAudioPerSecondAbove128kTokens: pricing.InputCostPerAudioPerSecondAbove128kTokens, - OutputCostPerTokenAbove128kTokens: pricing.OutputCostPerTokenAbove128kTokens, - - // Costs - Cache - CacheCreationInputTokenCost: pricing.CacheCreationInputTokenCost, - CacheReadInputTokenCost: pricing.CacheReadInputTokenCost, - CacheCreationInputTokenCostAbove200kTokens: pricing.CacheCreationInputTokenCostAbove200kTokens, - CacheReadInputTokenCostAbove200kTokens: pricing.CacheReadInputTokenCostAbove200kTokens, - CacheReadInputTokenCostAbove200kTokensPriority: pricing.CacheReadInputTokenCostAbove200kTokensPriority, - CacheCreationInputTokenCostAbove1hr: pricing.CacheCreationInputTokenCostAbove1hr, - CacheCreationInputTokenCostAbove1hrAbove200kTokens: pricing.CacheCreationInputTokenCostAbove1hrAbove200kTokens, - CacheCreationInputAudioTokenCost: pricing.CacheCreationInputAudioTokenCost, - CacheReadInputTokenCostPriority: pricing.CacheReadInputTokenCostPriority, - CacheReadInputTokenCostFlex: pricing.CacheReadInputTokenCostFlex, - CacheReadInputImageTokenCost: pricing.CacheReadInputImageTokenCost, - CacheReadInputTokenCostAbove272kTokens: pricing.CacheReadInputTokenCostAbove272kTokens, - CacheReadInputTokenCostAbove272kTokensPriority: pricing.CacheReadInputTokenCostAbove272kTokensPriority, - - // Costs - Image - InputCostPerImage: pricing.InputCostPerImage, - InputCostPerPixel: pricing.InputCostPerPixel, - OutputCostPerImage: pricing.OutputCostPerImage, - OutputCostPerPixel: pricing.OutputCostPerPixel, - OutputCostPerImagePremiumImage: pricing.OutputCostPerImagePremiumImage, - OutputCostPerImageAbove512x512Pixels: pricing.OutputCostPerImageAbove512x512Pixels, - OutputCostPerImageAbove512x512PixelsPremium: pricing.OutputCostPerImageAbove512x512PixelsPremium, - OutputCostPerImageAbove1024x1024Pixels: pricing.OutputCostPerImageAbove1024x1024Pixels, - OutputCostPerImageAbove1024x1024PixelsPremium: pricing.OutputCostPerImageAbove1024x1024PixelsPremium, - OutputCostPerImageAbove2048x2048Pixels: pricing.OutputCostPerImageAbove2048x2048Pixels, - OutputCostPerImageAbove4096x4096Pixels: pricing.OutputCostPerImageAbove4096x4096Pixels, - OutputCostPerImageLowQuality: pricing.OutputCostPerImageLowQuality, - OutputCostPerImageMediumQuality: pricing.OutputCostPerImageMediumQuality, - OutputCostPerImageHighQuality: pricing.OutputCostPerImageHighQuality, - OutputCostPerImageAutoQuality: pricing.OutputCostPerImageAutoQuality, - // Costs - Image Token - InputCostPerImageToken: pricing.InputCostPerImageToken, - OutputCostPerImageToken: pricing.OutputCostPerImageToken, - - // Costs - Audio/Video - InputCostPerAudioToken: pricing.InputCostPerAudioToken, - InputCostPerAudioPerSecond: pricing.InputCostPerAudioPerSecond, - InputCostPerSecond: pricing.InputCostPerSecond, - InputCostPerVideoPerSecond: pricing.InputCostPerVideoPerSecond, - OutputCostPerAudioToken: pricing.OutputCostPerAudioToken, - OutputCostPerVideoPerSecond: pricing.OutputCostPerVideoPerSecond, - OutputCostPerSecond: pricing.OutputCostPerSecond, - - // Costs - Other - SearchContextCostPerQuery: pricing.SearchContextCostPerQuery, - CodeInterpreterCostPerSession: pricing.CodeInterpreterCostPerSession, - - // Costs - OCR - OCRCostPerPage: pricing.OCRCostPerPage, - AnnotationCostPerPage: pricing.AnnotationCostPerPage, - } - return &PricingEntry{ - BaseModel: pricing.BaseModel, - Provider: pricing.Provider, - Mode: pricing.Mode, - ContextLength: pricing.ContextLength, - MaxInputTokens: pricing.MaxInputTokens, - MaxOutputTokens: pricing.MaxOutputTokens, - Architecture: pricing.Architecture, - AdditionalAttributes: pricing.AdditionalAttributes, - PricingOptions: options, - } -} - -// convertTablePricingOverrideToPricingOverride converts a TablePricingOverride to a PricingOverride. -func convertTablePricingOverrideToPricingOverride(override *configstoreTables.TablePricingOverride) (PricingOverride, error) { - var options PricingOptions - if err := sonic.Unmarshal([]byte(override.PricingPatchJSON), &options); err != nil { - return PricingOverride{}, err - } - return PricingOverride{ - ID: override.ID, - Name: override.Name, - ScopeKind: ScopeKind(override.ScopeKind), - VirtualKeyID: override.VirtualKeyID, - ProviderID: override.ProviderID, - ProviderKeyID: override.ProviderKeyID, - MatchType: MatchType(override.MatchType), - Pattern: override.Pattern, - RequestTypes: override.RequestTypes, - Options: options, - }, nil -} - -// normalizeEndpointToOutputType converts a supported_endpoints URL path to a normalized output type. -// Returns empty string for unrecognized endpoints. -func normalizeEndpointToOutputType(endpoint string) string { - switch { - case strings.Contains(endpoint, "/chat/completions"): - return "chat_completion" - case strings.Contains(endpoint, "/responses"): - return "responses" - case strings.Contains(endpoint, "/completions"): - return "text_completion" - default: - return "" - } -} - -// normalizeModeToOutputType converts mode to a normalized output type. -func normalizeModeToOutputType(mode string) string { - switch mode { - case "chat": - return "chat_completion" - case "completion": - return "text_completion" - case "responses": - return "responses" - default: - return "" - } -} - -// modelParametersParseResult is the parsed result type used by buildSupportedOutputsIndex. -type modelParametersParseResult struct { - Mode *string `json:"mode,omitempty"` - SupportedEndpoints []string `json:"supported_endpoints,omitempty"` - ModelParameters []struct { - ID string `json:"id"` - } `json:"model_parameters,omitempty"` - SupportsAssistantPrefill *bool `json:"supports_assistant_prefill,omitempty"` - SupportsFunctionCalling *bool `json:"supports_function_calling,omitempty"` - SupportsParallelFunctionCalling *bool `json:"supports_parallel_function_calling,omitempty"` - SupportsToolChoice *bool `json:"supports_tool_choice,omitempty"` - SupportsReasoning *bool `json:"supports_reasoning,omitempty"` - SupportsResponseSchema *bool `json:"supports_response_schema,omitempty"` - SupportsServiceTier *bool `json:"supports_service_tier,omitempty"` - SupportsPromptCaching *bool `json:"supports_prompt_caching,omitempty"` - VertexMultiRegionOnly *bool `json:"vertex_multi_region_only,omitempty"` -} - -// extractSupportedParams builds a list of supported OpenAI-compatible parameter -// names from model_parameters[].id values and supports_* boolean flags. -func extractSupportedParams(parsed *modelParametersParseResult) []string { - var supported []string - addParam := func(name string) { - if !slices.Contains(supported, name) { - supported = append(supported, name) - } - } - - // From model_parameters[].id — map IDs to request param names - for _, mp := range parsed.ModelParameters { - switch mp.ID { - case "reasoning_effort", "reasoning_summary": - addParam("reasoning") - case "web_search": - addParam("web_search_options") - case "promptTools", "image_detail", "stream": - // skip — not top-level request parameters - default: - addParam(mp.ID) - } - } - - // From supports_* boolean flags - if parsed.SupportsAssistantPrefill != nil && *parsed.SupportsAssistantPrefill { - // not an actual model parameter; if present, trailing assistant messages - // for anthropic and bedrock's anthropic models will not be trimmed - addParam("assistant_prefill") - } - if parsed.SupportsFunctionCalling != nil && *parsed.SupportsFunctionCalling { - addParam("tools") - } - if parsed.SupportsParallelFunctionCalling != nil && *parsed.SupportsParallelFunctionCalling { - addParam("parallel_tool_calls") - } - if parsed.SupportsToolChoice != nil && *parsed.SupportsToolChoice { - addParam("tool_choice") - } - if parsed.SupportsReasoning != nil && *parsed.SupportsReasoning { - addParam("reasoning") - } - if parsed.SupportsResponseSchema != nil && *parsed.SupportsResponseSchema { - addParam("response_format") - addParam("text") - } - if parsed.SupportsServiceTier != nil && *parsed.SupportsServiceTier { - addParam("service_tier") - } - if parsed.SupportsPromptCaching != nil && *parsed.SupportsPromptCaching { - addParam("cachePoint") - addParam("cache_control") - addParam("prompt_cache_key") - addParam("prompt_cache_retention") - } - - return supported -}