From 7fb8358942a514025519f84a1106d5aba02e7c72 Mon Sep 17 00:00:00 2001 From: vrivri Date: Thu, 30 Jul 2026 17:10:51 +0530 Subject: [PATCH] smenatic router: budgeting+telemetry management --- core/schemas/bifrost.go | 19 + framework/modelcatalog/datasheet/cost.go | 62 ++- framework/modelcatalog/datasheet/cost_test.go | 150 ++++++++ framework/modelcatalog/pricing.go | 9 + plugins/governance/embedding.go | 196 +++++++++- plugins/governance/embedding_test.go | 363 +++++++++++++++++- plugins/governance/main.go | 10 + plugins/telemetry/main.go | 71 +++- plugins/telemetry/main_test.go | 110 ++++++ transports/bifrost-http/server/server.go | 23 ++ 10 files changed, 982 insertions(+), 31 deletions(-) diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 8492531d01f..04b78b2e881 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -1728,6 +1728,7 @@ type BifrostResponseExtraFields struct { RawRequest interface{} `json:"raw_request,omitempty"` RawResponse interface{} `json:"raw_response,omitempty"` CacheDebug *BifrostCacheDebug `json:"cache_debug,omitempty"` + RoutingDebug *BifrostRoutingDebug `json:"routing_debug,omitempty"` GuardrailDebug *BifrostGuardrailDebug `json:"guardrail_debug,omitempty"` ParseErrors []BatchError `json:"parse_errors,omitempty"` // errors encountered while parsing JSONL batch results ConvertedRequestType RequestType `json:"converted_request_type,omitempty"` @@ -1805,6 +1806,24 @@ type BifrostCacheDebug struct { CacheHitLatency *int64 `json:"cache_hit_latency,omitempty"` } +// BifrostRoutingDebug records routing-classification overhead attached to the +// triggering request — today the embedding call semantic complexity routing +// makes before provider selection. It is stamped whenever a routing embedding +// ran, independent of budget attribution, so routing cost stays observable. +// Routing-mechanism fields (tier, mechanism, similarity) may be added here as +// classification logging grows. +type BifrostRoutingDebug struct { + ProviderUsed *string `json:"provider_used,omitempty"` + ModelUsed *string `json:"model_used,omitempty"` + InputTokens *int `json:"input_tokens,omitempty"` + + // CountTowardBudgets carries the governance count_toward_budgets flag to + // cost calculation, which cannot see governance config. When true, the + // routing embedding cost is added to the request's calculated cost (and so + // to its budget attribution); it is never budget-enforced. + CountTowardBudgets bool `json:"count_toward_budgets,omitempty"` +} + const ( RequestCancelled = "request_cancelled" RequestTimedOut = "request_timed_out" diff --git a/framework/modelcatalog/datasheet/cost.go b/framework/modelcatalog/datasheet/cost.go index eccff2ac8e8..f4f2cecc211 100644 --- a/framework/modelcatalog/datasheet/cost.go +++ b/framework/modelcatalog/datasheet/cost.go @@ -11,7 +11,7 @@ import ( ) // CalculateCost calculates the cost of a Bifrost response. -// It handles all request types, cache and guardrail billing, and tiered pricing. +// It handles all request types, cache, guardrail, and routing 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 { @@ -26,20 +26,60 @@ func (s *Store) CalculateCost(result *schemas.BifrostResponse, scopes *LookupSco extraFields := result.GetExtraFields() - // Handle semantic cache billing - cacheDebug := extraFields.CacheDebug - var requestCost float64 - if cacheDebug != nil { - requestCost = s.calculateCostWithCache(result, cacheDebug, lookupScopes) + // The main request and each internal sidecar are independently billable. + // Keep a single accumulator so cache, guardrail, and routing debug metadata + // can coexist without one branch hiding another. + var cost float64 + if extraFields != nil && extraFields.CacheDebug != nil { + cost = s.calculateCostWithCache(result, extraFields.CacheDebug, lookupScopes) } else { - requestCost = s.calculateBaseCost(result, lookupScopes) + cost = s.calculateBaseCost(result, lookupScopes) + } + + if extraFields != nil && extraFields.GuardrailDebug != nil { + cost += s.CalculateGuardrailCost(extraFields.GuardrailDebug, &lookupScopes) + } + + // Routing-classification embedding cost is budget-attributed only when + // governance explicitly opted the request in. Telemetry can still price it + // independently through RoutingEmbeddingCost. + if extraFields != nil && extraFields.RoutingDebug != nil && extraFields.RoutingDebug.CountTowardBudgets { + cost += s.RoutingEmbeddingCost(extraFields.RoutingDebug, &lookupScopes) + } + + return cost +} + +// RoutingEmbeddingCost calculates the embedding cost of a semantic routing +// classification from its RoutingDebug stamp. Exported so telemetry can price +// routing overhead unconditionally, while CalculateCost folds it into the +// request cost only when RoutingDebug.CountTowardBudgets is set. If scopes is +// nil, an empty LookupScopes is used. +func (s *Store) RoutingEmbeddingCost(routingDebug *schemas.BifrostRoutingDebug, scopes *LookupScopes) float64 { + if routingDebug == nil || routingDebug.ProviderUsed == nil || routingDebug.ModelUsed == nil || routingDebug.InputTokens == nil { + return 0 + } + // Malformed usage must never create a negative sidecar cost that subtracts + // from the request's budget attribution. + if *routingDebug.InputTokens < 0 { + return 0 } - // Handle guardrail judge-call billing - if extraFields.GuardrailDebug == nil { - return requestCost + var lookupScopes LookupScopes + if scopes != nil { + lookupScopes = *scopes + } + // The embedding can use a different provider from the main request, so its + // provider-scoped overrides must be resolved against the embedding provider. + lookupScopes.Provider = *routingDebug.ProviderUsed + pricing := s.resolvePricing(schemas.RoutingInfo{ + Provider: schemas.ModelProvider(*routingDebug.ProviderUsed), + Model: *routingDebug.ModelUsed, + }, schemas.EmbeddingRequest, lookupScopes) + if pricing == nil { + return 0 } - return requestCost + s.CalculateGuardrailCost(extraFields.GuardrailDebug, &lookupScopes) + return float64(*routingDebug.InputTokens) * tieredInputRate(pricing, *routingDebug.InputTokens, serviceTier{}) } // CalculateCostForUsage computes the dollar cost from a bare usage object plus diff --git a/framework/modelcatalog/datasheet/cost_test.go b/framework/modelcatalog/datasheet/cost_test.go index 0291ded28db..b9dade734d5 100644 --- a/framework/modelcatalog/datasheet/cost_test.go +++ b/framework/modelcatalog/datasheet/cost_test.go @@ -1854,6 +1854,156 @@ func TestCalculateGuardrailCostUsesJudgeProviderWithoutCallerSelectedKey(t *test assert.InDelta(t, 30.0, cost, 1e-12) } +// ========================================================================= +// 10b. Semantic routing billing (RoutingDebug) +// ========================================================================= + +// routingDebugFor builds a RoutingDebug stamp for the given embedding call. +func routingDebugFor(provider, model string, inputTokens int, countTowardBudgets bool) *schemas.BifrostRoutingDebug { + return &schemas.BifrostRoutingDebug{ + ProviderUsed: &provider, + ModelUsed: &model, + InputTokens: &inputTokens, + CountTowardBudgets: countTowardBudgets, + } +} + +// routingBillingTestStore has chat pricing for the request model and embedding +// pricing for the routing classifier's model. +func routingBillingTestStore() *Store { + return testStoreWithPricing(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), + }, + }) +} + +func TestCalculateCost_RoutingDebugFlagOff(t *testing.T) { + s := routingBillingTestStore() + + 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"), + RoutingDebug: routingDebugFor("openai", "text-embedding-3-small", 500, false), + }, + }, + } + + // count_toward_budgets off: base cost only, routing embed adds zero. + cost := s.CalculateCost(resp, nil) + assert.InDelta(t, 0.0125, cost, 1e-12) +} + +func TestCalculateCost_RoutingDebugFlagOn(t *testing.T) { + s := routingBillingTestStore() + + 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"), + RoutingDebug: routingDebugFor("openai", "text-embedding-3-small", 500, true), + }, + }, + } + + // Base cost: 1000*0.000005 + 500*0.000015 = 0.0125 + // Routing embedding cost: 500 * 0.00000002 = 0.00001 + cost := s.CalculateCost(resp, nil) + assert.InDelta(t, 0.01251, cost, 1e-12) +} + +func TestCalculateCost_RoutingDebugComposesWithCacheDebug(t *testing.T) { + embProvider := "openai" + embModel := "text-embedding-3-small" + embTokens := 500 + + s := routingBillingTestStore() + + 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, + }, + RoutingDebug: routingDebugFor("openai", "text-embedding-3-small", 250, true), + }, + }, + } + + // Base cost: 0.0125 + // Cache embedding cost: 500 * 0.00000002 = 0.00001 + // Routing embedding cost: 250 * 0.00000002 = 0.000005 + cost := s.CalculateCost(resp, nil) + assert.InDelta(t, 0.012515, cost, 1e-12) +} + +func TestRoutingEmbeddingCost_Standalone(t *testing.T) { + s := routingBillingTestStore() + + // Priced regardless of the CountTowardBudgets flag — telemetry uses this + // to report routing overhead unconditionally. + cost := s.RoutingEmbeddingCost(routingDebugFor("openai", "text-embedding-3-small", 500, false), nil) + assert.InDelta(t, 0.00001, cost, 1e-12) +} + +func TestRoutingEmbeddingCost_MissingFields(t *testing.T) { + s := routingBillingTestStore() + + assert.Equal(t, 0.0, s.RoutingEmbeddingCost(nil, nil)) + assert.Equal(t, 0.0, s.RoutingEmbeddingCost(&schemas.BifrostRoutingDebug{}, nil)) + + provider := "openai" + model := "text-embedding-3-small" + assert.Equal(t, 0.0, s.RoutingEmbeddingCost(&schemas.BifrostRoutingDebug{ + ProviderUsed: &provider, + ModelUsed: &model, + // InputTokens missing + }, nil)) +} + +// A malformed provider usage payload must never produce a negative routing +// cost: that would subtract from the request's billed cost and under-count its +// budget attribution. +func TestRoutingEmbeddingCost_NegativeInputTokensRejected(t *testing.T) { + s := routingBillingTestStore() + + assert.Equal(t, 0.0, s.RoutingEmbeddingCost(routingDebugFor("openai", "text-embedding-3-small", -500, true), nil)) + + 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"), + RoutingDebug: routingDebugFor("openai", "text-embedding-3-small", -500, true), + }, + }, + } + + // Base cost only — the negative routing embed contributes nothing. + assert.InDelta(t, 0.0125, s.CalculateCost(resp, nil), 1e-12) + + // Zero stays valid and free; positive counts are unaffected. + assert.Equal(t, 0.0, s.RoutingEmbeddingCost(routingDebugFor("openai", "text-embedding-3-small", 0, true), nil)) + assert.InDelta(t, 0.00001, s.RoutingEmbeddingCost(routingDebugFor("openai", "text-embedding-3-small", 500, true), nil), 1e-12) +} + // ========================================================================= // 11. CalculateCost integration — end-to-end // ========================================================================= diff --git a/framework/modelcatalog/pricing.go b/framework/modelcatalog/pricing.go index 91a7a65b0f3..22fe5837e95 100644 --- a/framework/modelcatalog/pricing.go +++ b/framework/modelcatalog/pricing.go @@ -64,6 +64,15 @@ func (mc *ModelCatalog) CalculateCostForUsage(usage *schemas.BifrostLLMUsage, pr return mc.datasheet.CalculateCostForUsage(usage, provider, model, requestType, (*datasheet.LookupScopes)(scopes)) } +// CalculateRoutingEmbeddingCost prices the embedding call recorded in a +// response's RoutingDebug stamp, independent of its CountTowardBudgets flag — +// telemetry uses it to report routing overhead unconditionally, while +// CalculateCost folds the same amount into the request's cost only when the +// flag is set. +func (mc *ModelCatalog) CalculateRoutingEmbeddingCost(routingDebug *schemas.BifrostRoutingDebug, scopes *PricingLookupScopes) float64 { + return mc.datasheet.RoutingEmbeddingCost(routingDebug, (*datasheet.LookupScopes)(scopes)) +} + // CalculateGuardrailCost computes the aggregate cost of guardrail judge calls. func (mc *ModelCatalog) CalculateGuardrailCost(debug *schemas.BifrostGuardrailDebug, scopes *PricingLookupScopes) float64 { return mc.datasheet.CalculateGuardrailCost(debug, (*datasheet.LookupScopes)(scopes)) diff --git a/plugins/governance/embedding.go b/plugins/governance/embedding.go index 70b4d61f792..9bf32909aed 100644 --- a/plugins/governance/embedding.go +++ b/plugins/governance/embedding.go @@ -26,6 +26,21 @@ var ErrEmbeddingTimeout = errors.New("embedding request timed out") // classification. It mirrors the signature of bifrost.Client.EmbeddingRequest. type EmbeddingRequestExecutor func(ctx *schemas.BifrostContext, req *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) +// routingEmbedUsageContextKey carries a *routingEmbedUsage on the triggering +// request's context from classification (PreRequestHook) to PostLLMHook, where +// it is stamped onto the response as ExtraFields.RoutingDebug. It mirrors the +// semantic cache's per-request state handoff between its pre and post hooks. +const routingEmbedUsageContextKey schemas.BifrostContextKey = "bf-governance-routing-embed-usage" + +// routingEmbedUsage is the recorded usage of one semantic classification +// embedding call made on behalf of a request. +type routingEmbedUsage struct { + Provider string + Model string + InputTokens int + CountTowardBudgets bool +} + // EmbeddingExecutorSetter is implemented by governance plugins that accept an // embedding request executor. The HTTP server wires the executor after the // bifrost client is constructed (the plugin itself is built while the client @@ -35,12 +50,37 @@ type EmbeddingExecutorSetter interface { SetEmbeddingRequestExecutor(EmbeddingRequestExecutor) } +// WarmupEmbedUsageObserver receives the usage of every warmup/boot embedding +// call made by semantic complexity routing. Warmup embeds have no triggering +// request — there is no response to stamp — so this callback is how their cost +// reaches telemetry. The HTTP server wires it to the telemetry plugin's routing +// overhead counters. Budget attribution is separate: settleWarmupEmbedUsage +// bills the admin-owned provider/model-level budgets directly when +// count_toward_budgets is on. +type WarmupEmbedUsageObserver func(provider, model string, inputTokens int) + +// WarmupEmbedUsageObserverSetter is implemented by governance plugins that +// accept a warmup embedding usage observer. Wired by the HTTP server like +// EmbeddingExecutorSetter; wrappers that embed *GovernancePlugin satisfy this +// via method promotion. +type WarmupEmbedUsageObserverSetter interface { + SetWarmupEmbedUsageObserver(WarmupEmbedUsageObserver) +} + // ComplexityVectorStoreSetter is implemented by governance plugins that accept // Bifrost's configured VectorStore for semantic complexity routing. type ComplexityVectorStoreSetter interface { SetComplexityVectorStore(vectorstore.VectorStore) } +// warmupEmbeddingTimeout bounds one warmup embedding call, whether that is a +// batch of exemplars or a single-input fallback. Warmup runs in the background +// with no request waiting on it, so it must NOT inherit semantic.Timeout — that +// is the hot-path budget (100ms by default), which a 32-exemplar batch cannot +// possibly meet. It stays bounded so a hung provider cannot pin the warmup +// worker forever. +const warmupEmbeddingTimeout = 60 * time.Second + // SetEmbeddingRequestExecutor wires up the function used to call out to the // embedding provider. Without it, semantic complexity classification publishes // no tier. Safe for concurrent use with classification and plugin reloads. @@ -68,26 +108,111 @@ func (p *GovernancePlugin) SetComplexityVectorStore(store vectorstore.VectorStor } // embedComplexityText adapts Governance's Bifrost-aware embedding path to the -// classifier's context-based dependency without attributing token usage here. +// classifier's context-based dependency. It records the embed's usage on the +// triggering request's context so PostLLMHook can stamp RoutingDebug; budget +// attribution itself happens later in cost calculation, never here. func (p *GovernancePlugin) embedComplexityText(ctx context.Context, semantic *complexity.SemanticConfig, text string) ([]float32, error) { + // A *schemas.BifrostContext means a live request is blocked on this embed; a + // plain context means warmup's single-input fallback (batch embeds + // unsupported by the provider/model). The two get very different budgets: + // the hot path must not wait, warmup can. + _, isRequest := ctx.(*schemas.BifrostContext) + timeout := warmupEmbeddingTimeout + if isRequest { + timeout = requestEmbeddingTimeout(semantic) + } + embeddingCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) defer embeddingCtx.Cancel() - embedding, _, err := p.generateEmbedding(embeddingCtx, semantic, text) + embedding, inputTokens, err := p.generateEmbedding(embeddingCtx, semantic, text, timeout) if err != nil { return nil, err } + if isRequest { + recordRoutingEmbedUsage(ctx, semantic, inputTokens) + } else { + p.settleWarmupEmbedUsage(semantic, inputTokens) + } return embedding, nil } +// requestEmbeddingTimeout is the configured hot-path budget for one +// classification embed, which a live request is waiting on. +func requestEmbeddingTimeout(semantic *complexity.SemanticConfig) time.Duration { + if semantic != nil && semantic.Timeout > 0 { + return semantic.Timeout + } + return configstore.DefaultComplexitySemanticTimeout +} + +// recordRoutingEmbedUsage stashes one classification embed's usage on the +// triggering request's context. Warmup embeds arrive on plain background +// contexts (never a *schemas.BifrostContext), so they are naturally excluded — +// boot/warmup embedding cost is never stamped or attributed to any request. +func recordRoutingEmbedUsage(ctx context.Context, semantic *complexity.SemanticConfig, inputTokens int) { + bfCtx, ok := ctx.(*schemas.BifrostContext) + if !ok || semantic == nil { + return + } + bfCtx.SetValue(routingEmbedUsageContextKey, &routingEmbedUsage{ + Provider: string(semantic.Provider), + Model: semantic.EmbeddingModel, + InputTokens: inputTokens, + CountTowardBudgets: semantic.CountTowardBudgets, + }) +} + +// stampRoutingDebug attaches routing-classification telemetry to the response +// when this request ran a semantic routing embed. Stamped on every such +// response for visibility, independent of count_toward_budgets — the flag rides +// in the struct because cost calculation (modelcatalog) cannot see governance +// config. For streams, only the final chunk is stamped, matching where cost is +// billed and mirroring the semantic cache's stamping. +func stampRoutingDebug(ctx *schemas.BifrostContext, result *schemas.BifrostResponse, requestType schemas.RequestType, isFinalChunk bool) { + if result == nil { + return + } + if bifrost.IsStreamRequestType(requestType) && !isFinalChunk { + return + } + usage, ok := ctx.Value(routingEmbedUsageContextKey).(*routingEmbedUsage) + if !ok || usage == nil { + return + } + extraFields := result.GetExtraFields() + if extraFields == nil { + return + } + // InputTokens is provider-derived and must be non-negative before it reaches + // cost calculation: a negative count prices to a negative routing charge, + // which would subtract from the request's cost and its budget attribution + // when CountTowardBudgets is set. generateEmbeddings already drops negative + // provider usage; this is the invariant at the point of stamping, so any + // other writer of the usage key cannot bypass it. + inputTokens := usage.InputTokens + if inputTokens < 0 { + inputTokens = 0 + } + extraFields.RoutingDebug = &schemas.BifrostRoutingDebug{ + ProviderUsed: bifrost.Ptr(usage.Provider), + ModelUsed: bifrost.Ptr(usage.Model), + InputTokens: &inputTokens, + CountTowardBudgets: usage.CountTowardBudgets, + } +} + // embedComplexityTexts adapts the same internal embedding path for bounded -// warmup batches. It preserves response order by EmbeddingData.Index. +// warmup batches. It preserves response order by EmbeddingData.Index. Batch +// embeds are warmup-only, so usage always goes to the warmup observer and is +// never attributed to a request. func (p *GovernancePlugin) embedComplexityTexts(ctx context.Context, semantic *complexity.SemanticConfig, texts []string) ([][]float32, error) { embeddingCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline) defer embeddingCtx.Cancel() - embeddings, _, err := p.generateEmbeddings(embeddingCtx, semantic, texts) + embeddings, inputTokens, err := p.generateEmbeddings(embeddingCtx, semantic, texts, warmupEmbeddingTimeout) if err != nil { return nil, err } + p.settleWarmupEmbedUsage(semantic, inputTokens) return embeddings, nil } @@ -99,6 +224,56 @@ func (p *GovernancePlugin) embeddingExecutor() EmbeddingRequestExecutor { return nil } +// SetWarmupEmbedUsageObserver wires (or clears, with nil) the callback that +// receives warmup embedding usage. Safe for concurrent use with warmup and +// plugin reloads. +func (p *GovernancePlugin) SetWarmupEmbedUsageObserver(observer WarmupEmbedUsageObserver) { + if observer == nil { + p.warmupEmbedUsageObserver.Store(nil) + return + } + p.warmupEmbedUsageObserver.Store(&observer) +} + +// settleWarmupEmbedUsage settles one warmup embedding call: it reports the +// usage to the wired observer (telemetry, always) and, when +// count_toward_budgets is on, attributes the cost to the admin-owned +// provider-level and global model-level budgets — the same ledger the tracker +// uses for usage with no virtual key. Warmup has no triggering request, so no +// VK/team/customer budget is ever touched: there is no tenant to bill, only +// the platform-level budgets on the embedding provider/model. Called only from +// paths with no triggering request. +func (p *GovernancePlugin) settleWarmupEmbedUsage(semantic *complexity.SemanticConfig, inputTokens int) { + if semantic == nil { + return + } + if ptr := p.warmupEmbedUsageObserver.Load(); ptr != nil { + (*ptr)(string(semantic.Provider), semantic.EmbeddingModel, inputTokens) + } + + if !semantic.CountTowardBudgets || p.modelCatalog == nil || p.store == nil { + return + } + provider := string(semantic.Provider) + model := semantic.EmbeddingModel + tokens := inputTokens + cost := p.modelCatalog.CalculateRoutingEmbeddingCost(&schemas.BifrostRoutingDebug{ + ProviderUsed: &provider, + ModelUsed: &model, + InputTokens: &tokens, + }, nil) + if cost <= 0 { + return + } + ctx := p.ctx + if ctx == nil { + ctx = context.Background() + } + if err := p.store.UpdateProviderAndModelBudgetUsageInMemory(ctx, model, semantic.Provider, cost); err != nil && p.logger != nil { + p.logger.Error("failed to attribute warmup embedding cost to provider/model budgets: %v", err) + } +} + // CanClassifySemantically reports whether semantic classification is currently // viable. The executor alone is not a sufficient gate — the server wires it // unconditionally; the semantic config decides whether classification is @@ -114,8 +289,8 @@ func (p *GovernancePlugin) CanClassifySemantically(semantic *complexity.Semantic // returns the vector plus the input token count (fed to routing-cost // attribution). The call is bounded by the configured semantic timeout: the // router hot path must never wait on a slow embedding provider. -func (p *GovernancePlugin) generateEmbedding(ctx *schemas.BifrostContext, semantic *complexity.SemanticConfig, text string) ([]float32, int, error) { - embeddings, inputTokens, err := p.generateEmbeddings(ctx, semantic, []string{text}) +func (p *GovernancePlugin) generateEmbedding(ctx *schemas.BifrostContext, semantic *complexity.SemanticConfig, text string, timeout time.Duration) ([]float32, int, error) { + embeddings, inputTokens, err := p.generateEmbeddings(ctx, semantic, []string{text}, timeout) if err != nil { return nil, 0, err } @@ -128,7 +303,7 @@ func (p *GovernancePlugin) generateEmbedding(ctx *schemas.BifrostContext, semant // generateEmbeddings sends one embedding request for an ordered set of texts. // A multi-input response must contain exactly one uniquely indexed vector per // input; otherwise warmup can safely retry through the single-input adapter. -func (p *GovernancePlugin) generateEmbeddings(ctx *schemas.BifrostContext, semantic *complexity.SemanticConfig, texts []string) ([][]float32, int, error) { +func (p *GovernancePlugin) generateEmbeddings(ctx *schemas.BifrostContext, semantic *complexity.SemanticConfig, texts []string, timeout time.Duration) ([][]float32, int, error) { executor := p.embeddingExecutor() if executor == nil { return nil, 0, fmt.Errorf("embedding request executor is not configured") @@ -140,7 +315,6 @@ func (p *GovernancePlugin) generateEmbeddings(ctx *schemas.BifrostContext, seman return nil, 0, fmt.Errorf("embedding input is empty") } - timeout := semantic.Timeout if timeout <= 0 { timeout = configstore.DefaultComplexitySemanticTimeout } @@ -188,7 +362,11 @@ func (p *GovernancePlugin) generateEmbeddings(ctx *schemas.BifrostContext, seman return nil, 0, fmt.Errorf("no embeddings returned from provider") } inputTokens := 0 - if response.Usage != nil { + // Provider-reported usage is untrusted input: a negative count would flow + // into the RoutingDebug stamp and from there into cost calculation and + // warmup budget attribution, where it would subtract from billed usage. + // Drop it to zero — the embed still happened, we just cannot price it. + if response.Usage != nil && response.Usage.TotalTokens > 0 { inputTokens = response.Usage.TotalTokens } diff --git a/plugins/governance/embedding_test.go b/plugins/governance/embedding_test.go index 46da35d1fe7..5c58c9f0fdb 100644 --- a/plugins/governance/embedding_test.go +++ b/plugins/governance/embedding_test.go @@ -1,11 +1,17 @@ package governance import ( + "context" "errors" + "path/filepath" "testing" "time" "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" + "github.com/maximhq/bifrost/framework/modelcatalog/datasheet" "github.com/maximhq/bifrost/plugins/governance/complexity" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -49,7 +55,7 @@ func TestGenerateEmbeddingDecodesAllEncodings(t *testing.T) { ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) defer ctx.Cancel() - vector, tokens, err := plugin.generateEmbedding(ctx, testEmbeddingSemanticConfig(), "hello") + vector, tokens, err := plugin.generateEmbedding(ctx, testEmbeddingSemanticConfig(), "hello", requestEmbeddingTimeout(testEmbeddingSemanticConfig())) require.NoError(t, err) assert.Equal(t, tt.want, vector) assert.Equal(t, 42, tokens) @@ -78,7 +84,7 @@ func TestGenerateEmbeddingRequestShape(t *testing.T) { callerDeadline := before.Add(50 * cfg.Timeout) ctx := schemas.NewBifrostContext(t.Context(), callerDeadline) defer ctx.Cancel() - _, _, err := plugin.generateEmbedding(ctx, cfg, "classify me") + _, _, err := plugin.generateEmbedding(ctx, cfg, "classify me", requestEmbeddingTimeout(cfg)) require.NoError(t, err) require.NotNil(t, gotReq) @@ -116,7 +122,7 @@ func TestGenerateEmbeddingsBatchesAndRestoresInputOrder(t *testing.T) { ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) defer ctx.Cancel() - embeddings, tokens, err := plugin.generateEmbeddings(ctx, testEmbeddingSemanticConfig(), []string{"first", "second"}) + embeddings, tokens, err := plugin.generateEmbeddings(ctx, testEmbeddingSemanticConfig(), []string{"first", "second"}, requestEmbeddingTimeout(testEmbeddingSemanticConfig())) require.NoError(t, err) assert.Equal(t, [][]float32{{1, 0}, {0, 1}}, embeddings) assert.Equal(t, 7, tokens) @@ -132,7 +138,7 @@ func TestGenerateEmbeddingsSignalsUnsupportedBatchShape(t *testing.T) { ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) defer ctx.Cancel() - _, _, err := plugin.generateEmbeddings(ctx, testEmbeddingSemanticConfig(), []string{"first", "second"}) + _, _, err := plugin.generateEmbeddings(ctx, testEmbeddingSemanticConfig(), []string{"first", "second"}, requestEmbeddingTimeout(testEmbeddingSemanticConfig())) require.Error(t, err) assert.True(t, errors.Is(err, complexity.ErrBatchEmbeddingsUnsupported)) } @@ -156,7 +162,7 @@ func TestGenerateEmbeddingTimeoutCancelsCall(t *testing.T) { ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) defer ctx.Cancel() start := time.Now() - _, _, err := plugin.generateEmbedding(ctx, cfg, "slow") + _, _, err := plugin.generateEmbedding(ctx, cfg, "slow", requestEmbeddingTimeout(cfg)) require.Error(t, err) // Tolerant of scheduler jitter, but still tight enough that only the 20ms // configuration can satisfy it — a second-scale budget would not. @@ -208,13 +214,73 @@ func TestGenerateEmbeddingDistinguishesTimeoutFromOtherFailures(t *testing.T) { ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) defer ctx.Cancel() - _, _, err := plugin.generateEmbedding(ctx, testEmbeddingSemanticConfig(), "classify me") + _, _, err := plugin.generateEmbedding(ctx, testEmbeddingSemanticConfig(), "classify me", requestEmbeddingTimeout(testEmbeddingSemanticConfig())) require.Error(t, err) assert.Equal(t, tt.wantTimeout, errors.Is(err, ErrEmbeddingTimeout)) }) } } +// TestWarmupEmbedsDoNotInheritTheRequestTimeout is a regression guard: warmup +// used to run through semantic.Timeout, the hot-path budget (100ms by default). +// A 32-exemplar batch cannot finish in that window, so every warmup failed with +// a 504 and semantic routing silently served its fallback forever. +func TestWarmupEmbedsDoNotInheritTheRequestTimeout(t *testing.T) { + semantic := testEmbeddingSemanticConfig() + semantic.Timeout = 10 * time.Millisecond + + t.Run("batch warmup", func(t *testing.T) { + plugin := &GovernancePlugin{} + var deadline time.Time + plugin.SetEmbeddingRequestExecutor(func(ctx *schemas.BifrostContext, _ *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { + deadline, _ = ctx.Deadline() + return &schemas.BifrostEmbeddingResponse{ + Data: []schemas.EmbeddingData{ + {Index: 0, Embedding: schemas.EmbeddingStruct{EmbeddingArray: []float64{1}}}, + {Index: 1, Embedding: schemas.EmbeddingStruct{EmbeddingArray: []float64{2}}}, + }, + Usage: &schemas.BifrostLLMUsage{TotalTokens: 2}, + }, nil + }) + + before := time.Now() + _, err := plugin.embedComplexityTexts(context.Background(), semantic, []string{"a", "b"}) + require.NoError(t, err) + assert.Greater(t, deadline.Sub(before), time.Second, "warmup batch must not run on the hot-path budget") + }) + + t.Run("single-input warmup fallback", func(t *testing.T) { + plugin := &GovernancePlugin{} + var deadline time.Time + plugin.SetEmbeddingRequestExecutor(func(ctx *schemas.BifrostContext, _ *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { + deadline, _ = ctx.Deadline() + return embeddingResponse(schemas.EmbeddingStruct{EmbeddingArray: []float64{1}}, 1), nil + }) + + // A plain context is what the classifier passes during warmup. + before := time.Now() + _, err := plugin.embedComplexityText(context.Background(), semantic, "exemplar") + require.NoError(t, err) + assert.Greater(t, deadline.Sub(before), time.Second, "warmup fallback must not run on the hot-path budget") + }) + + t.Run("request classification still honors the configured budget", func(t *testing.T) { + plugin := &GovernancePlugin{} + var deadline time.Time + plugin.SetEmbeddingRequestExecutor(func(ctx *schemas.BifrostContext, _ *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { + deadline, _ = ctx.Deadline() + return embeddingResponse(schemas.EmbeddingStruct{EmbeddingArray: []float64{1}}, 1), nil + }) + + requestCtx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) + defer requestCtx.Cancel() + before := time.Now() + _, err := plugin.embedComplexityText(requestCtx, semantic, "classify me") + require.NoError(t, err) + assert.LessOrEqual(t, deadline.Sub(before), 100*time.Millisecond, "a live request must stay on its configured budget") + }) +} + func TestGenerateEmbeddingGuards(t *testing.T) { okExecutor := func(ctx *schemas.BifrostContext, req *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { return embeddingResponse(schemas.EmbeddingStruct{EmbeddingArray: []float64{1}}, 1), nil @@ -224,7 +290,7 @@ func TestGenerateEmbeddingGuards(t *testing.T) { plugin := &GovernancePlugin{} ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) defer ctx.Cancel() - _, _, err := plugin.generateEmbedding(ctx, testEmbeddingSemanticConfig(), "x") + _, _, err := plugin.generateEmbedding(ctx, testEmbeddingSemanticConfig(), "x", requestEmbeddingTimeout(testEmbeddingSemanticConfig())) require.ErrorContains(t, err, "executor is not configured") }) @@ -233,7 +299,7 @@ func TestGenerateEmbeddingGuards(t *testing.T) { plugin.SetEmbeddingRequestExecutor(okExecutor) ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) defer ctx.Cancel() - _, _, err := plugin.generateEmbedding(ctx, nil, "x") + _, _, err := plugin.generateEmbedding(ctx, nil, "x", requestEmbeddingTimeout(testEmbeddingSemanticConfig())) require.ErrorContains(t, err, "not configured") }) @@ -244,7 +310,7 @@ func TestGenerateEmbeddingGuards(t *testing.T) { }) ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) defer ctx.Cancel() - _, _, err := plugin.generateEmbedding(ctx, testEmbeddingSemanticConfig(), "x") + _, _, err := plugin.generateEmbedding(ctx, testEmbeddingSemanticConfig(), "x", requestEmbeddingTimeout(testEmbeddingSemanticConfig())) require.ErrorContains(t, err, "no embeddings returned") }) @@ -254,11 +320,288 @@ func TestGenerateEmbeddingGuards(t *testing.T) { plugin.SetEmbeddingRequestExecutor(nil) ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) defer ctx.Cancel() - _, _, err := plugin.generateEmbedding(ctx, testEmbeddingSemanticConfig(), "x") + _, _, err := plugin.generateEmbedding(ctx, testEmbeddingSemanticConfig(), "x", requestEmbeddingTimeout(testEmbeddingSemanticConfig())) require.ErrorContains(t, err, "executor is not configured") }) } +func TestEmbedComplexityTextRecordsRoutingUsage(t *testing.T) { + plugin := &GovernancePlugin{} + plugin.SetEmbeddingRequestExecutor(func(_ *schemas.BifrostContext, _ *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { + return embeddingResponse(schemas.EmbeddingStruct{EmbeddingArray: []float64{1, 0}}, 42), nil + }) + + cfg := testEmbeddingSemanticConfig() + cfg.CountTowardBudgets = true + + ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) + defer ctx.Cancel() + _, err := plugin.embedComplexityText(ctx, cfg, "classify me") + require.NoError(t, err) + + usage, ok := ctx.Value(routingEmbedUsageContextKey).(*routingEmbedUsage) + require.True(t, ok, "classification embed must record usage on the request context") + assert.Equal(t, "openai", usage.Provider) + assert.Equal(t, "text-embedding-3-small", usage.Model) + assert.Equal(t, 42, usage.InputTokens) + assert.True(t, usage.CountTowardBudgets) +} + +// A provider that reports a negative token count must not have it recorded: +// the stamp feeds cost calculation and warmup budget attribution, where a +// negative count would subtract from billed usage. +func TestEmbedComplexityTextDropsNegativeProviderUsage(t *testing.T) { + plugin := &GovernancePlugin{} + plugin.SetEmbeddingRequestExecutor(func(_ *schemas.BifrostContext, _ *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { + return embeddingResponse(schemas.EmbeddingStruct{EmbeddingArray: []float64{1, 0}}, -42), nil + }) + + cfg := testEmbeddingSemanticConfig() + cfg.CountTowardBudgets = true + + ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) + defer ctx.Cancel() + _, err := plugin.embedComplexityText(ctx, cfg, "classify me") + require.NoError(t, err) + + usage, ok := ctx.Value(routingEmbedUsageContextKey).(*routingEmbedUsage) + require.True(t, ok) + assert.Equal(t, 0, usage.InputTokens, "negative provider usage must not reach budget accounting") +} + +// warmupObservation captures one WarmupEmbedUsageObserver invocation. +type warmupObservation struct { + Provider string + Model string + InputTokens int +} + +func TestEmbedComplexityTextWarmupPathObservesInsteadOfRecording(t *testing.T) { + plugin := &GovernancePlugin{} + plugin.SetEmbeddingRequestExecutor(func(_ *schemas.BifrostContext, _ *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { + return embeddingResponse(schemas.EmbeddingStruct{EmbeddingArray: []float64{1, 0}}, 42), nil + }) + var observed []warmupObservation + plugin.SetWarmupEmbedUsageObserver(func(provider, model string, inputTokens int) { + observed = append(observed, warmupObservation{provider, model, inputTokens}) + }) + + // Warmup's single-input fallback runs on plain background contexts, never a + // *schemas.BifrostContext — its embeds go to the warmup observer, not to + // request attribution. + _, err := plugin.embedComplexityText(t.Context(), testEmbeddingSemanticConfig(), "warmup exemplar") + require.NoError(t, err) + require.Len(t, observed, 1) + assert.Equal(t, warmupObservation{"openai", "text-embedding-3-small", 42}, observed[0]) +} + +func TestEmbedComplexityTextWarmupPathWithoutObserver(t *testing.T) { + plugin := &GovernancePlugin{} + plugin.SetEmbeddingRequestExecutor(func(_ *schemas.BifrostContext, _ *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { + return embeddingResponse(schemas.EmbeddingStruct{EmbeddingArray: []float64{1, 0}}, 42), nil + }) + + // No observer wired (SDK usage, or before the server wires it): the warmup + // path must still work and must not panic. + _, err := plugin.embedComplexityText(t.Context(), testEmbeddingSemanticConfig(), "warmup exemplar") + require.NoError(t, err) +} + +func TestEmbedComplexityTextsObservesWarmupNeverRecordsRoutingUsage(t *testing.T) { + plugin := &GovernancePlugin{} + plugin.SetEmbeddingRequestExecutor(func(_ *schemas.BifrostContext, _ *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { + return &schemas.BifrostEmbeddingResponse{ + Data: []schemas.EmbeddingData{ + {Index: 0, Embedding: schemas.EmbeddingStruct{EmbeddingArray: []float64{1, 0}}}, + {Index: 1, Embedding: schemas.EmbeddingStruct{EmbeddingArray: []float64{0, 1}}}, + }, + Usage: &schemas.BifrostLLMUsage{TotalTokens: 7}, + }, nil + }) + var observed []warmupObservation + plugin.SetWarmupEmbedUsageObserver(func(provider, model string, inputTokens int) { + observed = append(observed, warmupObservation{provider, model, inputTokens}) + }) + + // Batch embeds are warmup-only: even on a request context they observe as + // warmup and never attribute usage — only per-request classification + // embeds do. + ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) + defer ctx.Cancel() + _, err := plugin.embedComplexityTexts(ctx, testEmbeddingSemanticConfig(), []string{"a", "b"}) + require.NoError(t, err) + assert.Nil(t, ctx.Value(routingEmbedUsageContextKey)) + require.Len(t, observed, 1) + assert.Equal(t, warmupObservation{"openai", "text-embedding-3-small", 7}, observed[0]) +} + +func TestRequestClassificationEmbedDoesNotObserveWarmup(t *testing.T) { + plugin := &GovernancePlugin{} + plugin.SetEmbeddingRequestExecutor(func(_ *schemas.BifrostContext, _ *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { + return embeddingResponse(schemas.EmbeddingStruct{EmbeddingArray: []float64{1, 0}}, 42), nil + }) + var observed []warmupObservation + plugin.SetWarmupEmbedUsageObserver(func(provider, model string, inputTokens int) { + observed = append(observed, warmupObservation{provider, model, inputTokens}) + }) + + // A classification embed on a request context records usage for the + // RoutingDebug stamp; it must NOT also fire the warmup observer, or the + // request phase would double-count in telemetry. + ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) + defer ctx.Cancel() + _, err := plugin.embedComplexityText(ctx, testEmbeddingSemanticConfig(), "classify me") + require.NoError(t, err) + assert.NotNil(t, ctx.Value(routingEmbedUsageContextKey)) + assert.Empty(t, observed) +} + +func TestStampRoutingDebug(t *testing.T) { + newCtxWithUsage := func(t *testing.T, countTowardBudgets bool) *schemas.BifrostContext { + t.Helper() + ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) + t.Cleanup(ctx.Cancel) + ctx.SetValue(routingEmbedUsageContextKey, &routingEmbedUsage{ + Provider: "openai", + Model: "text-embedding-3-small", + InputTokens: 42, + CountTowardBudgets: countTowardBudgets, + }) + return ctx + } + newChatResult := func() *schemas.BifrostResponse { + return &schemas.BifrostResponse{ChatResponse: &schemas.BifrostChatResponse{}} + } + + t.Run("stamps regardless of budget flag", func(t *testing.T) { + for _, flag := range []bool{false, true} { + result := newChatResult() + stampRoutingDebug(newCtxWithUsage(t, flag), result, schemas.ChatCompletionRequest, false) + + rd := result.GetExtraFields().RoutingDebug + require.NotNil(t, rd, "routing debug must be stamped whenever a routing embed ran (flag=%v)", flag) + require.NotNil(t, rd.ProviderUsed) + assert.Equal(t, "openai", *rd.ProviderUsed) + require.NotNil(t, rd.ModelUsed) + assert.Equal(t, "text-embedding-3-small", *rd.ModelUsed) + require.NotNil(t, rd.InputTokens) + assert.Equal(t, 42, *rd.InputTokens) + assert.Equal(t, flag, rd.CountTowardBudgets) + } + }) + + t.Run("stream stamps only the final chunk", func(t *testing.T) { + ctx := newCtxWithUsage(t, false) + + intermediate := newChatResult() + stampRoutingDebug(ctx, intermediate, schemas.ChatCompletionStreamRequest, false) + assert.Nil(t, intermediate.GetExtraFields().RoutingDebug) + + final := newChatResult() + stampRoutingDebug(ctx, final, schemas.ChatCompletionStreamRequest, true) + assert.NotNil(t, final.GetExtraFields().RoutingDebug) + }) + + t.Run("no usage recorded means no stamp", func(t *testing.T) { + ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) + defer ctx.Cancel() + result := newChatResult() + stampRoutingDebug(ctx, result, schemas.ChatCompletionRequest, false) + assert.Nil(t, result.GetExtraFields().RoutingDebug) + }) + + t.Run("nil result is a no-op", func(t *testing.T) { + stampRoutingDebug(newCtxWithUsage(t, true), nil, schemas.ChatCompletionRequest, false) + }) + + // The stamp feeds cost calculation, so a negative token count must never + // leave this function — it would price to a negative routing charge and + // subtract from the request's budget attribution. + t.Run("negative input tokens are rejected", func(t *testing.T) { + ctx := schemas.NewBifrostContext(t.Context(), schemas.NoDeadline) + defer ctx.Cancel() + ctx.SetValue(routingEmbedUsageContextKey, &routingEmbedUsage{ + Provider: "openai", + Model: "text-embedding-3-small", + InputTokens: -42, + CountTowardBudgets: true, + }) + + result := newChatResult() + stampRoutingDebug(ctx, result, schemas.ChatCompletionRequest, false) + + rd := result.GetExtraFields().RoutingDebug + require.NotNil(t, rd, "the embed still ran, so it stays observable") + require.NotNil(t, rd.InputTokens) + assert.Equal(t, 0, *rd.InputTokens) + }) +} + +// newOfflinePricingCatalog builds a ModelCatalog from the committed pricing +// testdata via a file:// URL (no network). The testdata includes +// text-embedding-3-small at $0.00000002 per input token. +func newOfflinePricingCatalog(t *testing.T) *modelcatalog.ModelCatalog { + t.Helper() + abs, err := filepath.Abs("../../framework/modelcatalog/datasheet/testdata/pricing.json") + require.NoError(t, err) + ds := datasheet.New(nil, NewMockLogger(), datasheet.Config{URL: "file://" + abs}) + require.NoError(t, ds.LoadFromURLIntoMemory(context.Background())) + return modelcatalog.NewTestCatalogWithDatasheet(ds) +} + +// newWarmupBudgetFixture wires a plugin over a store whose "openai" provider +// carries a provider-level budget — the admin-owned ledger warmup embeds are +// attributed to when count_toward_budgets is on. +func newWarmupBudgetFixture(t *testing.T) (*GovernancePlugin, GovernanceStore) { + t.Helper() + logger := NewMockLogger() + budget := buildBudgetWithUsage("provider-budget", 1000.0, 0.0, "1d") + budgetID := budget.ID + store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ + Budgets: []configstoreTables.TableBudget{*budget}, + Providers: []configstoreTables.TableProvider{{Name: "openai", BudgetID: &budgetID}}, + }, nil) + require.NoError(t, err) + plugin := &GovernancePlugin{ + ctx: context.Background(), + store: store, + modelCatalog: newOfflinePricingCatalog(t), + logger: logger, + } + return plugin, store +} + +func TestSettleWarmupEmbedUsageAttributesProviderBudget(t *testing.T) { + plugin, store := newWarmupBudgetFixture(t) + cfg := testEmbeddingSemanticConfig() + cfg.CountTowardBudgets = true + + plugin.settleWarmupEmbedUsage(cfg, 1000) + + // 1000 tokens × $0.00000002/token (text-embedding-3-small in testdata). + usage := store.GetGovernanceData(context.Background()).Budgets["provider-budget"].CurrentUsage + assert.InDelta(t, 0.00002, usage, 1e-12) +} + +func TestSettleWarmupEmbedUsageFlagOffLeavesBudgetsUntouched(t *testing.T) { + plugin, store := newWarmupBudgetFixture(t) + + // count_toward_budgets defaults to off: warmup cost stays telemetry-only. + plugin.settleWarmupEmbedUsage(testEmbeddingSemanticConfig(), 1000) + + usage := store.GetGovernanceData(context.Background()).Budgets["provider-budget"].CurrentUsage + assert.Equal(t, 0.0, usage) +} + +func TestSettleWarmupEmbedUsageWithoutStoreOrCatalog(t *testing.T) { + // Bare plugin (no store, no catalog, no observer): flag on must be a + // harmless no-op, not a panic — SDK callers may never wire these. + plugin := &GovernancePlugin{} + cfg := testEmbeddingSemanticConfig() + cfg.CountTowardBudgets = true + plugin.settleWarmupEmbedUsage(cfg, 1000) +} + func TestCanClassifySemantically(t *testing.T) { executor := func(ctx *schemas.BifrostContext, req *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { return nil, nil diff --git a/plugins/governance/main.go b/plugins/governance/main.go index b0316b2363b..6f0cbf63c87 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -103,6 +103,12 @@ type GovernancePlugin struct { // then and during teardown. Atomic because classification reads it on the // request hot path while plugin reloads may re-wire it. embeddingRequestExecutor atomic.Pointer[EmbeddingRequestExecutor] + + // warmupEmbedUsageObserver is wired by the HTTP server (to the telemetry + // plugin's routing overhead counters) via SetWarmupEmbedUsageObserver; nil + // until then. Atomic because warmup fires it from a background worker while + // plugin reloads may re-wire it. + warmupEmbedUsageObserver atomic.Pointer[WarmupEmbedUsageObserver] } // Init initializes and returns a governance plugin instance. @@ -1434,6 +1440,10 @@ func (p *GovernancePlugin) PostLLMHook(ctx *schemas.BifrostContext, result *sche isFinalChunk := bifrost.IsFinalChunk(ctx) + // Stamp routing-classification telemetry before postHookWorker runs so its + // CalculateCost call (and every later post-hook, e.g. telemetry) sees it. + stampRoutingDebug(ctx, result, requestType, isFinalChunk) + // Build pricing scopes from context using the governance VK ID (not the raw VK token) pricingScopes := modelcatalog.PricingLookupScopesFromContext(ctx, string(provider)) diff --git a/plugins/telemetry/main.go b/plugins/telemetry/main.go index f62d713ee65..66b628d84ab 100644 --- a/plugins/telemetry/main.go +++ b/plugins/telemetry/main.go @@ -170,6 +170,8 @@ type PrometheusPlugin struct { CacheWriteInputTokens5mTotal *prometheus.CounterVec CacheWriteInputTokens1hTotal *prometheus.CounterVec CostTotal *prometheus.CounterVec + RoutingEmbeddingRequestsTotal *prometheus.CounterVec + RoutingEmbeddingCostTotal *prometheus.CounterVec StreamInterTokenLatencySeconds *prometheus.HistogramVec StreamFirstTokenLatencySeconds *prometheus.HistogramVec RequestRetries *prometheus.HistogramVec @@ -288,6 +290,15 @@ var defaultMCPLabelNames = []string{ // so both exporters report the same quantiles for the same operation. var mcpOperationDurationBuckets = []float64{0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1, 2, 5, 10, 30, 60, 120, 300} +// Values of the phase label on bifrost_routing_embedding_* counters: "request" +// is a per-request classification embed (recorded off the RoutingDebug response +// stamp), "warmup" is a boot/config-change exemplar embed (recorded via +// ObserveWarmupRoutingEmbedding — no request or response exists for those). +const ( + routingEmbeddingPhaseRequest = "request" + routingEmbeddingPhaseWarmup = "warmup" +) + func Init(config *Config, pricingManager *modelcatalog.ModelCatalog, logger schemas.Logger) (*PrometheusPlugin, error) { if config == nil { return nil, fmt.Errorf("config is required") @@ -471,6 +482,29 @@ func Init(config *Config, pricingManager *modelcatalog.ModelCatalog, logger sche append(defaultBifrostLabels, filteredCustomLabels...), ) + // Routing-classification overhead (semantic complexity router embeddings). + // Standalone label set on purpose: these count Bifrost's own routing + // overhead, keyed by the embedding provider/model — not the request's — so + // the canonical defaultBifrostLabelNames (conformance-checked against + // schemas.EnrichmentDims) don't apply. phase distinguishes per-request + // classification embeds from warmup/boot exemplar embeds; summing over + // phase gives total routing overhead. + bifrostRoutingEmbeddingRequestsTotal := factory.NewCounterVec( + prometheus.CounterOpts{ + Name: "bifrost_routing_embedding_requests_total", + Help: "Total number of embedding calls made by semantic routing, labeled by the embedding provider/model and phase (request classification vs warmup).", + }, + []string{"provider", "model", "phase"}, + ) + + bifrostRoutingEmbeddingCostTotal := factory.NewCounterVec( + prometheus.CounterOpts{ + Name: "bifrost_routing_embedding_cost_total", + Help: "Total cost in USD of semantic routing embeddings, labeled by the embedding provider/model and phase (request classification vs warmup). Recorded regardless of whether the cost counts toward budgets.", + }, + []string{"provider", "model", "phase"}, + ) + bifrostStreamInterTokenLatencySeconds := factory.NewHistogramVec( prometheus.HistogramOpts{ Name: "bifrost_stream_inter_token_latency_seconds", @@ -562,6 +596,8 @@ func Init(config *Config, pricingManager *modelcatalog.ModelCatalog, logger sche CacheWriteInputTokens5mTotal: bifrostCacheWriteInputTokens5mTotal, CacheWriteInputTokens1hTotal: bifrostCacheWriteInputTokens1hTotal, CostTotal: bifrostCostTotal, + RoutingEmbeddingRequestsTotal: bifrostRoutingEmbeddingRequestsTotal, + RoutingEmbeddingCostTotal: bifrostRoutingEmbeddingCostTotal, StreamInterTokenLatencySeconds: bifrostStreamInterTokenLatencySeconds, StreamFirstTokenLatencySeconds: bifrostStreamFirstTokenLatencySeconds, RequestRetries: bifrostRequestRetries, @@ -676,6 +712,27 @@ func (p *PrometheusPlugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas. return nil } +// ObserveWarmupRoutingEmbedding records one warmup/boot embedding call made by +// semantic routing (exemplar warmup has no request or response, so it cannot +// ride the RoutingDebug stamp path that request-phase embeds use). The HTTP +// server wires this into the governance plugin's warmup embed usage observer. +// This method is telemetry-only; budget attribution for warmup (provider/model- +// level budgets, gated on count_toward_budgets) happens inside governance. +func (p *PrometheusPlugin) ObserveWarmupRoutingEmbedding(provider, model string, inputTokens int) { + p.RoutingEmbeddingRequestsTotal.WithLabelValues(provider, model, routingEmbeddingPhaseWarmup).Inc() + if p.pricingManager == nil { + return + } + routingDebug := &schemas.BifrostRoutingDebug{ + ProviderUsed: &provider, + ModelUsed: &model, + InputTokens: &inputTokens, + } + if cost := p.pricingManager.CalculateRoutingEmbeddingCost(routingDebug, nil); cost > 0 { + p.RoutingEmbeddingCostTotal.WithLabelValues(provider, model, routingEmbeddingPhaseWarmup).Add(cost) + } +} + // PreLLMHook records the start time of the request in the context. // This time is used later in PostLLMHook to calculate request duration. func (p *PrometheusPlugin) PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { @@ -1029,6 +1086,19 @@ func (p *PrometheusPlugin) PostLLMHook(ctx *schemas.BifrostContext, result *sche } if result != nil { + // Record routing-classification overhead (always-on: independent of + // the count_toward_budgets flag, which only controls whether the + // cost also folds into bifrost_cost_total via CalculateCost). + extraFields := result.GetExtraFields() + if rd := extraFields.RoutingDebug; rd != nil && rd.ProviderUsed != nil && rd.ModelUsed != nil { + p.RoutingEmbeddingRequestsTotal.WithLabelValues(*rd.ProviderUsed, *rd.ModelUsed, routingEmbeddingPhaseRequest).Inc() + if p.pricingManager != nil { + if embeddingCost := p.pricingManager.CalculateRoutingEmbeddingCost(rd, pricingScopes); embeddingCost > 0 { + p.RoutingEmbeddingCostTotal.WithLabelValues(*rd.ProviderUsed, *rd.ModelUsed, routingEmbeddingPhaseRequest).Add(embeddingCost) + } + } + } + // Record input and output tokens var inputTokens, outputTokens int @@ -1099,7 +1169,6 @@ func (p *PrometheusPlugin) PostLLMHook(ctx *schemas.BifrostContext, result *sche } // Record cache hits with cache type - extraFields := result.GetExtraFields() if extraFields.CacheDebug != nil && extraFields.CacheDebug.CacheHit { cacheType := "unknown" if extraFields.CacheDebug.HitType != nil { diff --git a/plugins/telemetry/main_test.go b/plugins/telemetry/main_test.go index 31c1f0190fa..fd78983682d 100644 --- a/plugins/telemetry/main_test.go +++ b/plugins/telemetry/main_test.go @@ -199,6 +199,116 @@ func TestPostLLMHookRequiresStartTime(t *testing.T) { // TestMetricsEnabledGating covers the pull-gateway (/metrics scrape) on/off switch: default-on // when the config omits the field (back-compat), and honoring an explicit value. +// TestRoutingEmbeddingCounters: a response stamped with RoutingDebug increments +// bifrost_routing_embedding_requests_total regardless of the count_toward_budgets +// flag — the flag only controls budget folding (CalculateCost), never telemetry. +// Cost stays unrecorded here because the test plugin has no pricing manager; +// the routing embedding cost math is covered in modelcatalog's datasheet tests. +func TestRoutingEmbeddingCounters(t *testing.T) { + for _, countTowardBudgets := range []bool{false, true} { + p := newTestPlugin(t) + + provider := "openai" + model := "text-embedding-3-small" + inputTokens := 42 + resp := &schemas.BifrostResponse{ChatResponse: &schemas.BifrostChatResponse{ + Usage: &schemas.BifrostLLMUsage{PromptTokens: 11, CompletionTokens: 7, TotalTokens: 18}, + ExtraFields: schemas.BifrostResponseExtraFields{ + RequestType: schemas.ChatCompletionRequest, + RoutingDebug: &schemas.BifrostRoutingDebug{ + ProviderUsed: &provider, + ModelUsed: &model, + InputTokens: &inputTokens, + CountTowardBudgets: countTowardBudgets, + }, + }, + }} + resp.PopulateExtraFields(schemas.ChatCompletionRequest, schemas.ModelProvider(provider), model, model) + + ctx := newHookContext(schemas.ChatCompletionRequest) + if _, _, err := p.PostLLMHook(ctx, resp, nil); err != nil { + t.Fatalf("PostLLMHook (flag=%v): %v", countTowardBudgets, err) + } + + waitForCounter(t, p.registry, "bifrost_routing_embedding_requests_total", 1) + if got := counterTotalWithLabel(t, p.registry, "bifrost_routing_embedding_requests_total", "phase", "request"); got != 1 { + t.Fatalf("request-phase requests counter = %v, want 1", got) + } + if got := counterTotal(t, p.registry, "bifrost_routing_embedding_cost_total"); got != 0 { + t.Fatalf("cost counter without pricing manager = %v, want 0", got) + } + } +} + +// TestObserveWarmupRoutingEmbedding: warmup embeds report through the direct +// observer method (no request/response exists for them) and land under +// phase="warmup", separate from the request-phase series. Cost stays +// unrecorded without a pricing manager, same as the request phase. +func TestObserveWarmupRoutingEmbedding(t *testing.T) { + p := newTestPlugin(t) + + p.ObserveWarmupRoutingEmbedding("openai", "text-embedding-3-small", 7) + p.ObserveWarmupRoutingEmbedding("openai", "text-embedding-3-small", 9) + + if got := counterTotalWithLabel(t, p.registry, "bifrost_routing_embedding_requests_total", "phase", "warmup"); got != 2 { + t.Fatalf("warmup-phase requests counter = %v, want 2", got) + } + if got := counterTotalWithLabel(t, p.registry, "bifrost_routing_embedding_requests_total", "phase", "request"); got != 0 { + t.Fatalf("request-phase requests counter = %v, want 0 (warmup must not leak into it)", got) + } + if got := counterTotal(t, p.registry, "bifrost_routing_embedding_cost_total"); got != 0 { + t.Fatalf("cost counter without pricing manager = %v, want 0", got) + } +} + +// counterTotalWithLabel sums every series of the named counter family whose +// labels include name=value. +func counterTotalWithLabel(t *testing.T, reg *prometheus.Registry, name, labelName, labelValue string) float64 { + t.Helper() + fams, err := reg.Gather() + if err != nil { + t.Fatalf("Gather: %v", err) + } + var sum float64 + for _, mf := range fams { + if mf.GetName() != name { + continue + } + for _, m := range mf.GetMetric() { + for _, lp := range m.GetLabel() { + if lp.GetName() == labelName && lp.GetValue() == labelValue { + sum += m.GetCounter().GetValue() + break + } + } + } + } + return sum +} + +// TestRoutingEmbeddingCountersAbsentWithoutStamp: responses without RoutingDebug +// (no routing embed ran) must not touch the routing counters. +func TestRoutingEmbeddingCountersAbsentWithoutStamp(t *testing.T) { + p := newTestPlugin(t) + + resp := &schemas.BifrostResponse{ChatResponse: &schemas.BifrostChatResponse{ + Usage: &schemas.BifrostLLMUsage{PromptTokens: 11, CompletionTokens: 7, TotalTokens: 18}, + }} + resp.PopulateExtraFields(schemas.ChatCompletionRequest, "openai", "test-model", "test-model") + + ctx := newHookContext(schemas.ChatCompletionRequest) + if _, _, err := p.PostLLMHook(ctx, resp, nil); err != nil { + t.Fatalf("PostLLMHook: %v", err) + } + + // Token counters record after the routing check in the same goroutine, so + // once they land we know the routing check already ran without recording. + waitForCounter(t, p.registry, "bifrost_input_tokens_total", 11) + if got := counterTotal(t, p.registry, "bifrost_routing_embedding_requests_total"); got != 0 { + t.Fatalf("routing requests counter = %v, want 0", got) + } +} + func TestMetricsEnabledGating(t *testing.T) { cases := []struct { name string diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index 2de6b1c6072..ed526f30d39 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -1984,6 +1984,9 @@ func (s *BifrostHTTPServer) ReloadPlugin(ctx context.Context, name string, path if governanceEmbeddingPlugin, ok := plugin.(governance.EmbeddingExecutorSetter); ok { governanceEmbeddingPlugin.SetEmbeddingRequestExecutor(s.Client.EmbeddingRequest) } + if governanceWarmupObserverPlugin, ok := plugin.(governance.WarmupEmbedUsageObserverSetter); ok { + governanceWarmupObserverPlugin.SetWarmupEmbedUsageObserver(s.observeWarmupRoutingEmbedding) + } return s.SyncLoadedPlugin(ctx, name, plugin, placement, order) } @@ -2243,6 +2246,19 @@ func (s *BifrostHTTPServer) RegisterAPIRoutes(ctx context.Context, callbacks Ser return nil } +// observeWarmupRoutingEmbedding forwards semantic-routing warmup embedding +// usage from the governance plugin to the telemetry plugin's routing overhead +// counters. The telemetry plugin is resolved per call (warmup is rare — boot +// and config changes only) so a reloaded telemetry instance, with its fresh +// registry, is picked up without re-wiring governance. +func (s *BifrostHTTPServer) observeWarmupRoutingEmbedding(provider, model string, inputTokens int) { + plugin, err := lib.FindPluginAs[*telemetry.PrometheusPlugin](s.Config, telemetry.PluginName) + if err != nil || plugin == nil { + return + } + plugin.ObserveWarmupRoutingEmbedding(provider, model, inputTokens) +} + // RegisterUIRoutes registers the UI handler with the specified router func (s *BifrostHTTPServer) RegisterUIRoutes(middlewares ...schemas.BifrostHTTPMiddleware) { // WARNING: This UI handler needs to be registered after all the other handlers @@ -2604,6 +2620,13 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { if err == nil && governanceEmbeddingPlugin != nil { governanceEmbeddingPlugin.SetEmbeddingRequestExecutor(s.Client.EmbeddingRequest) } + // Route semantic-routing warmup embedding usage to the telemetry counters. + // Warmup embeds have no request/response, so they cannot ride the + // RoutingDebug stamp path that per-request classification embeds use. + governanceWarmupObserverPlugin, err := lib.FindPluginAs[governance.WarmupEmbedUsageObserverSetter](s.Config, s.getGovernancePluginName()) + if err == nil && governanceWarmupObserverPlugin != nil { + governanceWarmupObserverPlugin.SetWarmupEmbedUsageObserver(s.observeWarmupRoutingEmbedding) + } // Initialize Sidekiq runner for background jobs if s.Config != nil && s.Config.ConfigStore != nil {