From bdb85c822e5072e5c3effd3e1bb4777e1d31dba2 Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Wed, 3 Jun 2026 16:08:45 +0530 Subject: [PATCH] feat: adds routinginfo in respose/error extra fields --- core/bifrost.go | 36 ++++++++ core/schemas/account.go | 35 ++++++++ core/schemas/bifrost.go | 155 +++++++++++++++++++++++++++++++++-- framework/streaming/types.go | 1 + 4 files changed, 222 insertions(+), 5 deletions(-) diff --git a/core/bifrost.go b/core/bifrost.go index d19bc21834..5a6956ea8c 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -4759,6 +4759,11 @@ func (bifrost *Bifrost) handleRequest(ctx *schemas.BifrostContext, req *schemas. // Try the fallback provider result, fallbackErr := bifrost.tryRequest(ctx, fallbackReq) + // Layer on Primary/IsFallback — the per-attempt code populates only + // attempt-level RoutingInfo (Provider/Model/Key/ResolvedKeyAlias); + // fallback-relative signals belong to the orchestrator scope. + result.SetFallbackRoutingInfo(provider, model) + fallbackErr.SetFallbackRoutingInfo(provider, model) if fallbackErr == nil { bifrost.logger.Debug(fmt.Sprintf("successfully used fallback provider %s with model %s", fallback.Provider, fallback.Model)) tracer.EndSpan(handle, schemas.SpanStatusOk, "") @@ -4864,6 +4869,12 @@ func (bifrost *Bifrost) handleStreamRequest(ctx *schemas.BifrostContext, req *sc // Try the fallback provider result, fallbackErr := bifrost.tryStreamRequest(ctx, fallbackReq) + // Layer on Primary/IsFallback on errors. For the success case the + // result is a chan of stream chunks emitted asynchronously — those + // chunks already carry per-attempt RoutingInfo populated upstream, + // but Primary/IsFallback aren't reachable from here without wrapping + // the channel. See SetFallbackRoutingInfo doc. + fallbackErr.SetFallbackRoutingInfo(provider, model) if fallbackErr == nil { bifrost.logger.Debug(fmt.Sprintf("successfully used fallback provider %s with model %s", fallback.Provider, fallback.Model)) tracer.EndSpan(handle, schemas.SpanStatusOk, "") @@ -6151,6 +6162,20 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas // line 5653). Streaming postHookRunner must NOT capture this var by reference — it // snapshots its own attemptResolvedModel inside the per-attempt closure. var resolvedModel string + // attemptRoutingInfo holds the LAST attempt's RoutingInfo. Same single-writer/ + // single-reader contract as resolvedModel — assigned inside the per-attempt + // closure, read after retries finish by the post-retry populate below. + // Streaming postHookRunner must NOT capture by reference — it snapshots its + // own copy inside the per-attempt closure. + // Pre-seeded with the provider/model the orchestrator already knows so that + // retries that fail before the per-attempt closure ever runs (e.g. key + // selection error) still produce a populated RoutingInfo on the error — + // otherwise the post-retry populate at line ~6221 would clobber RoutingInfo + // to a zero value, leaving new consumers without provider/model context. + attemptRoutingInfo := schemas.RoutingInfo{ + Provider: provider.GetProviderKey(), + Model: originalModelRequested, + } // lastAttemptFinalizer captures the LAST attempt's postHookSpanFinalizer for the // worker-level error fallback below. Single-threaded write (assigned by the retry // loop's per-attempt closure) and single-threaded read (after retries finish), so @@ -6178,6 +6203,10 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas // Snapshot per-attempt so postHookRunner doesn't observe a later retry's // alias while this attempt's provider goroutine is still emitting chunks. attemptResolvedModel := resolvedModel + attemptRoutingInfo = schemas.BuildRoutingInfo(req.Context, provider.GetProviderKey(), originalModelRequested, k) + // Per-attempt snapshot for the async postHookRunner closure (it must + // not capture the outer var by reference — a later retry would race). + perAttemptRoutingInfo := attemptRoutingInfo // Snapshot RequestType before the closure. After tryStreamRequest receives // the stream channel it releases the *ChannelMessage back to the pool; // a concurrent request can then reuse it and overwrite RequestType. @@ -6191,9 +6220,11 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas // reference would let a later retry's alias bleed into this attempt's chunks. if result != nil { result.PopulateExtraFields(attemptRequestType, provider.GetProviderKey(), originalModelRequested, attemptResolvedModel) + result.PopulateRoutingInfo(perAttemptRoutingInfo) } if err != nil { err.PopulateExtraFields(attemptRequestType, provider.GetProviderKey(), originalModelRequested, attemptResolvedModel) + err.PopulateRoutingInfo(perAttemptRoutingInfo) } resp, bifrostErr := pipeline.RunPostLLMHooks(ctx, result, err, len(*bifrost.llmPlugins.Load())) if IsFinalChunk(ctx) { @@ -6201,9 +6232,11 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas } if bifrostErr != nil { bifrostErr.PopulateExtraFields(attemptRequestType, provider.GetProviderKey(), originalModelRequested, attemptResolvedModel) + bifrostErr.PopulateRoutingInfo(perAttemptRoutingInfo) return nil, bifrostErr } else if resp != nil { resp.PopulateExtraFields(attemptRequestType, provider.GetProviderKey(), originalModelRequested, attemptResolvedModel) + resp.PopulateRoutingInfo(perAttemptRoutingInfo) } return resp, nil } @@ -6240,6 +6273,7 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas req.Context.SetValue(schemas.BifrostContextKeyResolvedAlias, nil) } req.SetModel(resolvedModel) + attemptRoutingInfo = schemas.BuildRoutingInfo(req.Context, provider.GetProviderKey(), originalModelRequested, k) return bifrost.handleProviderRequest(provider, config, req, k, keys) }, keyProvider, req.RequestType, provider.GetProviderKey(), model, &req.BifrostRequest, bifrost.logger) } @@ -6262,6 +6296,7 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas if bifrostError != nil { bifrostError.PopulateExtraFields(req.RequestType, provider.GetProviderKey(), originalModelRequested, resolvedModel) + bifrostError.PopulateRoutingInfo(attemptRoutingInfo) // Send error with context awareness to prevent deadlock select { @@ -6277,6 +6312,7 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas } else { if result != nil { result.PopulateExtraFields(req.RequestType, provider.GetProviderKey(), originalModelRequested, resolvedModel) + result.PopulateRoutingInfo(attemptRoutingInfo) } if IsStreamRequestType(req.RequestType) { // Send stream with context awareness to prevent deadlock diff --git a/core/schemas/account.go b/core/schemas/account.go index 8e37cec1b1..6a9d159170 100644 --- a/core/schemas/account.go +++ b/core/schemas/account.go @@ -482,6 +482,41 @@ func IsVeoModelFamily(ctx *BifrostContext, model string) bool { return ResolveFamily(ctx, model) == ModelFamilyVeo } +// BuildRoutingInfo constructs a RoutingInfo for the current attempt from this +// attempt's chosen provider/model/key and the resolved alias stashed in ctx. +// +// Populates only the per-attempt fields (Provider, Model, Key, +// ResolvedKeyAlias). IsFallback and PrimaryProvider/PrimaryModel are layered +// on later by the orchestrator (handleRequest / handleStreamRequest) via +// SetFallbackRoutingInfo on the final response/error, since those signals +// belong to the orchestrator scope rather than the per-attempt one. +// +// ResolvedKeyAlias.ModelFamily reflects the family explicitly configured on +// the alias (nil when the admin didn't set one) — not the substring-resolved +// family used for routing. +func BuildRoutingInfo(ctx *BifrostContext, attemptProvider ModelProvider, attemptModel string, attemptKey Key) RoutingInfo { + info := RoutingInfo{ + Provider: attemptProvider, + Model: attemptModel, + Key: attemptKey.Name, + } + if ra := GetResolvedAlias(ctx); ra != nil && ra.Config != nil { + rka := &ResolvedKeyAlias{ + ModelID: ra.Config.ModelID, + } + if ra.Config.ModelName != nil { + mn := *ra.Config.ModelName + rka.ModelName = &mn + } + if ra.Config.ModelFamily != nil { + f := *ra.Config.ModelFamily + rka.ModelFamily = &f + } + info.ResolvedKeyAlias = rka + } + return info +} + // ResolveConfig returns the AliasConfig for the given user-facing model name, // or nil if no alias matches. Case-insensitive fallback matches Resolve. func (ka KeyAliases) ResolveConfig(model string) *AliasConfig { diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 15f3793b2d..8f6d9f2721 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -1095,6 +1095,107 @@ func (r *BifrostResponse) GetExtraFields() *BifrostResponseExtraFields { return &BifrostResponseExtraFields{} } +// syncDeprecatedFromRoutingInfo backfills the deprecated Provider / +// OriginalModelRequested / ResolvedModelUsed triplet on an ExtraFields-like +// target from a finalized RoutingInfo, applying the rules documented on each +// deprecated field. Centralized so PopulateRoutingInfo and +// SetFallbackRoutingInfo cannot drift apart. +func syncDeprecatedFromRoutingInfo(info RoutingInfo, provider *ModelProvider, originalModelRequested, resolvedModelUsed *string) { + if provider != nil && info.Provider != "" { + *provider = info.Provider + } + // OriginalModelRequested: collapses to the caller-sent model. On a fallback + // attempt that's the primary's model (the user never asked for the fallback's); + // otherwise it's this attempt's model. + if originalModelRequested != nil { + if info.IsFallback && info.PrimaryModel != nil && *info.PrimaryModel != "" { + *originalModelRequested = *info.PrimaryModel + } else if info.Model != "" { + *originalModelRequested = info.Model + } + } + // ResolvedModelUsed: the wire model. Alias's ModelID when an alias matched, + // otherwise the attempt's Model. + if resolvedModelUsed != nil { + if info.ResolvedKeyAlias != nil && info.ResolvedKeyAlias.ModelID != "" { + *resolvedModelUsed = info.ResolvedKeyAlias.ModelID + } else if info.Model != "" { + *resolvedModelUsed = info.Model + } + } +} + +// PopulateRoutingInfo sets ExtraFields.RoutingInfo on the active sub-response +// and keeps the deprecated Provider/OriginalModelRequested/ResolvedModelUsed +// triplet in sync per their documented derivation rules. +// Core always calls this both before and after RunPostLLMHooks so any plugin +// modifications are no-ops — tampering with RoutingInfo inside plugins is +// discouraged. +func (r *BifrostResponse) PopulateRoutingInfo(info RoutingInfo) { + if r == nil { + return + } + if ef := r.GetExtraFields(); ef != nil { + ef.RoutingInfo = info + syncDeprecatedFromRoutingInfo(info, &ef.Provider, &ef.OriginalModelRequested, &ef.ResolvedModelUsed) + } +} + +// PopulateRoutingInfo sets ExtraFields.RoutingInfo on the error and syncs the +// deprecated triplet. Core calls this both before and after RunPostLLMHooks +// alongside PopulateExtraFields. +func (e *BifrostError) PopulateRoutingInfo(info RoutingInfo) { + if e == nil { + return + } + e.ExtraFields.RoutingInfo = info + syncDeprecatedFromRoutingInfo(info, &e.ExtraFields.Provider, &e.ExtraFields.OriginalModelRequested, &e.ExtraFields.ResolvedModelUsed) +} + +// SetFallbackRoutingInfo marks the active sub-response's RoutingInfo as a +// fallback attempt and records the primary attempt's provider/model. Also +// re-syncs the deprecated OriginalModelRequested to the primary model per +// its documented derivation rule. +// Called by the orchestrator (handleRequest) on each fallback attempt's +// result/error — the per-attempt code never sets these fields itself. +func (r *BifrostResponse) SetFallbackRoutingInfo(primaryProvider ModelProvider, primaryModel string) { + if r == nil { + return + } + ef := r.GetExtraFields() + if ef == nil { + return + } + ef.RoutingInfo.IsFallback = true + if primaryProvider != "" { + p := primaryProvider + ef.RoutingInfo.PrimaryProvider = &p + } + if primaryModel != "" { + m := primaryModel + ef.RoutingInfo.PrimaryModel = &m + } + syncDeprecatedFromRoutingInfo(ef.RoutingInfo, &ef.Provider, &ef.OriginalModelRequested, &ef.ResolvedModelUsed) +} + +// SetFallbackRoutingInfo is the BifrostError counterpart — see the +// BifrostResponse method for semantics. +func (e *BifrostError) SetFallbackRoutingInfo(primaryProvider ModelProvider, primaryModel string) { + if e == nil { + return + } + e.ExtraFields.RoutingInfo.IsFallback = true + if primaryProvider != "" { + p := primaryProvider + e.ExtraFields.RoutingInfo.PrimaryProvider = &p + } + if primaryModel != "" { + m := primaryModel + e.ExtraFields.RoutingInfo.PrimaryModel = &m + } + syncDeprecatedFromRoutingInfo(e.ExtraFields.RoutingInfo, &e.ExtraFields.Provider, &e.ExtraFields.OriginalModelRequested, &e.ExtraFields.ResolvedModelUsed) +} + // PopulateExtraFields sets RequestType, Provider, OriginalModelRequested, and ResolvedModelUsed on the // active sub-response. Core always calls this both before and after RunPostLLMHooks, so any plugin // modifications to these 4 fields are no-ops — tampering with them inside plugins is discouraged. @@ -1448,9 +1549,20 @@ func (r *BifrostMCPResponse) PopulateExtraFields(mcpRequestType MCPRequestType, // BifrostResponseExtraFields contains additional fields in a response. type BifrostResponseExtraFields struct { RequestType RequestType `json:"request_type"` - Provider ModelProvider `json:"provider,omitempty"` - OriginalModelRequested string `json:"original_model_requested,omitempty"` // the model alias the caller sent in the request - ResolvedModelUsed string `json:"resolved_model_used,omitempty"` // the actual provider API identifier used (equals OriginalModelRequested when no alias mapping exists) + RoutingInfo RoutingInfo `json:"routing_info"` + // Deprecated: use RoutingInfo.Provider. Still populated for backward + // compatibility; new consumers should read from RoutingInfo. + Provider ModelProvider `json:"provider,omitempty"` + // Deprecated: use RoutingInfo.PrimaryModel when RoutingInfo.IsFallback + // is true, otherwise RoutingInfo.Model — both branches collapse to the + // model string the caller sent in the request. Still populated for + // backward compatibility; new consumers should read from RoutingInfo. + OriginalModelRequested string `json:"original_model_requested,omitempty"` + // Deprecated: use RoutingInfo.ResolvedKeyAlias.ModelID when an alias + // matched (i.e. RoutingInfo.ResolvedKeyAlias != nil), otherwise + // RoutingInfo.Model. Still populated for backward compatibility; new + // consumers should read from RoutingInfo. + ResolvedModelUsed string `json:"resolved_model_used,omitempty"` Latency int64 `json:"latency"` // in milliseconds (for streaming responses this will be each chunk latency, and the last chunk latency will be the total latency) ChunkIndex int `json:"chunk_index"` // used for streaming responses to identify the chunk index, will be 0 for non-streaming responses RawRequest interface{} `json:"raw_request,omitempty"` @@ -1463,6 +1575,28 @@ type BifrostResponseExtraFields struct { PassthroughPath string `json:"passthrough_path,omitempty"` // Stripped provider path for passthrough requests, e.g. "/v1/chat/completions" } +type RoutingInfo struct { + // What actually handled this attempt + Provider ModelProvider `json:"provider,omitempty"` + Model string `json:"model,omitempty"` // model name passed to this attempt's key + Key string `json:"key,omitempty"` // KeyName of the key used + + // Populated only when Model matched an entry in this key's Aliases map + ResolvedKeyAlias *ResolvedKeyAlias `json:"resolved_key_alias,omitempty"` + + IsFallback bool `json:"is_fallback,omitempty"` + + // What the caller asked for, before any fallback resolution (populated only when fallback resolution occurred) + PrimaryProvider *ModelProvider `json:"primary_provider,omitempty"` + PrimaryModel *string `json:"primary_model,omitempty"` +} + +type ResolvedKeyAlias struct { + ModelID string `json:"model_id"` // wire model identifier actually sent to the provider + ModelName *string `json:"model_name,omitempty"` // canonical name (used for pricing/logs) + ModelFamily *ModelFamily `json:"model_family,omitempty"` // resolved family for routing +} + type BifrostMCPResponseExtraFields struct { MCPRequestType MCPRequestType `json:"mcp_request_type"` // request type this response corresponds to — lets PostMCPHook discriminate ping/list_tools from tool execute on success too ClientName string `json:"client_name"` @@ -1698,8 +1832,19 @@ func (e *ErrorField) UnmarshalJSON(data []byte) error { // BifrostErrorExtraFields contains additional fields in an error response. type BifrostErrorExtraFields struct { - Provider ModelProvider `json:"provider,omitempty"` - OriginalModelRequested string `json:"original_model_requested,omitempty"` + RoutingInfo RoutingInfo `json:"routing_info"` + // Deprecated: use RoutingInfo.Provider. Still populated for backward + // compatibility; new consumers should read from RoutingInfo. + Provider ModelProvider `json:"provider,omitempty"` + // Deprecated: use RoutingInfo.PrimaryModel when RoutingInfo.IsFallback + // is true, otherwise RoutingInfo.Model — both branches collapse to the + // model string the caller sent in the request. Still populated for + // backward compatibility; new consumers should read from RoutingInfo. + OriginalModelRequested string `json:"original_model_requested,omitempty"` + // Deprecated: use RoutingInfo.ResolvedKeyAlias.ModelID when an alias + // matched (i.e. RoutingInfo.ResolvedKeyAlias != nil), otherwise + // RoutingInfo.Model. Still populated for backward compatibility; new + // consumers should read from RoutingInfo. ResolvedModelUsed string `json:"resolved_model_used,omitempty"` RequestType RequestType `json:"request_type,omitempty"` MCPRequestType MCPRequestType `json:"mcp_request_type,omitempty"` diff --git a/framework/streaming/types.go b/framework/streaming/types.go index c673963331..9dcb3b1708 100644 --- a/framework/streaming/types.go +++ b/framework/streaming/types.go @@ -244,6 +244,7 @@ type ProcessedStreamResponse struct { Provider schemas.ModelProvider RequestedModel string // original model requested by the caller ResolvedModel string // actual model used by the provider (equals RequestedModel when no alias mapping exists) + RoutingInfo schemas.RoutingInfo Data *AccumulatedData RawRequest *interface{} }