Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions core/bifrost.go
Original file line number Diff line number Diff line change
Expand Up @@ -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, "")
Expand Down Expand Up @@ -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, "")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -6191,19 +6220,23 @@ 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) {
drainAndAttachPluginLogs(ctx)
}
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
}
Expand Down Expand Up @@ -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)
}
Expand All @@ -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 {
Expand All @@ -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
Expand Down
35 changes: 35 additions & 0 deletions core/schemas/account.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
155 changes: 150 additions & 5 deletions core/schemas/bifrost.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Comment thread
greptile-apps[bot] marked this conversation as resolved.

// 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.
Expand Down Expand Up @@ -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
Comment thread
Pratham-Mishra04 marked this conversation as resolved.
// 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"`
Expand All @@ -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"`
Expand Down Expand Up @@ -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"`
Expand Down
1 change: 1 addition & 0 deletions framework/streaming/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Data *AccumulatedData
RawRequest *interface{}
}
Expand Down
Loading