From 735f72b945a8837dcced8f25611143b407c22af0 Mon Sep 17 00:00:00 2001 From: tejas ghatte Date: Mon, 8 Jun 2026 17:08:43 +0530 Subject: [PATCH 001/108] fix: private network toggle in custom provider form --- .../dialogs/addNewCustomProviderSheet.tsx | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/ui/app/workspace/providers/dialogs/addNewCustomProviderSheet.tsx b/ui/app/workspace/providers/dialogs/addNewCustomProviderSheet.tsx index f3728bc32a..a9a5668e0a 100644 --- a/ui/app/workspace/providers/dialogs/addNewCustomProviderSheet.tsx +++ b/ui/app/workspace/providers/dialogs/addNewCustomProviderSheet.tsx @@ -23,6 +23,7 @@ const formSchema = z.object({ allowed_requests: allowedRequestsSchema, request_path_overrides: z.record(z.string(), z.string().optional()).optional(), is_key_less: z.boolean().optional(), + allow_private_network: z.boolean().optional(), }); type FormData = z.infer; @@ -79,6 +80,7 @@ export function AddCustomProviderSheetContent({ show = true, onClose, onSave }: }, request_path_overrides: undefined, is_key_less: false, + allow_private_network: false, }, }); @@ -99,6 +101,7 @@ export function AddCustomProviderSheetContent({ show = true, onClose, onSave }: }, network_config: { base_url: data.base_url, + allow_private_network: data.allow_private_network ?? false, default_request_timeout_in_seconds: 30, max_retries: 0, retry_backoff_initial: 500, @@ -194,6 +197,32 @@ export function AddCustomProviderSheetContent({ show = true, onClose, onSave }: )} /> + ( + +
+
+ +

+ Allow connecting to private network IPs (e.g. 192.168.x.x, 10.x.x.x). Link-local addresses remain blocked. +

+
+ +
+ + )} + /> {!isKeyLessDisabled && ( Date: Tue, 9 Jun 2026 15:42:26 +0530 Subject: [PATCH 002/108] feat: add `PreRequestHook` to `LLMPlugin` interface for once-per-request provider/model routing (#4175) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Introduces a new `PreRequestHook` phase to the `LLMPlugin` interface. This hook runs exactly once per top-level request — after `HTTPTransportPreHook` and before `PreLLMHook` — and is the canonical place for plugins to resolve provider, model, and fallback routing decisions. Previously, routing logic had to be shoehorned into `PreLLMHook`, which runs on every fallback attempt and whose mutations have incidental cross-fallback visibility. `PreRequestHook` mutations are committed to the shared `*BifrostRequest` before any fan-out and are observed by every subsequent plugin, every `PreLLMHook` invocation, the provider call, and every fallback. As part of this change, the `filterProvidersByContext` helper (used in `ListAllModels`) is removed, and request validation is moved to after `PreRequestHook` runs so that plugins have the opportunity to populate provider/model before the empty-field check fires. Error messages for missing provider/model are updated to reflect that auto-resolution was attempted. ## Changes - Added `PreRequestHook(ctx *BifrostContext, req *BifrostRequest) error` to the `LLMPlugin` interface with non-blocking error semantics (logged as warning, pipeline continues). - Added `RunPreRequestHooks` to `PluginPipeline`, executing the hook in registration order once per request with tracing and plugin-scope isolation. - Added `RunPreRequestHooks` as a public method on `Bifrost` for callers (e.g., realtime WebSocket handlers) that bypass the normal inference path. - Moved `validateRequest` to after `PreRequestHook` execution in both `handleRequest` and `handleStreamRequest`, renamed to `validateRequestAfterPreRequestHooks` with updated error messages. - Added primary-provider error logging to `handleStreamRequest` to match `handleRequest` behavior. - Removed `filterProvidersByContext` and its tests from `ListAllModels`. - Updated `DynamicPlugin` (shared-object loader) to optionally load `PreRequestHook` from `.so` plugins; legacy plugins without the export get a no-op passthrough, preserving backward compatibility. - Updated `AsLLMPlugin` to recognize `preRequestHook` as sufficient to qualify a `DynamicPlugin` as an `LLMPlugin`. - Added no-op `PreRequestHook` implementations to all existing plugins (`compat`, `governance`, `jsonparser`, `logging`, `maxim`, `mocker`, `prompts`, `semanticcache`, `telemetry`) and all example/test plugins to satisfy the updated interface. - Updated plugin execution-order documentation in `plugin.go` to describe per-request vs. per-attempt semantics. ## Type of change - [ ] Bug fix - [x] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [ ] Providers/Integrations - [x] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./... ``` Validate that: - A plugin implementing `PreRequestHook` can mutate `req.Provider` and `req.Model` before the provider call, and those mutations are visible to subsequent plugins and fallback attempts. - A plugin that returns a non-nil error from `PreRequestHook` does not abort the request; the pipeline continues to the next plugin and a warning is logged. - Existing plugins with no-op `PreRequestHook` implementations behave identically to before. - Legacy `.so` plugins that do not export `PreRequestHook` load and function correctly with the no-op passthrough. - Requests with no provider set (and no plugin resolving one) return the updated error message: `"could not auto resolve a provider for the request, please specify a provider explicitly"`. ## Breaking changes - [x] Yes - [ ] No The `LLMPlugin` interface gains a new required method `PreRequestHook`. Any external plugin implementing `LLMPlugin` must add a `PreRequestHook` method. Plugins that do not participate in routing should return `nil`. Shared-object (`.so`) plugins are exempt — the loader treats `PreRequestHook` as optional and provides a no-op default. ## Security considerations `PreRequestHook` runs with `BlockRestrictedWrites` active on the context (same as `RunLLMPreHooks`), preventing plugins from writing to restricted context keys during the hook. Plugins cannot abort or gate requests via error return from this hook; authorization and content-policy enforcement must remain in `HTTPTransportPreHook` or via a short-circuit in `PreLLMHook`. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [x] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable ## Summary by CodeRabbit ## Release Notes * **New Features** * Added `PreRequestHook` plugin phase enabling plugins to perform per-request routing decisions before provider and model validation. * **Refactor** * Request validation now occurs after plugin hooks execute, allowing automatic resolution of routing parameters. --- core/bifrost.go | 180 +++++++++++++----- core/bifrost_test.go | 45 ----- core/schemas/plugin.go | 38 +++- core/utils.go | 8 +- examples/plugins/hello-world/main.go | 4 + examples/plugins/llm-only/main.go | 5 + examples/plugins/multi-interface/main.go | 5 + framework/plugins/main.go | 2 +- framework/plugins/soloader.go | 9 + framework/plugins/soplugin.go | 18 +- framework/tracing/tracer_test.go | 3 + plugins/compat/main.go | 7 +- plugins/governance/main.go | 5 + plugins/jsonparser/main.go | 5 + plugins/logging/main.go | 5 + plugins/maxim/main.go | 5 + plugins/mocker/main.go | 31 +-- plugins/prompts/main.go | 5 + plugins/semanticcache/main.go | 5 + .../semanticcache/plugin_no_mutation_test.go | 4 + plugins/telemetry/main.go | 5 + .../handlers/realtime_client_secrets_test.go | 4 + transports/bifrost-http/lib/config_test.go | 4 + 23 files changed, 286 insertions(+), 116 deletions(-) diff --git a/core/bifrost.go b/core/bifrost.go index 06e6d6beae..f63bf3ebec 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -458,8 +458,6 @@ func (bifrost *Bifrost) ListAllModels(ctx *schemas.BifrostContext, req *schemas. }, } } - providerKeys = filterProvidersByContext(ctx, providerKeys) - startTime := time.Now() // Result structure for collecting provider responses @@ -604,35 +602,6 @@ func (bifrost *Bifrost) ListAllModels(ctx *schemas.BifrostContext, req *schemas. return response, nil } -func filterProvidersByContext(ctx *schemas.BifrostContext, providerKeys []schemas.ModelProvider) []schemas.ModelProvider { - if ctx == nil { - return providerKeys - } - - rawAvailableProviders := ctx.Value(schemas.BifrostContextKeyAvailableProviders) - if rawAvailableProviders == nil { - return providerKeys - } - - availableProviders, ok := rawAvailableProviders.([]schemas.ModelProvider) - if !ok { - return []schemas.ModelProvider{} - } - - if len(availableProviders) == 0 || len(providerKeys) == 0 { - return []schemas.ModelProvider{} - } - - filteredProviders := make([]schemas.ModelProvider, 0, len(providerKeys)) - for _, providerKey := range providerKeys { - if slices.Contains(availableProviders, providerKey) { - filteredProviders = append(filteredProviders, providerKey) - } - } - - return filteredProviders -} - // TextCompletionRequest sends a text completion request to the specified provider. func (bifrost *Bifrost) TextCompletionRequest(ctx *schemas.BifrostContext, req *schemas.BifrostTextCompletionRequest) (*schemas.BifrostTextCompletionResponse, *schemas.BifrostError) { if req == nil { @@ -4197,6 +4166,33 @@ type RealtimeTurnHooks struct { Cleanup func() } +// RunPreRequestHooks acquires a plugin pipeline and runs PreRequestHook on each LLM plugin +// for callers that do not flow through handleRequest/handleStreamRequest — primarily realtime +// WebSocket upgrades, where the upgrade itself is the routing decision (once per WS connection) +// but the per-turn pipeline handles PreLLMHook/PostLLMHook separately. +// +// Mutations to req.Provider/req.Model/req.Fallbacks made by PreRequestHook plugins are committed +// to the shared *BifrostRequest. Plugin errors are non-blocking — they are logged as warnings +// and the pipeline continues to the next plugin (same semantics as RunLLMPreHooks). Callers +// should validate req.Provider after this returns if a provider is required. +func (bifrost *Bifrost) RunPreRequestHooks(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) { + if ctx == nil { + ctx = bifrost.ctx + } + + if _, ok := ctx.Value(schemas.BifrostContextKeyRequestID).(string); !ok { + ctx.SetValue(schemas.BifrostContextKeyRequestID, uuid.New().String()) + } + + pipeline := bifrost.getPluginPipeline() + defer bifrost.releasePluginPipeline(pipeline) + pipeline.RunPreRequestHooks(ctx, req) + // This path has no downstream post-hook cleanup, so drain any plugin logs + // emitted by PreRequestHook here to avoid them bleeding into a later request + // on a reused/long-lived context (e.g. realtime WS connections). + flushPluginLogs(ctx) +} + // RunStreamPreHooks acquires a plugin pipeline, sets up tracing context, runs PreLLMHooks, // and returns a PostHookRunner for per-chunk post-processing. // Used by WebSocket handlers that bypass the normal inference path but still need plugin hooks. @@ -4636,18 +4632,12 @@ func (bifrost *Bifrost) shouldContinueWithFallbacks(fallback schemas.Fallback, f func (bifrost *Bifrost) handleRequest(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostResponse, *schemas.BifrostError) { defer bifrost.releaseBifrostRequest(req) provider, model, fallbacks := req.GetRequestFields() - if err := validateRequest(req); err != nil { - err.PopulateExtraFields(req.RequestType, provider, model, model) - return nil, err - } // Handle nil context early to prevent blocking if ctx == nil { ctx = bifrost.ctx } - bifrost.logger.Debug(fmt.Sprintf("primary provider %s with model %s and %d fallbacks", provider, model, len(fallbacks))) - // Try the primary provider first ctx.SetValue(schemas.BifrostContextKeyFallbackIndex, 0) // Ensure request ID is set in context before PreHooks @@ -4655,6 +4645,27 @@ func (bifrost *Bifrost) handleRequest(ctx *schemas.BifrostContext, req *schemas. requestID := uuid.New().String() ctx.SetValue(schemas.BifrostContextKeyRequestID, requestID) } + + // PreRequestHook: once-per-request phase where plugins decide provider/model/fallbacks + // (and may mutate other request fields). Mutations commit to req and are observed by + // all downstream phases and fallbacks. Plugin errors are non-blocking (logged + skipped). + preReqPipeline := bifrost.getPluginPipeline() + preReqPipeline.RunPreRequestHooks(ctx, req) + bifrost.releasePluginPipeline(preReqPipeline) + // Re-read after PreRequestHook — provider/model/fallbacks may have changed. + provider, model, fallbacks = req.GetRequestFields() + // Empty provider/model after PreRequestHook means no plugin + // could pick a provider for this model — the caller's input is unresolvable. + if err := validateRequestAfterPreRequestHooks(req); err != nil { + // Returning before tryRequest skips the downstream log drain, so flush + // any PreRequestHook-emitted plugin logs here. + flushPluginLogs(ctx) + err.PopulateExtraFields(req.RequestType, provider, model, model) + return nil, err + } + + bifrost.logger.Debug(fmt.Sprintf("primary provider %s with model %s and %d fallbacks", provider, model, len(fallbacks))) + primaryResult, primaryErr := bifrost.tryRequest(ctx, req) if primaryErr != nil { if primaryErr.Error != nil { @@ -4727,15 +4738,8 @@ func (bifrost *Bifrost) handleRequest(ctx *schemas.BifrostContext, req *schemas. // It is the wrapper for all streaming public API methods. func (bifrost *Bifrost) handleStreamRequest(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { defer bifrost.releaseBifrostRequest(req) - provider, model, fallbacks := req.GetRequestFields() - if err := validateRequest(req); err != nil { - err.PopulateExtraFields(req.RequestType, provider, model, model) - err.StatusCode = schemas.Ptr(fasthttp.StatusBadRequest) - return nil, err - } - // Handle nil context early to prevent blocking if ctx == nil { ctx = bifrost.ctx @@ -4748,7 +4752,36 @@ func (bifrost *Bifrost) handleStreamRequest(ctx *schemas.BifrostContext, req *sc requestID := uuid.New().String() ctx.SetValue(schemas.BifrostContextKeyRequestID, requestID) } + + // PreRequestHook: once-per-request phase. See handleRequest for semantics. + preReqPipeline := bifrost.getPluginPipeline() + preReqPipeline.RunPreRequestHooks(ctx, req) + bifrost.releasePluginPipeline(preReqPipeline) + // Re-read after PreRequestHook — provider/model/fallbacks may have changed. + provider, model, fallbacks = req.GetRequestFields() + // Empty provider after PreRequestHook means no plugin + // could pick a provider for this model — the caller's input is unresolvable. + if err := validateRequestAfterPreRequestHooks(req); err != nil { + // Returning before tryStreamRequest skips the downstream log drain, so + // flush any PreRequestHook-emitted plugin logs here. + flushPluginLogs(ctx) + err.PopulateExtraFields(req.RequestType, provider, model, model) + return nil, err + } + + bifrost.logger.Debug(fmt.Sprintf("primary provider %s with model %s and %d fallbacks", provider, model, len(fallbacks))) + primaryResult, primaryErr := bifrost.tryStreamRequest(ctx, req) + if primaryErr != nil { + if primaryErr.Error != nil { + bifrost.logger.Debug(fmt.Sprintf("primary provider %s with model %s returned error: %s", provider, model, primaryErr.Error.Message)) + } else { + bifrost.logger.Debug(fmt.Sprintf("primary provider %s with model %s returned error: %v", provider, model, primaryErr)) + } + if len(fallbacks) > 0 { + bifrost.logger.Debug(fmt.Sprintf("check if we should try %d fallbacks", len(fallbacks))) + } + } // Check if we should proceed with fallbacks shouldTryFallbacks := bifrost.shouldTryFallbacks(req, primaryErr) @@ -6653,6 +6686,49 @@ func (p *PluginPipeline) RunLLMPreHooks(ctx *schemas.BifrostContext, req *schema return req, nil, p.executedPreHooks } +// RunPreRequestHooks executes PreRequestHook on each LLM plugin in registration order, once per +// top-level request. Plugins mutate req.Provider, req.Model, req.Fallbacks (and any other field +// they choose); mutations are committed to the shared *BifrostRequest and observed by every +// subsequent plugin, the provider call, and every fallback attempt. There is no short-circuit +// and errors are non-blocking — same semantics as RunLLMPreHooks: errors are logged as warnings +// and accumulated in p.preHookErrors, then the pipeline continues to the next plugin. The empty- +// provider validation in handleRequest/handleStreamRequest catches the case where no plugin +// successfully resolved a provider. +// +// Per-request semantics: unlike PreLLMHook (which runs again on every fallback), PreRequestHook +// runs exactly once at the top of handleRequest/handleStreamRequest, before any fan-out. +func (p *PluginPipeline) RunPreRequestHooks(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) { + // If the skip plugin pipeline flag is set, skip the plugin pipeline + if skipPluginPipeline, ok := ctx.Value(schemas.BifrostContextKeySkipPluginPipeline).(bool); ok && skipPluginPipeline { + return + } + ctx.BlockRestrictedWrites() + defer ctx.UnblockRestrictedWrites() + for _, plugin := range p.llmPlugins { + pluginName := plugin.GetName() + p.logger.Debug("running pre-request hook for plugin %s", pluginName) + spanCtx, handle := p.tracer.StartSpan(ctx, fmt.Sprintf("plugin.%s.prerequesthook", sanitizeSpanName(pluginName)), schemas.SpanKindPlugin) + if spanCtx != nil { + if spanID, ok := spanCtx.Value(schemas.BifrostContextKeySpanID).(string); ok { + ctx.SetValue(schemas.BifrostContextKeySpanID, spanID) + } + } + + pluginCtx := ctx.WithPluginScope(&pluginName) + err := plugin.PreRequestHook(pluginCtx, req) + pluginCtx.ReleasePluginScope() + + if err != nil { + p.tracer.SetAttribute(handle, "error", err.Error()) + p.tracer.EndSpan(handle, schemas.SpanStatusError, err.Error()) + p.preHookErrors = append(p.preHookErrors, err) + p.logger.Warn("error in PreRequestHook for plugin %s: %s", pluginName, err.Error()) + continue + } + p.tracer.EndSpan(handle, schemas.SpanStatusOk, "") + } +} + // RunPostLLMHooks executes PostHooks in reverse order for the plugins whose PreLLMHook ran. // Accepts the response and error, and allows plugins to transform either (e.g., recover from error, or invalidate a response). // Returns the final response and error after all hooks. If both are set, error takes precedence unless error is nil. @@ -7013,6 +7089,22 @@ func (p *PluginPipeline) resetPluginPipeline() { p.streamingMu.Unlock() } +// flushPluginLogs drains accumulated plugin logs from the BifrostContext and +// attaches them to the active trace when one exists. Unlike drainAndAttachPluginLogs, +// it always drains the buffer first, so logs emitted before any trace is established +// (e.g. by PreRequestHook) are not carried over to a later request on a reused context. +func flushPluginLogs(ctx *schemas.BifrostContext) { + logs := ctx.DrainPluginLogs() + if len(logs) == 0 { + return + } + tracer, traceID, err := GetTracerFromContext(ctx) + if err != nil || tracer == nil || traceID == "" { + return + } + tracer.AttachPluginLogs(traceID, logs) +} + // drainAndAttachPluginLogs drains accumulated plugin logs from the BifrostContext // and attaches them to the trace for later retrieval by observability plugins. func drainAndAttachPluginLogs(ctx *schemas.BifrostContext) { diff --git a/core/bifrost_test.go b/core/bifrost_test.go index 5c1300687f..d7893e3cde 100644 --- a/core/bifrost_test.go +++ b/core/bifrost_test.go @@ -789,51 +789,6 @@ func (t *countingTracer) CompleteAndFlushTrace(_ string) { t.flushed.Add(1) } -func TestFilterProvidersByContext(t *testing.T) { - providers := []schemas.ModelProvider{ - schemas.OpenAI, - schemas.Anthropic, - schemas.Mistral, - } - - t.Run("no context filter keeps all providers", func(t *testing.T) { - filtered := filterProvidersByContext(nil, providers) - if len(filtered) != len(providers) { - t.Fatalf("expected all providers, got %v", filtered) - } - }) - - t.Run("available providers restrict list models fanout", func(t *testing.T) { - ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - ctx.SetValue(schemas.BifrostContextKeyAvailableProviders, []schemas.ModelProvider{schemas.Anthropic}) - - filtered := filterProvidersByContext(ctx, providers) - if len(filtered) != 1 || filtered[0] != schemas.Anthropic { - t.Fatalf("expected only anthropic, got %v", filtered) - } - }) - - t.Run("empty available providers denies all providers", func(t *testing.T) { - ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - ctx.SetValue(schemas.BifrostContextKeyAvailableProviders, []schemas.ModelProvider{}) - - filtered := filterProvidersByContext(ctx, providers) - if len(filtered) != 0 { - t.Fatalf("expected no providers, got %v", filtered) - } - }) - - t.Run("malformed available providers fails closed", func(t *testing.T) { - ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - ctx.SetValue(schemas.BifrostContextKeyAvailableProviders, "openai") - - filtered := filterProvidersByContext(ctx, providers) - if len(filtered) != 0 { - t.Fatalf("expected no providers for malformed context value, got %v", filtered) - } - }) -} - func TestRunStreamPreHooks_FinalChunkFlushesTrace(t *testing.T) { ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) account := NewMockAccount() diff --git a/core/schemas/plugin.go b/core/schemas/plugin.go index 7ca4a0ec68..186af928b7 100644 --- a/core/schemas/plugin.go +++ b/core/schemas/plugin.go @@ -170,12 +170,22 @@ func ReleaseHTTPResponse(resp *HTTPResponse) { // PostHooks are executed in the reverse order of PreHooks. // // Execution order: -// 1. HTTPTransportPreHook (HTTP transport only, executed in registration order) -// 2. PreLLMHook (executed in registration order) -// 3. Provider call -// 4. PostLLMHook (executed in reverse order of PreHooks) -// 5. HTTPTransportPostHook (HTTP transport only, executed in reverse order) -// 5a. HTTPTransportStreamChunkHook (for streaming responses, called per-chunk in reverse order) +// 1. HTTPTransportPreHook (HTTP transport only, once per request, executed in registration order) +// 2. PreRequestHook (once per request, executed in registration order) +// 3. PreLLMHook (executed in registration order, runs again on each fallback attempt) +// 4. Provider call +// 5. PostLLMHook (executed in reverse order of PreHooks, runs on each fallback attempt) +// 6. HTTPTransportPostHook (HTTP transport only, once per request, executed in reverse order) +// 6a. HTTPTransportStreamChunkHook (for streaming responses, called per-chunk in reverse order) +// +// Per-request vs per-attempt phases: +// - HTTPTransportPreHook, PreRequestHook, HTTPTransportPostHook run ONCE per top-level request. +// - PreLLMHook, PostLLMHook run ONCE PER ATTEMPT: the primary provider call, plus once per +// fallback attempt. Mutations a PreLLMHook makes to the request only carry to later +// fallbacks where prepareFallbackRequest happens to share pointers (shallow copy) — +// visibility across fallbacks is incidental. PreRequestHook is the explicit phase whose +// mutations are committed to the request before any fan-out and are observed by every +// subsequent plugin, every PreLLMHook invocation, the provider call, and every fallback. // // Common use cases: rate limiting, caching, logging, monitoring, request transformation, governance. // @@ -257,6 +267,22 @@ type HTTPTransportPlugin interface { type LLMPlugin interface { BasePlugin + // PreRequestHook is called once per top-level request, after HTTPTransportPreHook and before + // PreLLMHook. It is the canonical phase for deciding which provider/model/fallbacks the + // request should be sent to. Plugins are free to mutate any field on req (Provider, Model, + // Fallbacks, Input, Params, Tools, ...) — unlike PreLLMHook, mutations made here are + // committed to the request and are observed by all subsequent plugins, the provider call, + // and every fallback attempt. + // + // Error semantics match PreLLMHook: a non-nil error is non-blocking — it is logged as a + // warning, the request continues, and the pipeline moves on to the next plugin. PreRequestHook + // CANNOT abort the request via error return. Plugins that need to gate or reject a request + // (e.g., authorization, content policy) must do so in HTTPTransportPreHook or via a + // short-circuit response in PreLLMHook — not by returning an error here. + // + // Plugins that don't participate in routing should return nil. + PreRequestHook(ctx *BifrostContext, req *BifrostRequest) error + PreLLMHook(ctx *BifrostContext, req *BifrostRequest) (*BifrostRequest, *LLMPluginShortCircuit, error) PostLLMHook(ctx *BifrostContext, resp *BifrostResponse, bifrostErr *BifrostError) (*BifrostResponse, *BifrostError, error) } diff --git a/core/utils.go b/core/utils.go index 3993b2dd11..167d4b890d 100644 --- a/core/utils.go +++ b/core/utils.go @@ -135,17 +135,17 @@ func calculateBackoff(attempt int, config *schemas.ProviderConfig) time.Duration return min(result, config.NetworkConfig.RetryBackoffMax) } -// validateRequest validates the given request. -func validateRequest(req *schemas.BifrostRequest) *schemas.BifrostError { +// validateRequestAfterPreRequestHooks validates the provider and model fields of the given request. +func validateRequestAfterPreRequestHooks(req *schemas.BifrostRequest) *schemas.BifrostError { if req == nil { return newBifrostErrorFromMsg("bifrost request cannot be nil") } provider, model, _ := req.GetRequestFields() if provider == "" { - return newBifrostErrorFromMsg("provider is required") + return newBifrostErrorFromMsg("could not auto resolve a provider for the request, please specify a provider explicitly") } if isModelRequired(req.RequestType) && model == "" { - return newBifrostErrorFromMsg("model is required") + return newBifrostErrorFromMsg("could not auto resolve a model for the request, please specify a model explicitly") } return nil } diff --git a/examples/plugins/hello-world/main.go b/examples/plugins/hello-world/main.go index 2f464e2a7d..a08fe013f0 100644 --- a/examples/plugins/hello-world/main.go +++ b/examples/plugins/hello-world/main.go @@ -57,6 +57,10 @@ func HTTPTransportStreamChunkHook(ctx *schemas.BifrostContext, req *schemas.HTTP return chunk, nil } +func PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + func PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { value1 := ctx.Value(transportPreHookKey) fmt.Println("value1:", value1) diff --git a/examples/plugins/llm-only/main.go b/examples/plugins/llm-only/main.go index 65d47c2020..90ec9da184 100644 --- a/examples/plugins/llm-only/main.go +++ b/examples/plugins/llm-only/main.go @@ -67,6 +67,11 @@ func GetName() string { return "llm-only" } +// PreRequestHook is the per-request routing phase. This example plugin doesn't route. +func PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + // PreLLMHook is called before the LLM provider is invoked // This example demonstrates request modification and logging func PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { diff --git a/examples/plugins/multi-interface/main.go b/examples/plugins/multi-interface/main.go index d593b778ef..e8b3a14292 100644 --- a/examples/plugins/multi-interface/main.go +++ b/examples/plugins/multi-interface/main.go @@ -155,6 +155,11 @@ func HTTPTransportPostHook(ctx *schemas.BifrostContext, req *schemas.HTTPRequest // LLMPlugin Interface // ============================================================================ +// PreRequestHook is the per-request routing phase. This example plugin doesn't route. +func PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + // PreLLMHook is called before the LLM provider is invoked func PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { if !pluginConfig.EnableLLMHooks { diff --git a/framework/plugins/main.go b/framework/plugins/main.go index b7bc20bb06..5ce955aad2 100644 --- a/framework/plugins/main.go +++ b/framework/plugins/main.go @@ -27,7 +27,7 @@ func AsLLMPlugin(plugin schemas.BasePlugin) schemas.LLMPlugin { // Check if it's a DynamicPlugin first if dp, ok := plugin.(*DynamicPlugin); ok { // Only return as LLMPlugin if it actually has LLM hooks - if dp.preLLMHook != nil || dp.postLLMHook != nil { + if dp.preRequestHook != nil || dp.preLLMHook != nil || dp.postLLMHook != nil { return dp } return nil diff --git a/framework/plugins/soloader.go b/framework/plugins/soloader.go index face5772a7..afbec9d6f2 100644 --- a/framework/plugins/soloader.go +++ b/framework/plugins/soloader.go @@ -94,6 +94,15 @@ func (l *SharedObjectPluginLoader) LoadPlugin(path string, config any) (schemas. } } + // Optional: PreRequestHook — new .so plugins built against LLMPlugin can export this + // to participate in routing. Legacy plugins predating PreRequestHook keep working; + // DynamicPlugin's default PreRequestHook is a no-op passthrough. + if sym, err := pluginObj.Lookup("PreRequestHook"); err == nil { + if dp.preRequestHook, ok = sym.(func(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) error); !ok { + return nil, fmt.Errorf("failed to cast PreRequestHook to expected signature") + } + } + // Optional: PreLLMHook (with backward compatibility for legacy PreHook) if sym, err := pluginObj.Lookup("PreLLMHook"); err == nil { if dp.preLLMHook, ok = sym.(func(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error)); !ok { diff --git a/framework/plugins/soplugin.go b/framework/plugins/soplugin.go index 196a1ceed2..8b73e521ae 100644 --- a/framework/plugins/soplugin.go +++ b/framework/plugins/soplugin.go @@ -27,8 +27,12 @@ type DynamicPlugin struct { httpTransportStreamChunkHook func(ctx *schemas.BifrostContext, req *schemas.HTTPRequest, stream *schemas.BifrostStreamChunk) (*schemas.BifrostStreamChunk, error) // LLMPlugin (optional) - preLLMHook func(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) - postLLMHook func(ctx *schemas.BifrostContext, resp *schemas.BifrostResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError, error) + // preRequestHook is forward-compat: new .so plugins built against LLMPlugin can export + // PreRequestHook to participate in the per-request routing phase. Legacy plugins predating + // PreRequestHook leave it nil and silently no-op for routing. + preRequestHook func(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) error + preLLMHook func(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) + postLLMHook func(ctx *schemas.BifrostContext, resp *schemas.BifrostResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError, error) // MCPPlugin (optional) preMCPHook func(ctx *schemas.BifrostContext, req *schemas.BifrostMCPRequest) (*schemas.BifrostMCPRequest, *schemas.MCPPluginShortCircuit, error) @@ -79,6 +83,16 @@ func (dp *DynamicPlugin) HTTPTransportStreamChunkHook(ctx *schemas.BifrostContex return dp.httpTransportStreamChunkHook(ctx, req, stream) } +// PreRequestHook is invoked once per top-level request to decide provider/model/fallbacks +// (LLMPlugin interface). Defaults to a no-op passthrough for legacy plugins that don't +// export PreRequestHook. +func (dp *DynamicPlugin) PreRequestHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) error { + if dp.preRequestHook == nil { + return nil + } + return dp.preRequestHook(ctx, req) +} + // PreLLMHook is invoked before LLM provider calls (LLMPlugin interface) func (dp *DynamicPlugin) PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { if dp.preLLMHook == nil { diff --git a/framework/tracing/tracer_test.go b/framework/tracing/tracer_test.go index 2b467509af..d8e9e6b8cc 100644 --- a/framework/tracing/tracer_test.go +++ b/framework/tracing/tracer_test.go @@ -14,6 +14,9 @@ type testRealtimeObservabilityPlugin struct { func (p *testRealtimeObservabilityPlugin) GetName() string { return "test-observability" } func (p *testRealtimeObservabilityPlugin) Cleanup() error { return nil } +func (p *testRealtimeObservabilityPlugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} func (p *testRealtimeObservabilityPlugin) PreLLMHook(_ *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { return req, nil, nil } diff --git a/plugins/compat/main.go b/plugins/compat/main.go index 71e83bb95b..64b7f2303f 100644 --- a/plugins/compat/main.go +++ b/plugins/compat/main.go @@ -89,6 +89,11 @@ func (p *CompatPlugin) HTTPTransportStreamChunkHook(ctx *schemas.BifrostContext, return chunk, nil } +// PreRequestHook implements schemas.LLMPlugin (no-op — required for plugin indexing). +func (p *CompatPlugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + // PreLLMHook intercepts requests and applies LiteLLM-compatible request normalization. func (p *CompatPlugin) PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { if ctx == nil || req == nil { @@ -186,4 +191,4 @@ func (p *CompatPlugin) markForConversion(ctx *schemas.BifrostContext, provider s if shouldConvert { ctx.SetValue(schemas.BifrostContextKeyChangeRequestType, targetType) } -} \ No newline at end of file +} diff --git a/plugins/governance/main.go b/plugins/governance/main.go index e2c61db700..c005f84f61 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -1440,6 +1440,11 @@ func (p *GovernancePlugin) isMCPToolAllowedByVKWith(vk *configstoreTables.TableV return false } +// PreRequestHook implements schemas.LLMPlugin (no-op — required for plugin indexing). +func (p *GovernancePlugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + // PreLLMHook intercepts requests before they are processed (governance decision point) // Parameters: // - ctx: The Bifrost context diff --git a/plugins/jsonparser/main.go b/plugins/jsonparser/main.go index c0de696b0d..5cef86e4a8 100644 --- a/plugins/jsonparser/main.go +++ b/plugins/jsonparser/main.go @@ -98,6 +98,11 @@ func (p *JsonParserPlugin) HTTPTransportStreamChunkHook(ctx *schemas.BifrostCont return chunk, nil } +// PreRequestHook implements schemas.LLMPlugin (no-op — required for plugin indexing). +func (p *JsonParserPlugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + // PreLLMHook is not used for this plugin as we only process responses // Parameters: // - ctx: The Bifrost context diff --git a/plugins/logging/main.go b/plugins/logging/main.go index 3e21b23120..d9c8202ddf 100644 --- a/plugins/logging/main.go +++ b/plugins/logging/main.go @@ -503,6 +503,11 @@ func (p *LoggerPlugin) captureLoggingHeaders(ctx *schemas.BifrostContext) map[st return metadata } +// PreRequestHook implements schemas.LLMPlugin (no-op — required for plugin indexing). +func (p *LoggerPlugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + // PreLLMHook is called before a request is processed - FULLY ASYNC, NO DATABASE I/O // Parameters: // - ctx: The Bifrost context diff --git a/plugins/maxim/main.go b/plugins/maxim/main.go index ace1329e6a..14022998a6 100644 --- a/plugins/maxim/main.go +++ b/plugins/maxim/main.go @@ -229,6 +229,11 @@ func (plugin *Plugin) getOrCreateLogger(logRepoID string) (*logging.Logger, erro return logger, nil } +// PreRequestHook implements schemas.LLMPlugin (no-op — required for plugin indexing). +func (plugin *Plugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + // PreLLMHook is called before a request is processed by Bifrost. // It manages trace and generation tracking for incoming requests by either: // - Creating a new trace if none exists diff --git a/plugins/mocker/main.go b/plugins/mocker/main.go index 53ce896615..06c79a6612 100644 --- a/plugins/mocker/main.go +++ b/plugins/mocker/main.go @@ -495,6 +495,11 @@ func (p *MockerPlugin) HTTPTransportStreamChunkHook(ctx *schemas.BifrostContext, return chunk, nil } +// PreRequestHook implements schemas.LLMPlugin (no-op — required for plugin indexing). +func (p *MockerPlugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + // PreLLMHook intercepts requests and applies mocking rules based on configuration // This is called before the actual provider request and can short-circuit the flow func (p *MockerPlugin) PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { @@ -853,10 +858,10 @@ func (p *MockerPlugin) generateSuccessShortCircuit(req *schemas.BifrostRequest, }, }, ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: req.RequestType, - Provider: provider, + RequestType: req.RequestType, + Provider: provider, OriginalModelRequested: model, - Latency: int64(time.Since(startTime).Milliseconds()), + Latency: int64(time.Since(startTime).Milliseconds()), }, } } else if req.RequestType == schemas.ResponsesRequest { @@ -877,10 +882,10 @@ func (p *MockerPlugin) generateSuccessShortCircuit(req *schemas.BifrostRequest, TotalTokens: usage.TotalTokens, }, ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ResponsesRequest, - Provider: provider, + RequestType: schemas.ResponsesRequest, + Provider: provider, OriginalModelRequested: model, - Latency: int64(time.Since(startTime).Milliseconds()), + Latency: int64(time.Since(startTime).Milliseconds()), }, } } else if req.RequestType == schemas.ResponsesStreamRequest { @@ -905,10 +910,10 @@ func (p *MockerPlugin) generateSuccessShortCircuit(req *schemas.BifrostRequest, }, }, ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ResponsesStreamRequest, - Provider: provider, + RequestType: schemas.ResponsesStreamRequest, + Provider: provider, OriginalModelRequested: model, - Latency: int64(time.Since(startTime).Milliseconds()), + Latency: int64(time.Since(startTime).Milliseconds()), }, } } @@ -959,8 +964,8 @@ func (p *MockerPlugin) generateErrorShortCircuit(req *schemas.BifrostRequest, re }, AllowFallbacks: allowFallbacks, ExtraFields: schemas.BifrostErrorExtraFields{ - RequestType: req.RequestType, - Provider: provider, + RequestType: req.RequestType, + Provider: provider, OriginalModelRequested: model, }, } @@ -1083,8 +1088,8 @@ func (p *MockerPlugin) handleDefaultBehavior(req *schemas.BifrostRequest) (*sche }, }, ExtraFields: schemas.BifrostResponseExtraFields{ - RequestType: schemas.ChatCompletionRequest, - Provider: provider, + RequestType: schemas.ChatCompletionRequest, + Provider: provider, OriginalModelRequested: model, }, }, diff --git a/plugins/prompts/main.go b/plugins/prompts/main.go index e710980dcb..d74a5b9fea 100644 --- a/plugins/prompts/main.go +++ b/plugins/prompts/main.go @@ -203,6 +203,11 @@ func (p *Plugin) HTTPTransportStreamChunkHook(ctx *schemas.BifrostContext, req * return chunk, nil } +// PreRequestHook implements schemas.LLMPlugin (no-op — required for plugin indexing). +func (p *Plugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + // PreLLMHook resolves the prompt via PromptResolver, loads the version from the in-memory // cache, sets governance/observability context (selected prompt name and version), merges // version ModelParams with the request (request overrides), converts stored messages to diff --git a/plugins/semanticcache/main.go b/plugins/semanticcache/main.go index 7dbad109dc..737ef7c96e 100644 --- a/plugins/semanticcache/main.go +++ b/plugins/semanticcache/main.go @@ -330,6 +330,11 @@ func (plugin *Plugin) HTTPTransportStreamChunkHook(ctx *schemas.BifrostContext, return chunk, nil } +// PreRequestHook implements schemas.LLMPlugin (no-op — required for plugin indexing). +func (plugin *Plugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + // PreLLMHook performs the cache lookup before the request reaches the // provider. It runs the direct hash path first (cheapest), falls back to // semantic similarity search when configured, and short-circuits the diff --git a/plugins/semanticcache/plugin_no_mutation_test.go b/plugins/semanticcache/plugin_no_mutation_test.go index d0a65b681f..a00b2a7bb9 100644 --- a/plugins/semanticcache/plugin_no_mutation_test.go +++ b/plugins/semanticcache/plugin_no_mutation_test.go @@ -30,6 +30,10 @@ type requestCapturer struct { func (p *requestCapturer) GetName() string { return "test-request-capturer" } func (p *requestCapturer) Cleanup() error { return nil } +func (p *requestCapturer) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + func (p *requestCapturer) PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { p.mu.Lock() // Snapshot the request via JSON round-trip so any later mutation by the diff --git a/plugins/telemetry/main.go b/plugins/telemetry/main.go index f28a6b0f2f..f823f1bbc8 100644 --- a/plugins/telemetry/main.go +++ b/plugins/telemetry/main.go @@ -613,6 +613,11 @@ func (p *PrometheusPlugin) HTTPTransportStreamChunkHook(ctx *schemas.BifrostCont return chunk, nil } +// PreRequestHook implements schemas.LLMPlugin (no-op — required for plugin indexing). +func (p *PrometheusPlugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + // 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) { diff --git a/transports/bifrost-http/handlers/realtime_client_secrets_test.go b/transports/bifrost-http/handlers/realtime_client_secrets_test.go index 8c1b83dfa8..b7ed942b23 100644 --- a/transports/bifrost-http/handlers/realtime_client_secrets_test.go +++ b/transports/bifrost-http/handlers/realtime_client_secrets_test.go @@ -311,6 +311,10 @@ func (m *mockRealtimeMintingGovernancePlugin) HTTPTransportPostHook(_ *schemas.B return nil } +func (m *mockRealtimeMintingGovernancePlugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + func (m *mockRealtimeMintingGovernancePlugin) PreLLMHook(_ *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { return req, nil, nil } diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go index a1b612d270..e9429c823a 100644 --- a/transports/bifrost-http/lib/config_test.go +++ b/transports/bifrost-http/lib/config_test.go @@ -12243,6 +12243,10 @@ type mockLLMPlugin struct { mockPlugin } +func (p *mockLLMPlugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { + return nil +} + func (p *mockLLMPlugin) PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { return req, nil, nil } From 27b1da540d6002f23511825c4d657c323d615668 Mon Sep 17 00:00:00 2001 From: Pratham Mishra <99235987+Pratham-Mishra04@users.noreply.github.com> Date: Tue, 9 Jun 2026 15:44:45 +0530 Subject: [PATCH 003/108] refactor: migrate governance routing from `HTTPTransportPreHook` to `PreRequestHook` with `BifrostRequest`-native mutations (#4176) ## Summary Governance routing logic has been migrated from `HTTPTransportPreHook` to `PreRequestHook`, operating directly on `BifrostRequest` structs rather than raw HTTP bodies. This eliminates the need to unmarshal/marshal JSON, parse multipart forms, or extract models from URL path parameters inside the governance plugin. Integration-specific normalization (Gemini, Bedrock, etc.) now happens upstream before the request reaches governance, so the plugin sees a clean `provider/model` pair regardless of the originating integration. ## Changes - `HTTPTransportPreHook` is now a no-op stub retained only to satisfy the `HTTPTransportPlugin` interface; all routing flows through `PreRequestHook` - `loadBalanceProvider` and `applyRoutingRules` now accept `*schemas.BifrostRequest` instead of `map[string]any` + `*schemas.HTTPRequest`, mutating `Provider`/`Model`/`Fallbacks` directly via typed setters - `addMCPIncludeTools` (header mutation) replaced by `computeMCPIncludeTools` (returns `[]string`); result is stored on context via `MCPContextKeyIncludeTools` rather than written into HTTP headers - `validateRequiredHeaders` and `stampGovernanceCtxFromVK` moved to `utils.go`; `stampGovernanceCtxFromVK` also fixes a bug where `Team.CustomerID`/`Team.Customer` were not propagated when the VK had a team association - `governLargePayload` and `governRealtimeQueryParam` removed; large-payload routing now runs through `runPreRequestRouting` (a thin wrapper that builds a synthetic `BifrostRequest` and calls the same helpers), and realtime WebSocket upgrades are routed via an explicit `RunPreRequestHooks` call in `wsrealtime.go` before the upgrade completes - `BifrostContextKeyRequestQuery` context key added; query params (lowercased) are now populated in `ConvertToBifrostContext` for normal HTTP requests and explicitly in the WS upgrade path, making them available to governance CEL routing rules via `params["..."]` - Fallbacks produced by `loadBalanceProvider` are now `[]schemas.Fallback` (typed) instead of `[]string` - Tests that exercised the old `HTTPTransportPreHook` body-parsing path are skipped with a note to rewrite them as `PreRequestHook` tests in Phase 3 ## Type of change - [ ] Bug fix - [ ] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [ ] Core (Go) - [x] Transports (HTTP) - [ ] Providers/Integrations - [x] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./plugins/governance/... go test ./transports/bifrost-http/... ``` Verify that: - Governance routing rules (CEL expressions reading `headers[...]` and `params[...]`) resolve correctly for both normal HTTP and WebSocket realtime upgrade requests - Virtual key load balancing selects a provider and populates typed fallbacks on the `BifrostRequest` - Large-payload streaming requests route via `LargePayloadMetadata.Model` and the rewritten model is visible to the upstream provider's body rewriter - Realtime WebSocket upgrades pick up the governance-routed `provider`/`model` before the connection is established ## Breaking changes - [x] Yes - [ ] No `applyRoutingRules` and `loadBalanceProvider` signatures have changed from `(ctx, *HTTPRequest, map[string]any, *TableVirtualKey)` to `(ctx, *BifrostRequest, *TableVirtualKey)`. Any code calling these methods directly (outside the governance plugin itself) must be updated. The `addMCPIncludeTools` method has been removed; callers should use `computeMCPIncludeTools` and store the result on context. ## Related issues Closes #2516 ## Security considerations No new auth surfaces introduced. Query params are now stored on context with lowercased keys, consistent with how request headers are already handled. No secrets or PII are added to context beyond what was already present. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [ ] I added/updated tests where appropriate - [ ] I updated documentation where needed - [ ] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable ## Summary by CodeRabbit * **New Features** * Governance routing and load balancing now support WebSocket realtime connections. * Request query parameters are now accessible for governance routing rules. * **Improvements** * Enhanced governance routing by moving rule evaluation to an earlier request processing phase for improved handling of large payloads. * Added required header validation in governance rules. * **Tests** * Migrated governance pre-hook tests; rewrite pending in Phase 3. * **Chores** * Updated indirect dependency version. --- core/schemas/bifrost.go | 1 + plugins/governance/go.mod | 1 + .../governance/httptransportprehook_test.go | 357 -------- plugins/governance/main.go | 829 +++++------------- plugins/governance/utils.go | 95 +- .../bifrost-http/handlers/wsrealtime.go | 62 ++ transports/bifrost-http/lib/ctx.go | 12 + 7 files changed, 355 insertions(+), 1002 deletions(-) diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 8cdcb8470c..e9f8c1c7e5 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -290,6 +290,7 @@ const ( BifrostContextKeyRealtimeVoice BifrostContextKey = "bifrost-realtime-voice" // string BifrostIsAsyncRequest BifrostContextKey = "bifrost-is-async-request" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) - whether the request is an async request (only used in gateway) BifrostContextKeyRequestHeaders BifrostContextKey = "bifrost-request-headers" // map[string]string (all request headers with lowercased keys) + BifrostContextKeyRequestQuery BifrostContextKey = "bifrost-request-query" // map[string]string (request query params with lowercased keys; consumed by governance routing CEL rules) BifrostContextKeyAllowPerRequestStorageOverride BifrostContextKey = "bifrost-allow-per-request-storage-override" // bool (set by transport from config — gates whether x-bf-disable-content-logging and x-bf-store-raw-request-response per-request overrides are honored) BifrostContextKeyAllowPerRequestRawOverride BifrostContextKey = "bifrost-allow-per-request-raw-override" // bool (set by transport from config — gates whether x-bf-send-back-raw-request and x-bf-send-back-raw-response per-request overrides are honored) BifrostContextKeyDisableContentLogging BifrostContextKey = "x-bf-disable-content-logging" // bool (per-request override for content logging; only honored when BifrostContextKeyAllowPerRequestStorageOverride is true) diff --git a/plugins/governance/go.mod b/plugins/governance/go.mod index 8e6d9d740d..44b0a3effc 100644 --- a/plugins/governance/go.mod +++ b/plugins/governance/go.mod @@ -55,6 +55,7 @@ require ( github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/buger/jsonparser v1.1.2 // indirect github.com/bytedance/gopkg v0.1.3 // indirect + github.com/bytedance/sonic v1.15.0 // indirect github.com/bytedance/sonic/loader v0.5.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cloudwego/base64x v0.1.6 // indirect diff --git a/plugins/governance/httptransportprehook_test.go b/plugins/governance/httptransportprehook_test.go index c05eb95707..7e3f1db452 100644 --- a/plugins/governance/httptransportprehook_test.go +++ b/plugins/governance/httptransportprehook_test.go @@ -2,10 +2,8 @@ package governance import ( "context" - "encoding/json" "testing" - bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore" configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" @@ -13,58 +11,6 @@ import ( "github.com/stretchr/testify/require" ) -// TestHTTPTransportPreHook_VirtualKeyReplicateRefinesNestedModel verifies that -// virtual-key provider pinning rewrites the request model to Replicate's nested provider slug. -func TestHTTPTransportPreHook_VirtualKeyReplicateRefinesNestedModel(t *testing.T) { - logger := NewMockLogger() - mc := modelcatalog.NewTestCatalog(map[string]string{ - "openai/gpt-5-nano": "gpt-5-nano", - }) - mc.UpsertModelDataForProvider(schemas.Replicate, &schemas.BifrostListModelsResponse{ - Data: []schemas.Model{ - {ID: "replicate/openai/gpt-5-nano"}, - }, - }, nil) - - virtualKey := buildVirtualKeyWithProviders( - "vk1", - "sk-bf-test", - "replicate-only", - []configstoreTables.TableVirtualKeyProviderConfig{ - buildProviderConfig("replicate", []string{"*"}), - }, - ) - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - }, mc) - require.NoError(t, err) - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, mc, nil, nil) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/v1/chat/completions" - req.Headers["Authorization"] = "Bearer sk-bf-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"model":"gpt-5-nano","messages":[{"role":"user","content":"Hello!"}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - var payload struct { - Model string `json:"model"` - } - require.NoError(t, json.Unmarshal(req.Body, &payload)) - require.Equal(t, "replicate/openai/gpt-5-nano", payload.Model) -} - func TestHTTPTransportPreHook_ModelOnlyVirtualKeySetsAvailableProviders(t *testing.T) { logger := NewMockLogger() @@ -398,306 +344,3 @@ func TestHTTPTransportPreHook_WildcardOpaqueProviderRespectsBlacklist(t *testing // Blacklisted model is excluded even though the provider is catalog-opaque under ["*"]. require.Empty(t, allowedProviders) } - -// TestHTTPTransportPreHook_GenAIRoutingRulePreservesTarget verifies that when a routing rule -// matches on the /genai path, governance load balancing does not override the routing-rule target -// with a provider from the VK pool (regression test for issue #2516). -func TestHTTPTransportPreHook_GenAIRoutingRulePreservesTarget(t *testing.T) { - logger := NewMockLogger() - - routingRule := configstoreTables.TableRoutingRule{ - ID: "rule-genai-1", - Name: "genai-repro-rule", - Enabled: bifrost.Ptr(true), - CelExpression: `model == "probe-genai-model" && provider == ""`, - Targets: []configstoreTables.TableRoutingTarget{ - { - RuleID: "rule-genai-1", - Provider: bifrost.Ptr("repro-openai-a"), - Model: bifrost.Ptr("error-test"), - Weight: 1.0, - }, - }, - Scope: "global", - Priority: 1, - } - - // VK with repro-openai-b at weight=1 — this is what governance LB would wrongly select without the fix - virtualKey := buildVirtualKeyWithProviders( - "vk-genai", - "sk-bf-genai-test", - "genai-repro-vk", - []configstoreTables.TableVirtualKeyProviderConfig{ - buildProviderConfig("repro-openai-b", []string{"*"}), - }, - ) - - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - RoutingRules: []configstoreTables.TableRoutingRule{routingRule}, - }, nil) - require.NoError(t, err) - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/genai/v1beta/models/probe-genai-model:generateContent" - req.PathParams["model"] = "probe-genai-model:generateContent" - req.Headers["Authorization"] = "Bearer sk-bf-genai-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - // Routing rule matched and set context model to "repro-openai-a/error-test:generateContent". - // Governance LB must NOT override this with "repro-openai-b/probe-genai-model:generateContent". - ctxModel, ok := bfCtx.Value("model").(string) - require.True(t, ok, "context model should be set") - require.Equal(t, "repro-openai-a/error-test:generateContent", ctxModel) -} - -// TestHTTPTransportPreHook_GenAIRoutingRulePreservesTarget_WithStore is a production-like variant -// of TestHTTPTransportPreHook_GenAIRoutingRulePreservesTarget that passes a non-nil inMemoryStore -// containing the routing-rule provider, confirming the fix holds when p.inMemoryStore != nil -// and the provider IS present in GetConfiguredProviders (the normal production code path). -func TestHTTPTransportPreHook_GenAIRoutingRulePreservesTarget_WithStore(t *testing.T) { - logger := NewMockLogger() - - routingRule := configstoreTables.TableRoutingRule{ - ID: "rule-genai-ws-1", - Name: "genai-repro-rule-with-store", - Enabled: bifrost.Ptr(true), - CelExpression: `model == "probe-genai-model" && provider == ""`, - Targets: []configstoreTables.TableRoutingTarget{ - { - RuleID: "rule-genai-ws-1", - Provider: bifrost.Ptr("repro-openai-a"), - Model: bifrost.Ptr("error-test"), - Weight: 1.0, - }, - }, - Scope: "global", - Priority: 1, - } - - virtualKey := buildVirtualKeyWithProviders( - "vk-genai-ws", - "sk-bf-genai-ws-test", - "genai-repro-vk-with-store", - []configstoreTables.TableVirtualKeyProviderConfig{ - buildProviderConfig("repro-openai-b", []string{"*"}), - }, - ) - - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - RoutingRules: []configstoreTables.TableRoutingRule{routingRule}, - }, nil) - require.NoError(t, err) - - // Register the fake provider so ParseModelString can split "repro-openai-a/model" - // the same way it would for a real provider in production. - schemas.RegisterKnownProvider("repro-openai-a") - t.Cleanup(func() { schemas.UnregisterKnownProvider("repro-openai-a") }) - - // Use a non-nil inMemoryStore that recognises the routing-rule provider, - // mirroring production where configured providers are always registered in the store. - inMemStore := &mockInMemoryStore{ - configuredProviders: map[schemas.ModelProvider]configstore.ProviderConfig{ - "repro-openai-a": {}, - }, - } - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, inMemStore) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/genai/v1beta/models/probe-genai-model:generateContent" - req.PathParams["model"] = "probe-genai-model:generateContent" - req.Headers["Authorization"] = "Bearer sk-bf-genai-ws-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - ctxModel, ok := bfCtx.Value("model").(string) - require.True(t, ok, "context model should be set") - require.Equal(t, "repro-openai-a/error-test:generateContent", ctxModel) -} - -// TestHTTPTransportPreHook_GenAINoRoutingRuleStillLoadBalances verifies that when no routing rule -// matches on the /genai path, governance load balancing still selects a provider from the VK pool. -func TestHTTPTransportPreHook_GenAINoRoutingRuleStillLoadBalances(t *testing.T) { - logger := NewMockLogger() - - // VK with repro-openai-b at weight=1 — LB should select this - virtualKey := buildVirtualKeyWithProviders( - "vk-genai-lb", - "sk-bf-genai-lb-test", - "genai-lb-vk", - []configstoreTables.TableVirtualKeyProviderConfig{ - buildProviderConfig("repro-openai-b", []string{"*"}), - }, - ) - - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - // No routing rules — governance LB should run normally - }, nil) - require.NoError(t, err) - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/genai/v1beta/models/probe-genai-model:generateContent" - req.PathParams["model"] = "probe-genai-model:generateContent" - req.Headers["Authorization"] = "Bearer sk-bf-genai-lb-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - // No routing rule: governance LB must still run and select repro-openai-b from the VK pool - ctxModel, ok := bfCtx.Value("model").(string) - require.True(t, ok, "context model should be set by governance LB") - require.Equal(t, "repro-openai-b/probe-genai-model:generateContent", ctxModel) -} - -// TestHTTPTransportPreHook_BedrockRoutingRulePreservesTarget verifies that when a routing rule -// matches on the /bedrock path, governance load balancing does not override the routing-rule target -// (regression test mirroring the GenAI fix for the Bedrock integration). -func TestHTTPTransportPreHook_BedrockRoutingRulePreservesTarget(t *testing.T) { - logger := NewMockLogger() - - routingRule := configstoreTables.TableRoutingRule{ - ID: "rule-bedrock-1", - Name: "bedrock-repro-rule", - Enabled: bifrost.Ptr(true), - CelExpression: `model == "probe-bedrock-model" && provider == ""`, - Targets: []configstoreTables.TableRoutingTarget{ - { - RuleID: "rule-bedrock-1", - Provider: bifrost.Ptr("repro-openai-a"), - Model: bifrost.Ptr("error-test"), - Weight: 1.0, - }, - }, - Scope: "global", - Priority: 1, - } - - // VK with repro-openai-b at weight=1 — this is what governance LB would wrongly select without the fix - virtualKey := buildVirtualKeyWithProviders( - "vk-bedrock", - "sk-bf-bedrock-test", - "bedrock-repro-vk", - []configstoreTables.TableVirtualKeyProviderConfig{ - buildProviderConfig("repro-openai-b", []string{"*"}), - }, - ) - - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - RoutingRules: []configstoreTables.TableRoutingRule{routingRule}, - }, nil) - require.NoError(t, err) - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/bedrock/model/probe-bedrock-model/converse" - req.PathParams["modelId"] = "probe-bedrock-model" - req.Headers["Authorization"] = "Bearer sk-bf-bedrock-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - // Routing rule matched and set context modelId to "repro-openai-a/error-test". - // Governance LB must NOT override this with "repro-openai-b/probe-bedrock-model". - ctxModelID, ok := bfCtx.Value("modelId").(string) - require.True(t, ok, "context modelId should be set") - require.Equal(t, "repro-openai-a/error-test", ctxModelID) -} - -// TestHTTPTransportPreHook_BedrockNoRoutingRuleStillLoadBalances verifies that when no routing rule -// matches on the /bedrock path, governance load balancing still selects a provider from the VK pool. -func TestHTTPTransportPreHook_BedrockNoRoutingRuleStillLoadBalances(t *testing.T) { - logger := NewMockLogger() - - // VK with repro-openai-b at weight=1 — LB should select this - virtualKey := buildVirtualKeyWithProviders( - "vk-bedrock-lb", - "sk-bf-bedrock-lb-test", - "bedrock-lb-vk", - []configstoreTables.TableVirtualKeyProviderConfig{ - buildProviderConfig("repro-openai-b", []string{"*"}), - }, - ) - - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - // No routing rules — governance LB should run normally - }, nil) - require.NoError(t, err) - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/bedrock/model/probe-bedrock-model/converse" - req.PathParams["modelId"] = "probe-bedrock-model" - req.Headers["Authorization"] = "Bearer sk-bf-bedrock-lb-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - // No routing rule: governance LB must still run and select repro-openai-b from the VK pool - ctxModelID, ok := bfCtx.Value("modelId").(string) - require.True(t, ok, "context modelId should be set by governance LB") - require.Equal(t, "repro-openai-b/probe-bedrock-model", ctxModelID) -} diff --git a/plugins/governance/main.go b/plugins/governance/main.go index c005f84f61..fb47d59d0b 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -6,17 +6,13 @@ import ( "errors" "fmt" "math/rand/v2" - "net/url" "sort" "strings" "sync" "time" - "github.com/bytedance/sonic" "github.com/google/uuid" bifrost "github.com/maximhq/bifrost/core" - "github.com/maximhq/bifrost/core/network" - "github.com/maximhq/bifrost/core/providers/gemini" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore" configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" @@ -345,331 +341,54 @@ func (p *GovernancePlugin) UpdateEnforceAuthOnInference(enforceAuthOnInference b p.isVkMandatory = new(enforceAuthOnInference) } -// HTTPTransportPreHook intercepts requests before they are processed (governance decision point) -// It modifies the request in-place and returns nil to continue, or an HTTPResponse to short-circuit. -// Optimized to skip unnecessary operations: only unmarshals/marshals when needed +// HTTPTransportPreHook is retained as a no-op so governance still satisfies the +// HTTPTransportPlugin interface (used by the enterprise wrapper's 503 gate delegation). +// All routing now flows through PreRequestHook: body-having requests via handleRequest, +// large-payload requests via PreRequestHook reading LargePayloadMetadata, and realtime WS +// upgrades via the realtime handler's explicit RunPreRequestHooks call. func (p *GovernancePlugin) HTTPTransportPreHook(ctx *schemas.BifrostContext, req *schemas.HTTPRequest) (*schemas.HTTPResponse, error) { - virtualKeyValue := parseVirtualKeyFromHTTPRequest(req) - hasRoutingRules := p.store.HasRoutingRules(ctx) - - if strings.Contains(req.Path, "passthrough") { - return nil, nil - } - - // If no virtual key and no routing rules configured, skip all processing - if virtualKeyValue == nil && !hasRoutingRules { - return nil, nil - } - - // If no body, check if the request carries a model via query params (e.g. realtime - // WebSocket upgrades: GET /v1/realtime?model=... or Azure preview ?deployment=...) - // or if large payload mode is active. - // For query-param-based models we build a synthetic payload so routing rules and VK - // load-balancing can rewrite provider/model, then propagate changes back to the query. - if len(req.Body) == 0 { - if modelParam := realtimeModelQueryParam(req); modelParam != "" { - return p.governRealtimeQueryParam(ctx, req, virtualKeyValue, hasRoutingRules) - } - isLargePayload, _ := ctx.Value(schemas.BifrostContextKeyLargePayloadMode).(bool) - if !isLargePayload { - return nil, nil - } - return p.governLargePayload(ctx, req, virtualKeyValue, hasRoutingRules) - } - - // Only unmarshal if we have VK or routing rules - var payload map[string]any - var virtualKey *configstoreTables.TableVirtualKey - var ok bool - var needsMarshal bool - - contentType := req.CaseInsensitiveHeaderLookup("Content-Type") - lowerCT := strings.ToLower(contentType) - // Strip parameters (e.g., "; charset=utf-8") for clean media type comparison - mediaType := lowerCT - if idx := strings.IndexByte(mediaType, ';'); idx >= 0 { - mediaType = strings.TrimSpace(mediaType[:idx]) - } - isMultipart := strings.HasPrefix(mediaType, "multipart/form-data") - isJSON := mediaType == "" || mediaType == "application/json" || strings.HasSuffix(mediaType, "+json") - - if !isMultipart && !isJSON { - // Non-parseable body (e.g., application/sdp for WebRTC signaling) — skip governance - return nil, nil - } - - var err error - if isMultipart { - payload, err = network.ParseMultipartFormFields(contentType, req.Body) - if err != nil { - p.logger.Warn("failed to parse multipart form in governance plugin: %v", err) - return nil, nil - } - } else { - err = sonic.Unmarshal(req.Body, &payload) - if err != nil { - p.logger.Error("failed to unmarshal request body: %v", err) - return nil, nil - } - } - - // Process virtual key if provided - if virtualKeyValue != nil { - virtualKey, ok = p.store.GetVirtualKey(ctx, *virtualKeyValue) - if !ok || virtualKey == nil || !virtualKey.IsActiveValue() { - return nil, nil - } - } - - // Attaching team and customer based on the virtual key - if virtualKey != nil { - if virtualKey.TeamID != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamID, *virtualKey.TeamID) - } - if virtualKey.Team != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamName, virtualKey.Team.Name) - } - if virtualKey.CustomerID != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerID, *virtualKey.CustomerID) - } - if virtualKey.Customer != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerName, virtualKey.Customer.Name) - } - } - - //1. Apply routing rules only if we have rules or matched decision - var routingDecision *RoutingDecision - if hasRoutingRules { - var err error - payload, routingDecision, err = p.applyRoutingRules(ctx, req, payload, virtualKey) - if err != nil { - return nil, err - } - // Mark for marshal if a routing rule matched - if routingDecision != nil { - needsMarshal = true - } - } - - // Process virtual key if provided - if virtualKey != nil { - //2. Load balance provider - payload, err = p.loadBalanceProvider(ctx, req, payload, virtualKey) - if err != nil { - return nil, err - } - //3. Add MCP tools only when auto-inject is enabled and header not already set by the caller - p.cfgMutex.RLock() - autoInjectDisabled := p.disableAutoToolInject != nil && *p.disableAutoToolInject - p.cfgMutex.RUnlock() - if !autoInjectDisabled { - // Treat an explicitly-present (even empty) x-bf-mcp-include-tools header as "present" - // so that callers can block auto-injection by sending an empty header value. - headerPresent := false - for k := range req.Headers { - if strings.EqualFold(k, "x-bf-mcp-include-tools") { - headerPresent = true - break - } - } - if !headerPresent { - req.Headers, err = p.addMCPIncludeTools(req.Headers, virtualKey) - if err != nil { - p.logger.Error("failed to add MCP include tools: %v", err) - return nil, nil - } - } - } - needsMarshal = true - } - - // Only marshal if something changed (VK processing or routing decision matched) - if needsMarshal { - if err := network.SerializePayloadToRequest(req, payload, isMultipart, contentType); err != nil { - p.logger.Error("failed to serialize request body in governance plugin: %v", err) - return nil, nil - } - } - return nil, nil } -// governLargePayload handles read-only governance for large payload requests. -// The request body is streaming and cannot be modified, so we build a synthetic payload -// from pre-extracted metadata and run VK validation, routing rules, and load balancing. -// Any model changes are propagated via the metadata in context (not body rewriting). -func (p *GovernancePlugin) governLargePayload(ctx *schemas.BifrostContext, req *schemas.HTTPRequest, virtualKeyValue *string, hasRoutingRules bool) (*schemas.HTTPResponse, error) { - metadata, _ := ctx.Value(schemas.BifrostContextKeyLargePayloadMetadata).(*schemas.LargePayloadMetadata) - if metadata == nil || metadata.Model == "" { - return nil, nil - } - - // Build synthetic payload from metadata — only the model field is needed - payload := map[string]any{ - "model": metadata.Model, +// runPreRequestRouting wraps a model string in a synthetic BifrostRequest, runs the same +// applyRoutingRules + loadBalanceProvider helpers used by the main PreRequestHook path, and +// returns the resolved model (provider-prefixed when a provider was selected, plain model +// otherwise). Used by PreRequestHook's large-payload branch where req.Model is empty because +// the body wasn't parsed. +func (p *GovernancePlugin) runPreRequestRouting(ctx *schemas.BifrostContext, virtualKey *configstoreTables.TableVirtualKey, hasRoutingRules bool, modelIn string, requestType schemas.RequestType) (string, error) { + synthetic := &schemas.BifrostRequest{ + RequestType: requestType, + ChatRequest: &schemas.BifrostChatRequest{Model: modelIn}, } - originalModel := metadata.Model - // Process virtual key if provided - var virtualKey *configstoreTables.TableVirtualKey - if virtualKeyValue != nil { - vk, ok := p.store.GetVirtualKey(ctx, *virtualKeyValue) - if !ok || vk == nil || !vk.IsActiveValue() { - return nil, nil - } - virtualKey = vk - } - - // Attaching team and customer based on the virtual key - if virtualKey != nil { - if virtualKey.TeamID != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamID, *virtualKey.TeamID) - } - if virtualKey.Team != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamName, virtualKey.Team.Name) - } - if virtualKey.CustomerID != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerID, *virtualKey.CustomerID) - } - if virtualKey.Customer != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerName, virtualKey.Customer.Name) - } - } - - // Apply routing rules (read-only: decisions still affect downstream evaluation) if hasRoutingRules { - var err error - payload, _, err = p.applyRoutingRules(ctx, req, payload, virtualKey) - if err != nil { - return nil, err + if _, err := p.applyRoutingRules(ctx, synthetic, virtualKey); err != nil { + return modelIn, err } } - // Process virtual key: load balance + MCP tool headers if virtualKey != nil { - var err error - payload, err = p.loadBalanceProvider(ctx, req, payload, virtualKey) - if err != nil { - return nil, err + if err := p.loadBalanceProvider(ctx, synthetic, virtualKey); err != nil { + return modelIn, err } - // MCP tool headers — apply the same auto-inject guard as the normal path: - // skip when DisableAutoToolInject is set or the caller already sent the header. + p.cfgMutex.RLock() autoInjectDisabled := p.disableAutoToolInject != nil && *p.disableAutoToolInject p.cfgMutex.RUnlock() if !autoInjectDisabled { - headerPresent := false - for k := range req.Headers { - if strings.EqualFold(k, "x-bf-mcp-include-tools") { - headerPresent = true - break + if existing := ctx.Value(schemas.MCPContextKeyIncludeTools); existing == nil { + if tools := p.computeMCPIncludeTools(virtualKey); tools != nil { + ctx.SetValue(schemas.MCPContextKeyIncludeTools, tools) } } - if !headerPresent { - req.Headers, err = p.addMCPIncludeTools(req.Headers, virtualKey) - if err != nil { - p.logger.Error("failed to add MCP include tools: %v", err) - return nil, nil - } - } - } - } - - // Propagate model changes to metadata so downstream hydration picks up - // the load-balanced/routed model (e.g., provider prefix added by LB). - if newModel, ok := payload["model"].(string); ok && newModel != originalModel { - metadata.Model = newModel - } - - // No body serialization — large payload body streams through unchanged - return nil, nil -} - -// realtimeModelQueryParam returns the query parameter used as the realtime model selector. -// Azure preview realtime uses `deployment`, while GA/OpenAI-compatible paths use `model`. -func realtimeModelQueryParam(req *schemas.HTTPRequest) string { - if req == nil || req.Query == nil { - return "" - } - if modelParam := req.Query["model"]; modelParam != "" { - return modelParam - } - return req.Query["deployment"] -} - -// governRealtimeQueryParam handles governance for bodyless realtime requests -// (e.g. WebSocket upgrade GET /v1/realtime?model=... or Azure preview -// /realtime?deployment=...) where the model lives in a query parameter instead -// of the JSON body. We build a synthetic payload so routing rules and VK -// load-balancing can evaluate normally, then propagate any model rewrite back -// to the original query param for the downstream handler to pick up. -func (p *GovernancePlugin) governRealtimeQueryParam(ctx *schemas.BifrostContext, req *schemas.HTTPRequest, virtualKeyValue *string, hasRoutingRules bool) (*schemas.HTTPResponse, error) { - modelQueryKey := "model" - modelParam := req.Query[modelQueryKey] - if modelParam == "" { - modelQueryKey = "deployment" - modelParam = req.Query[modelQueryKey] - } - if modelParam == "" { - return nil, nil - } - - payload := map[string]any{ - "model": modelParam, - } - originalModel := modelParam - - // Process virtual key if provided - var virtualKey *configstoreTables.TableVirtualKey - if virtualKeyValue != nil { - vk, ok := p.store.GetVirtualKey(ctx, *virtualKeyValue) - if !ok || vk == nil || !vk.IsActiveValue() { - return nil, nil - } - virtualKey = vk - } - - // Attaching team and customer based on the virtual key - if virtualKey != nil { - if virtualKey.TeamID != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamID, *virtualKey.TeamID) - } - if virtualKey.Team != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamName, virtualKey.Team.Name) - } - if virtualKey.CustomerID != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerID, *virtualKey.CustomerID) - } - if virtualKey.Customer != nil { - ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerName, virtualKey.Customer.Name) - } - } - - // Apply routing rules - if hasRoutingRules { - var err error - payload, _, err = p.applyRoutingRules(ctx, req, payload, virtualKey) - if err != nil { - return nil, err - } - } - - // Process virtual key: load balance provider - if virtualKey != nil { - var err error - payload, err = p.loadBalanceProvider(ctx, req, payload, virtualKey) - if err != nil { - return nil, err } } - // Propagate model changes back to the original query param so the downstream - // realtime handler sees the routed/load-balanced model. - if newModel, ok := payload["model"].(string); ok && newModel != originalModel { - req.Query[modelQueryKey] = newModel + provider, model, _ := synthetic.GetRequestFields() + if provider != "" { + return string(provider) + "/" + model, nil } - - return nil, nil + return model, nil } // HTTPTransportPostHook intercepts requests after they are processed (governance decision point) @@ -683,76 +402,26 @@ func (p *GovernancePlugin) HTTPTransportStreamChunkHook(ctx *schemas.BifrostCont return chunk, nil } -// loadBalanceProvider loads balances the provider for the request -// Parameters: -// - req: The HTTP request -// - body: The request body -// - virtualKey: The virtual key configuration -// -// Returns: -// - map[string]any: The updated request body -// - error: Any error that occurred during processing -func (p *GovernancePlugin) loadBalanceProvider(ctx *schemas.BifrostContext, req *schemas.HTTPRequest, body map[string]any, virtualKey *configstoreTables.TableVirtualKey) (map[string]any, error) { - // Check if the request has a model field - modelValue, hasModel := body["model"] - isGeminiPath := strings.Contains(req.Path, "/genai") - isBedrockPath := strings.Contains(req.Path, "/bedrock") - if !hasModel { - // For genai integration, model is present in URL path instead of the request body - if isGeminiPath { - // Prefer context value set by a routing rule (format: "provider/model:suffix") - if ctxModel, ok := ctx.Value("model").(string); ok && ctxModel != "" { - modelValue = ctxModel - } else { - modelValue = req.CaseInsensitivePathParamLookup("model") - } - } else if isBedrockPath { - // For bedrock integration, model is present in URL path as modelId - // Prefer context value set by a routing rule (format: "provider/model") - if ctxModelID, ok := ctx.Value("modelId").(string); ok && ctxModelID != "" { - modelValue = ctxModelID - } else { - rawModelID := req.CaseInsensitivePathParamLookup("modelId") - if rawModelID == "" { - return body, nil - } - // URL-decode the modelId (Bedrock model IDs may be URL-encoded, e.g. anthropic%2Fclaude-3-5-sonnet) - decoded, err := url.PathUnescape(rawModelID) - if err != nil { - decoded = rawModelID - } - modelValue = decoded - } - } else { - return body, nil - } - } - modelStr, ok := modelValue.(string) - if !ok || modelStr == "" { - return body, nil - } - var genaiRequestSuffix string - // Remove Google GenAI API endpoint suffixes if present - if isGeminiPath { - for _, sfx := range gemini.GeminiRequestSuffixPaths { - if before, ok := strings.CutSuffix(modelStr, sfx); ok { - modelStr = before - genaiRequestSuffix = sfx - break - } - } +// loadBalanceProvider picks a weighted provider from the VK's configs for req.Model +// and mutates req.Provider/req.Model with the refined provider/model. Also populates req.Fallbacks +// from the remaining weighted providers if no fallbacks were configured by the caller. +func (p *GovernancePlugin) loadBalanceProvider(ctx *schemas.BifrostContext, req *schemas.BifrostRequest, virtualKey *configstoreTables.TableVirtualKey) error { + _, modelStr, existingFallbacks := req.GetRequestFields() + if modelStr == "" { + return nil } - // Check if model already has provider prefix (contains "/") + + // Model already has provider prefix pointing to a configured provider → leave it alone. if strings.Contains(modelStr, "/") { provider, _ := schemas.ParseModelString(modelStr, "") // Checking valid provider when store is available; if store is nil, // assume the prefixed model should be left unchanged. if p.inMemoryStore != nil { if _, ok := p.inMemoryStore.GetConfiguredProviders()[provider]; ok { - return body, nil + return nil } } else { - return body, nil + return nil } } @@ -764,7 +433,7 @@ func (p *GovernancePlugin) loadBalanceProvider(ctx *schemas.BifrostContext, req ctx.SetValue(schemas.BifrostContextKeyAvailableProviders, []schemas.ModelProvider{}) ctx.AppendRoutingEngineLog(schemas.RoutingEngineGovernance, schemas.LogLevelWarn, fmt.Sprintf("No provider configs on virtual key %s for model %s, skipping load balancing", virtualKey.Name, modelStr)) // No provider configs, continue without modification - return body, nil + return nil } var configuredProviders []string @@ -838,9 +507,9 @@ func (p *GovernancePlugin) loadBalanceProvider(ctx *schemas.BifrostContext, req ctx.AppendRoutingEngineLog(schemas.RoutingEngineGovernance, schemas.LogLevelInfo, fmt.Sprintf("No eligible providers remaining after filtering for model %s, skipping load balancing", modelStr)) // TODO: Send proper error if (overall VK budget/rate limit) or (all provider budgets/rate limits) are violated // No allowed provider configs, continue without modification - return body, nil + return nil } - // Separate providers with weight set (participate in routing) from those without (nil weight = excluded from routing) + weightedConfigs := make([]configstoreTables.TableVirtualKeyProviderConfig, 0, len(allowedProviderConfigs)) for _, config := range allowedProviderConfigs { if config.Weight != nil { @@ -848,261 +517,168 @@ func (p *GovernancePlugin) loadBalanceProvider(ctx *schemas.BifrostContext, req } } - var selectedProvider schemas.ModelProvider + if len(weightedConfigs) == 0 { + // All allowed configs survived the model-allowance / budget / rate-limit filters, + // but none of them have a Weight set — there's nothing to feed weighted selection. + // Emit an explicit log so the routing trail explains why governance stops here + // instead of trailing off after "Allowed providers after filtering: [...]". + ctx.AppendRoutingEngineLog(schemas.RoutingEngineGovernance, schemas.LogLevelInfo, fmt.Sprintf("No weighted configs for model %s — none of the allowed VK provider configs have a weight assigned; skipping load balancing", modelStr)) + return nil + } - if len(weightedConfigs) > 0 { - // Weighted random selection from providers that have weight set - totalWeight := 0.0 - for _, config := range weightedConfigs { - totalWeight += getWeight(config.Weight) - } - // Generate random number between 0 and totalWeight - randomValue := rand.Float64() * totalWeight - // Select provider based on weighted random selection - currentWeight := 0.0 - for _, config := range weightedConfigs { - currentWeight += getWeight(config.Weight) - if randomValue <= currentWeight { - selectedProvider = schemas.ModelProvider(config.Provider) - break - } - } - // Fallback: if no provider was selected (shouldn't happen but guard against FP issues) - if selectedProvider == "" { - selectedProvider = schemas.ModelProvider(weightedConfigs[0].Provider) + var selectedProvider schemas.ModelProvider + totalWeight := 0.0 + for _, config := range weightedConfigs { + totalWeight += getWeight(config.Weight) + } + // Generate random number between 0 and totalWeight + randomValue := rand.Float64() * totalWeight + // Select provider based on weighted random selection + currentWeight := 0.0 + for _, config := range weightedConfigs { + currentWeight += getWeight(config.Weight) + if randomValue <= currentWeight { + selectedProvider = schemas.ModelProvider(config.Provider) + break } - } else { - // No providers have weight set - return body, nil + } + // Fallback: if no provider was selected (shouldn't happen but guard against FP issues) + if selectedProvider == "" { + selectedProvider = schemas.ModelProvider(weightedConfigs[0].Provider) } p.logger.Debug("[governance] Selected provider: %s", selectedProvider) ctx.AppendRoutingEngineLog(schemas.RoutingEngineGovernance, schemas.LogLevelInfo, fmt.Sprintf("Selected provider %s for model %s (from %d eligible: %v)", selectedProvider, modelStr, len(allowedProviderConfigs), allowedProviders)) - // For genai integration, model is present in URL path instead of the request body - if isGeminiPath { - newModelWithRequestSuffix := string(selectedProvider) + "/" + modelStr + genaiRequestSuffix - ctx.SetValue("model", newModelWithRequestSuffix) - } else if isBedrockPath { - // For bedrock integration, model is present in URL path as modelId - ctx.SetValue("modelId", string(selectedProvider)+"/"+modelStr) - } else { + refinedModel := modelStr + // Refine the model for the selected provider + if p.modelCatalog != nil { var err error - refinedModel := modelStr - // Refine the model for the selected provider - if p.modelCatalog != nil { - refinedModel, err = p.modelCatalog.RefineModelForProvider(selectedProvider, modelStr) - if err != nil { - return body, err - } + refinedModel, err = p.modelCatalog.RefineModelForProvider(selectedProvider, modelStr) + if err != nil { + return err } - // Update the model field in the request body - body["model"] = string(selectedProvider) + "/" + refinedModel } - // Append governance to routing engines used + + req.SetProvider(selectedProvider) + req.SetModel(refinedModel) + schemas.AppendToContextList(ctx, schemas.BifrostContextKeyRoutingEnginesUsed, schemas.RoutingEngineGovernance) - // Check if fallbacks field is already present - _, hasFallbacks := body["fallbacks"] - // Use the same candidate set that was used for primary selection - fallbackConfigs := weightedConfigs - if !hasFallbacks && len(fallbackConfigs) > 1 { - // Sort fallback configs by weight (descending) + if len(existingFallbacks) == 0 && len(weightedConfigs) > 1 { + fallbackConfigs := append([]configstoreTables.TableVirtualKeyProviderConfig(nil), weightedConfigs...) sort.Slice(fallbackConfigs, func(i, j int) bool { return getWeight(fallbackConfigs[i].Weight) > getWeight(fallbackConfigs[j].Weight) }) // Filter out the selected provider and create fallbacks array - fallbacks := make([]string, 0, len(fallbackConfigs)-1) + fallbacks := make([]schemas.Fallback, 0, len(fallbackConfigs)-1) for _, config := range fallbackConfigs { - if config.Provider != string(selectedProvider) { - var err error - refinedModel := modelStr - if p.modelCatalog != nil { - refinedModel, err = p.modelCatalog.RefineModelForProvider(schemas.ModelProvider(config.Provider), modelStr) - if err != nil { - // Skip fallback if model refinement fails - p.logger.Warn("failed to refine model for fallback, skipping fallback in governance plugin: %v", err) - continue - } + if config.Provider == string(selectedProvider) { + continue + } + fbProvider := schemas.ModelProvider(config.Provider) + fbModel := modelStr + if p.modelCatalog != nil { + refined, err := p.modelCatalog.RefineModelForProvider(fbProvider, modelStr) + if err != nil { + p.logger.Warn("failed to refine model for fallback, skipping fallback in governance plugin: %v", err) + continue } - fallbacks = append(fallbacks, string(schemas.ModelProvider(config.Provider))+"/"+refinedModel) + fbModel = refined } + fallbacks = append(fallbacks, schemas.Fallback{Provider: fbProvider, Model: fbModel}) } - - // Add fallbacks to request body - body["fallbacks"] = fallbacks - ctx.AppendRoutingEngineLog(schemas.RoutingEngineGovernance, schemas.LogLevelInfo, fmt.Sprintf("Added %d fallback providers: %v", len(fallbacks), fallbacks)) + req.SetFallbacks(fallbacks) + ctx.AppendRoutingEngineLog(schemas.RoutingEngineGovernance, schemas.LogLevelInfo, fmt.Sprintf("Added %d fallback providers", len(fallbacks))) } - return body, nil + return nil } -// applyRoutingRules evaluates routing rules and returns both the modified payload AND the routing decision -// This allows the caller to determine if marshaling is necessary (only if decision != nil or payload changed) -// Parameters: -// - ctx: Bifrost context -// - req: HTTP request -// - body: Request body (may be modified if routing rule matches) -// - virtualKey: Virtual key configuration (may be nil) -// -// Returns: -// - map[string]any: The potentially modified request body -// - *RoutingDecision: The matched routing decision (nil if no rule matched) -// - error: Any error that occurred during evaluation -func (p *GovernancePlugin) applyRoutingRules(ctx *schemas.BifrostContext, req *schemas.HTTPRequest, body map[string]any, virtualKey *configstoreTables.TableVirtualKey) (map[string]any, *RoutingDecision, error) { - // Check if the request has a model field - modelValue, hasModel := body["model"] - isGeminiPath := strings.Contains(req.Path, "/genai") - isBedrockPath := strings.Contains(req.Path, "/bedrock") - if !hasModel { - // For genai integration, model is present in URL path - if isGeminiPath { - modelValue = req.CaseInsensitivePathParamLookup("model") - } else if isBedrockPath { - // For bedrock integration, model is present in URL path as modelId - rawModelID := req.CaseInsensitivePathParamLookup("modelId") - if rawModelID == "" { - return body, nil, nil - } - // URL-decode the modelId (Bedrock model IDs may be URL-encoded) - decoded, err := url.PathUnescape(rawModelID) - if err != nil { - decoded = rawModelID - } - modelValue = decoded - } else { - return body, nil, nil - } - } - - modelStr, ok := modelValue.(string) - if !ok || modelStr == "" { - return body, nil, nil - } - - var genaiRequestSuffix string - if strings.Contains(req.Path, "/genai") { - for _, sfx := range gemini.GeminiRequestSuffixPaths { - if before, ok := strings.CutSuffix(modelStr, sfx); ok { - modelStr = before - genaiRequestSuffix = sfx - break - } - } +// applyRoutingRules evaluates routing rules against req and mutates +// req.Provider/req.Model/req.Fallbacks when a rule matches. Returns the matched RoutingDecision +// (nil if no rule matched). Integrations normalize req.Model (and Provider when applicable) before +// the BifrostRequest reaches this point. +func (p *GovernancePlugin) applyRoutingRules(ctx *schemas.BifrostContext, req *schemas.BifrostRequest, virtualKey *configstoreTables.TableVirtualKey) (*RoutingDecision, error) { + provider, model, _ := req.GetRequestFields() + if model == "" { + return nil, nil } - // Parse provider and model from modelStr (format: "provider/model" or just "model") - provider, model := schemas.ParseModelString(modelStr, "") - - // Extract normalized request type from context (set by HTTP middleware) - requestType := "" - if val := ctx.Value(schemas.BifrostContextKeyHTTPRequestType); val != nil { - if requestTypeEnum, ok := val.(schemas.RequestType); ok { - requestType = string(requestTypeEnum) - } else if requestTypeStr, ok := val.(string); ok { - requestType = requestTypeStr - } - } + requestType := string(req.RequestType) + headers, _ := ctx.Value(schemas.BifrostContextKeyRequestHeaders).(map[string]string) + queryParams, _ := ctx.Value(schemas.BifrostContextKeyRequestQuery).(map[string]string) - // Build routing context routingCtx := &RoutingContext{ VirtualKey: virtualKey, Provider: provider, Model: model, RequestType: requestType, - Headers: req.Headers, - QueryParams: req.Query, + Headers: headers, + QueryParams: queryParams, BudgetAndRateLimitStatus: p.store.GetBudgetAndRateLimitStatus(ctx, model, provider, virtualKey, nil, nil, nil), } - p.logger.Debug("[HTTPTransport] Built routing context: provider=%s, model=%s, requestType=%s, vk=%v, headerCount=%d, paramCount=%d", - provider, model, requestType, virtualKey != nil, len(req.Headers), len(req.Query)) + p.logger.Debug("[PreRequestHook] Built routing context: provider=%s, model=%s, requestType=%s, vk=%v", + provider, model, requestType, virtualKey != nil) // Evaluate routing rules decision, err := p.engine.EvaluateRoutingRules(ctx, routingCtx) if err != nil { p.logger.Error("failed to evaluate routing rules: %v", err) ctx.AppendRoutingEngineLog(schemas.RoutingEngineRoutingRule, schemas.LogLevelError, fmt.Sprintf("Routing rule evaluation error: %v", err)) - return body, nil, nil + return nil, nil + } + if decision == nil { + return nil, nil } - // If a routing rule matched, apply the decision - if decision != nil { - p.logger.Debug("[Governance] Routing rule matched: %s", decision.MatchedRuleName) + p.logger.Debug("[Governance] Routing rule matched: %s", decision.MatchedRuleName) - // Update model in request body - if strings.Contains(req.Path, "/genai") { - // For genai, model is in URL path - newModel := decision.Model + genaiRequestSuffix - // Add provider prefix if present (because there can be other routing rules down stream that can add the provider) - if decision.Provider != "" { - newModel = decision.Provider + "/" + newModel - } - ctx.SetValue("model", newModel) - } else if isBedrockPath { - // For bedrock, model is in URL path as modelId - // Set new modelId in context so bedrockPreCallback picks it up via ctx.UserValue("modelId") - newModel := decision.Model - if decision.Provider != "" { - newModel = decision.Provider + "/" + newModel + if decision.Provider != "" { + req.SetProvider(schemas.ModelProvider(decision.Provider)) + } + if decision.Model != "" { + req.SetModel(decision.Model) + } + + schemas.AppendToContextList(ctx, schemas.BifrostContextKeyRoutingEnginesUsed, schemas.RoutingEngineRoutingRule) + + // Add fallbacks if present; fill in the incoming model for fallbacks that omit it + if len(decision.Fallbacks) > 0 { + resolvedFallbacks := make([]schemas.Fallback, 0, len(decision.Fallbacks)) + for _, fb := range decision.Fallbacks { + fbProvider, fbModel := schemas.ParseModelString(fb, "") + trimmedFbProvider := strings.TrimSpace(string(fbProvider)) + trimmedFbModel := strings.TrimSpace(fbModel) + if trimmedFbProvider == "" { + continue } - ctx.SetValue("modelId", newModel) - } else { - // For regular requests, update in body - newModel := decision.Model - // Add provider prefix if present (because there can be other routing rules down stream that can add the provider) - if decision.Provider != "" { - newModel = decision.Provider + "/" + newModel + if trimmedFbModel == "" && model != "" { + trimmedFbModel = model } - body["model"] = newModel - } - // Append routing-rule to routing engines used - schemas.AppendToContextList(ctx, schemas.BifrostContextKeyRoutingEnginesUsed, schemas.RoutingEngineRoutingRule) - - // Add fallbacks if present; fill in the incoming model for fallbacks that omit it - if len(decision.Fallbacks) > 0 { - resolvedFallbacks := make([]string, 0, len(decision.Fallbacks)) - for _, fb := range decision.Fallbacks { - fbProvider, fbModel := schemas.ParseModelString(fb, "") - trimmedFbProvider := strings.TrimSpace(string(fbProvider)) - trimmedFbModel := strings.TrimSpace(fbModel) - if trimmedFbProvider == "" { - continue - } - if trimmedFbModel == "" && model != "" { - resolvedFallbacks = append(resolvedFallbacks, trimmedFbProvider+"/"+model) - } else { - resolvedFallbacks = append(resolvedFallbacks, trimmedFbProvider+"/"+trimmedFbModel) - } - } - body["fallbacks"] = resolvedFallbacks - } - - // Pin specific API key by ID if the routing rule specifies one - if decision.KeyID != "" { - ctx.SetValue(schemas.BifrostContextKeyAPIKeyID, decision.KeyID) + resolvedFallbacks = append(resolvedFallbacks, schemas.Fallback{ + Provider: schemas.ModelProvider(trimmedFbProvider), + Model: trimmedFbModel, + }) } + req.SetFallbacks(resolvedFallbacks) + } - p.logger.Debug("[Governance] Applied routing decision: provider=%s, model=%s, keyID=%s, fallbacks=%v", decision.Provider, decision.Model, decision.KeyID, decision.Fallbacks) + // Pin specific API key by ID if the routing rule specifies one + if decision.KeyID != "" { + ctx.SetValue(schemas.BifrostContextKeyAPIKeyID, decision.KeyID) } - return body, decision, nil + p.logger.Debug("[Governance] Applied routing decision: provider=%s, model=%s, keyID=%s, fallbacks=%v", decision.Provider, decision.Model, decision.KeyID, decision.Fallbacks) + return decision, nil } -// addMCPIncludeTools adds the x-bf-mcp-include-tools header to the request headers -// Parameters: -// - headers: The request headers -// - virtualKey: The virtual key configuration -// -// Returns: -// - map[string]string: The updated request headers -// - error: Any error that occurred during processing -func (p *GovernancePlugin) addMCPIncludeTools(headers map[string]string, virtualKey *configstoreTables.TableVirtualKey) (map[string]string, error) { - if headers == nil { - headers = make(map[string]string) - } - +// computeMCPIncludeTools builds the MCP include-tools list for a virtual key. Returns the list +// directly; callers store it via ctx.SetValue(schemas.MCPContextKeyIncludeTools, ...). VK-specific +// MCP configs take precedence over AllowOnAllVirtualKeys clients. +func (p *GovernancePlugin) computeMCPIncludeTools(virtualKey *configstoreTables.TableVirtualKey) []string { executeOnlyTools := make([]string, 0) // Build a lookup of AllowOnAllVirtualKeys clients: clientID -> clientName @@ -1145,39 +721,7 @@ func (p *GovernancePlugin) addMCPIncludeTools(headers map[string]string, virtual } } - // Set even when empty to exclude tools when no tools are present in the virtual key config - headers["x-bf-mcp-include-tools"] = strings.Join(executeOnlyTools, ",") - - return headers, nil -} - -// validateRequiredHeaders checks that all configured required headers are present in the request. -// Headers are compared case-insensitively (both sides lowercased). -// Returns a BifrostError with status 400 if any required headers are missing, or nil if all present. -func (p *GovernancePlugin) validateRequiredHeaders(ctx *schemas.BifrostContext) *schemas.BifrostError { - if p.requiredHeaders == nil || len(*p.requiredHeaders) == 0 { - return nil - } - headers, _ := ctx.Value(schemas.BifrostContextKeyRequestHeaders).(map[string]string) - if headers == nil { - headers = map[string]string{} - } - var missing []string - for _, h := range *p.requiredHeaders { - if _, ok := headers[strings.ToLower(h)]; !ok { - missing = append(missing, h) - } - } - if len(missing) > 0 { - return &schemas.BifrostError{ - Type: bifrost.Ptr("missing_required_headers"), - StatusCode: bifrost.Ptr(400), - Error: &schemas.ErrorField{ - Message: fmt.Sprintf("missing required headers: %s", strings.Join(missing, ", ")), - }, - } - } - return nil + return executeOnlyTools } // EvaluateGovernanceRequest is a common function that handles virtual key validation @@ -1440,8 +984,77 @@ func (p *GovernancePlugin) isMCPToolAllowedByVKWith(vk *configstoreTables.TableV return false } -// PreRequestHook implements schemas.LLMPlugin (no-op — required for plugin indexing). -func (p *GovernancePlugin) PreRequestHook(_ *schemas.BifrostContext, _ *schemas.BifrostRequest) error { +// PreRequestHook is the per-request governance phase. It runs for both normal body-having +// requests (route on req.Model) and large-payload streaming requests (route on +// LargePayloadMetadata.Model from ctx — the body is opaque mid-stream, so routing is +// constrained to same-protocol-family targets that the upstream provider can hydrate +// from the rewritten metadata). +// +// Realtime + generic streaming bypass handleRequest (see core/bifrost.go +// RunRealtimeTurnPreHooks / RunStreamPreHooks) and are still handled at HTTPTransportPreHook. +func (p *GovernancePlugin) PreRequestHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) error { + if req.RequestType == schemas.PassthroughRequest || req.RequestType == schemas.PassthroughStreamRequest { + return nil + } + + virtualKeyValue := bifrost.GetStringFromContext(ctx, schemas.BifrostContextKeyVirtualKey) + hasRoutingRules := p.store.HasRoutingRules(ctx) + if virtualKeyValue == "" && !hasRoutingRules { + return nil + } + + var virtualKey *configstoreTables.TableVirtualKey + if virtualKeyValue != "" { + var ok bool + virtualKey, ok = p.store.GetVirtualKey(ctx, virtualKeyValue) + if !ok || virtualKey == nil || !virtualKey.IsActiveValue() { + return nil + } + } + + stampGovernanceCtxFromVK(ctx, virtualKey) + + // Large-payload mode: the body streams to the provider unparsed, so req.Model is + // empty for routes where the model lives in the body (OpenAI/Anthropic chat, + // responses, etc.). Route on LargePayloadMetadata.Model — the provider's + // streaming body rewriter (ApplyLargePayloadRequestBodyWithModelNormalization) + // reads metadata.Model when it rewrites the model field in the body prefix, so + // mutating it here is what propagates the routing decision to the upstream call. + if metadata, _ := ctx.Value(schemas.BifrostContextKeyLargePayloadMetadata).(*schemas.LargePayloadMetadata); metadata != nil && metadata.Model != "" { + newModel, err := p.runPreRequestRouting(ctx, virtualKey, hasRoutingRules, metadata.Model, req.RequestType) + if err != nil { + return err + } + if newModel != "" && newModel != metadata.Model { + metadata.Model = newModel + } + return nil + } + + if hasRoutingRules { + if _, err := p.applyRoutingRules(ctx, req, virtualKey); err != nil { + return err + } + } + + if virtualKey != nil { + if err := p.loadBalanceProvider(ctx, req, virtualKey); err != nil { + return err + } + + p.cfgMutex.RLock() + autoInjectDisabled := p.disableAutoToolInject != nil && *p.disableAutoToolInject + p.cfgMutex.RUnlock() + if !autoInjectDisabled { + // Don't overwrite a caller-provided include-tools value (set via header in lib/ctx.go). + if existing := ctx.Value(schemas.MCPContextKeyIncludeTools); existing == nil { + if tools := p.computeMCPIncludeTools(virtualKey); tools != nil { + ctx.SetValue(schemas.MCPContextKeyIncludeTools, tools) + } + } + } + } + return nil } diff --git a/plugins/governance/utils.go b/plugins/governance/utils.go index b24940d343..95496148d9 100644 --- a/plugins/governance/utils.go +++ b/plugins/governance/utils.go @@ -3,11 +3,13 @@ package governance import ( "context" + "fmt" "slices" "strings" bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" + configstoreTables "github.com/maximhq/bifrost/framework/configstore/tables" "github.com/valyala/fasthttp" ) @@ -53,43 +55,6 @@ func IsModelRequiredForRequest(requestType schemas.RequestType) bool { return true } -// parseVirtualKeyFromHTTPRequest parses the virtual key from HTTP request headers. -// It checks multiple headers in order: x-bf-vk, Authorization (Bearer token), x-api-key, and x-goog-api-key. -// Parameters: -// - req: The HTTP request containing headers to parse -// -// Returns: -// - *string: The virtual key if found, nil otherwise -func parseVirtualKeyFromHTTPRequest(req *schemas.HTTPRequest) *string { - var virtualKeyValue string - vkHeader := req.CaseInsensitiveHeaderLookup("x-bf-vk") - if vkHeader != "" && strings.HasPrefix(strings.ToLower(vkHeader), VirtualKeyPrefix) { - return bifrost.Ptr(vkHeader) - } - authHeader := req.CaseInsensitiveHeaderLookup("Authorization") - if authHeader != "" { - if strings.HasPrefix(strings.ToLower(authHeader), "bearer ") { - authHeaderValue := strings.TrimSpace(authHeader[7:]) // Remove "Bearer " prefix - if authHeaderValue != "" && strings.HasPrefix(strings.ToLower(authHeaderValue), VirtualKeyPrefix) { - virtualKeyValue = authHeaderValue - } - } - } - if virtualKeyValue != "" { - return bifrost.Ptr(virtualKeyValue) - } - xAPIKey := req.CaseInsensitiveHeaderLookup("x-api-key") - if xAPIKey != "" && strings.HasPrefix(strings.ToLower(xAPIKey), VirtualKeyPrefix) { - return bifrost.Ptr(xAPIKey) - } - // Checking x-goog-api-key header - xGoogleAPIKey := req.CaseInsensitiveHeaderLookup("x-goog-api-key") - if xGoogleAPIKey != "" && strings.HasPrefix(strings.ToLower(xGoogleAPIKey), VirtualKeyPrefix) { - return bifrost.Ptr(xGoogleAPIKey) - } - return nil -} - // getWeight safely dereferences a *float64 weight pointer, returning 1.0 as default if nil. // This allows distinguishing between "not set" (nil -> 1.0) and "explicitly set to 0" (0.0). func getWeight(w *float64) float64 { @@ -129,6 +94,33 @@ func isModelBlockedByList(blacklist schemas.BlackList, model string) bool { return false } +// stampGovernanceCtxFromVK copies team/customer identifiers from the VK onto ctx so +// downstream plugins (logging, observability) see the governance scope. +func stampGovernanceCtxFromVK(ctx *schemas.BifrostContext, vk *configstoreTables.TableVirtualKey) { + if vk == nil { + return + } + if vk.TeamID != nil { + ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamID, *vk.TeamID) + } + if vk.Team != nil { + ctx.SetValue(schemas.BifrostContextKeyGovernanceTeamName, vk.Team.Name) + if vk.Team.CustomerID != nil { + ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerID, *vk.Team.CustomerID) + if vk.Team.Customer != nil { + ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerName, vk.Team.Customer.Name) + } + } + } else { + if vk.CustomerID != nil { + ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerID, *vk.CustomerID) + } + if vk.Customer != nil { + ctx.SetValue(schemas.BifrostContextKeyGovernanceCustomerName, vk.Customer.Name) + } + } +} + // filterModelsForVirtualKey filters models based on virtual key's provider configs // Returns only models that are allowed by the virtual key's ProviderConfigs func (p *GovernancePlugin) filterModelsForVirtualKey( @@ -195,3 +187,32 @@ func (p *GovernancePlugin) filterModelsForVirtualKey( return filteredModels } + +// validateRequiredHeaders checks that all configured required headers are present in the request. +// Headers are compared case-insensitively (both sides lowercased). +// Returns a BifrostError with status 400 if any required headers are missing, or nil if all present. +func (p *GovernancePlugin) validateRequiredHeaders(ctx *schemas.BifrostContext) *schemas.BifrostError { + if p.requiredHeaders == nil || len(*p.requiredHeaders) == 0 { + return nil + } + headers, _ := ctx.Value(schemas.BifrostContextKeyRequestHeaders).(map[string]string) + if headers == nil { + headers = map[string]string{} + } + var missing []string + for _, h := range *p.requiredHeaders { + if _, ok := headers[strings.ToLower(h)]; !ok { + missing = append(missing, h) + } + } + if len(missing) > 0 { + return &schemas.BifrostError{ + Type: bifrost.Ptr("missing_required_headers"), + StatusCode: bifrost.Ptr(400), + Error: &schemas.ErrorField{ + Message: fmt.Sprintf("missing required headers: %s", strings.Join(missing, ", ")), + }, + } + } + return nil +} diff --git a/transports/bifrost-http/handlers/wsrealtime.go b/transports/bifrost-http/handlers/wsrealtime.go index 81edd1496f..856ded4746 100644 --- a/transports/bifrost-http/handlers/wsrealtime.go +++ b/transports/bifrost-http/handlers/wsrealtime.go @@ -92,6 +92,68 @@ func (h *WSRealtimeHandler) handleUpgrade(ctx *fasthttp.RequestCtx) { return } + // Run PreRequestHook to give governance + LB a chance to route the realtime connection. + // Realtime bypasses handleRequest (per-turn pipelines instead), so we invoke the routing + // phase explicitly here. Mutations to provider/model are read back into the local vars + // and copied to fasthttp user values so snapshotRealtimeMiddlewareValues picks up any + // ctx changes (governance team/customer IDs, routing engine logs). + preReqCtx, preReqCancel := createBifrostContextFromAuth(h.handlerStore, auth) + if preReqCtx == nil { + preReqCancel() + upgrader := h.websocketUpgrader("") + upgradeErr := upgrader.Upgrade(ctx, func(conn *ws.Conn) { + defer conn.Close() + clientConn := newRealtimeClientConn(conn) + clientConn.writeRealtimeError(newRealtimeWireBifrostError(500, "server_error", "failed to create request context")) + }) + if upgradeErr != nil { + logger.Warn("websocket upgrade failed for %s: %v", path, upgradeErr) + } + return + } + preReqCtx.SetValue(schemas.BifrostContextKeyHTTPRequestType, schemas.RealtimeRequest) + if realtimeDefaultProviderForPath(path) == schemas.OpenAI { + preReqCtx.SetValue(schemas.BifrostContextKeyIntegrationType, "openai") + } + // Surface full request headers + query params on the pre-request context so governance + // CEL routing rules (which read headers[...] / params[...]) see the same shape they would + // for normal HTTP requests. Mirrors lib/ctx.go ConvertToBifrostContext; the normal HTTP + // path doesn't run for WS upgrades, so we populate these explicitly. Keys are lowercased. + allHeaders := make(map[string]string) + ctx.Request.Header.All()(func(key, value []byte) bool { + allHeaders[strings.ToLower(string(key))] = string(value) + return true + }) + preReqCtx.SetValue(schemas.BifrostContextKeyRequestHeaders, allHeaders) + if queryArgs := ctx.Request.URI().QueryArgs(); queryArgs.Len() > 0 { + allQuery := make(map[string]string, queryArgs.Len()) + queryArgs.All()(func(key, value []byte) bool { + allQuery[strings.ToLower(string(key))] = string(value) + return true + }) + preReqCtx.SetValue(schemas.BifrostContextKeyRequestQuery, allQuery) + } + preReq := &schemas.BifrostRequest{ + RequestType: schemas.RealtimeRequest, + ResponsesRequest: &schemas.BifrostResponsesRequest{ + Provider: providerKey, + Model: model, + }, + } + h.client.RunPreRequestHooks(preReqCtx, preReq) + if routedProvider, routedModel, _ := preReq.GetRequestFields(); routedProvider != "" { + providerKey = routedProvider + if routedModel != "" { + model = routedModel + } + } + // Mirror ctx values back to fasthttp user values so snapshotRealtimeMiddlewareValues + // (called below) picks them up — same mechanism TransportInterceptorMiddleware uses. + for k, v := range preReqCtx.GetUserValues() { + ctx.SetUserValue(k, v) + } + preReqCancel() + provider := h.client.GetProviderByKey(providerKey) rtProvider, ok := provider.(schemas.RealtimeProvider) if provider == nil || !ok || !rtProvider.SupportsRealtimeAPI() { diff --git a/transports/bifrost-http/lib/ctx.go b/transports/bifrost-http/lib/ctx.go index e94bbeabe1..79e1d50f7f 100644 --- a/transports/bifrost-http/lib/ctx.go +++ b/transports/bifrost-http/lib/ctx.go @@ -621,6 +621,18 @@ func ConvertToBifrostContext(ctx *fasthttp.RequestCtx, store HandlerStore) (*sch }) bifrostCtx.SetValue(schemas.BifrostContextKeyRequestHeaders, allHeaders) + // Collect all request query params for downstream use (e.g., governance routing CEL rules + // that read params["..."]). Keys are lowercased for case-insensitive lookup. + queryArgs := ctx.Request.URI().QueryArgs() + if queryArgs.Len() > 0 { + allQuery := make(map[string]string, queryArgs.Len()) + queryArgs.All()(func(key, value []byte) bool { + allQuery[strings.ToLower(string(key))] = string(value) + return true + }) + bifrostCtx.SetValue(schemas.BifrostContextKeyRequestQuery, allQuery) + } + // Build and set the MCP callback base URL. Used by per-user OAuth (appends // /api/oauth/callback) and per-user headers (appends the workspace submit // path) resolvers when initiating their respective auth flows. Bifrost is From 5282def56e9d33547f5355fc1a8cff904785c889 Mon Sep 17 00:00:00 2001 From: Pratham Mishra <99235987+Pratham-Mishra04@users.noreply.github.com> Date: Tue, 9 Jun 2026 15:47:45 +0530 Subject: [PATCH 004/108] refactor: extract provider resolution into `modelcatalogresolver` PreRequestHook plugin (#4177) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Provider resolution for unprefixed model strings (e.g. `gpt-4o` without a `provider/` prefix) was previously scattered across every integration router, every provider's `ToBifrost*` converter, and the `CheckAndSetDefaultProvider` utility — each doing its own inline catalog lookup or context-key dance. This PR consolidates all of that into a single, dedicated `modelcatalogresolver` built-in plugin that runs as the last `PreRequestHook` before the request reaches the LLM layer. ## Changes - **New `plugins/modelcatalogresolver` plugin**: A `PreRequestHook` that fills `req.Provider` from the model catalog when no provider was specified and no earlier routing plugin (governance routing rules, governance VK load balancing, enterprise LB) already set one. It also promotes remaining catalog candidates to fallbacks automatically when the caller didn't configure any. An integration-type hint (`BifrostContextKeyIntegrationType`) biases the pick toward the integration's canonical provider; Azure user-agent detection is handled here instead of inline in the OpenAI converters. - **Removed `CheckAndSetDefaultProvider`**: The `providerUtils.CheckAndSetDefaultProvider` helper and its `BifrostContextKeyAvailableProviders` / `BifrostContextKeyResolvedProvider` context keys are deleted. All `ToBifrost*` converters across Anthropic, Bedrock, Cohere, Gemini, OpenAI, and Vertex now pass `""` as the default provider to `ParseModelString`, deferring resolution entirely to the plugin layer. - **Removed `GetRequestModel` / `GetProvidersForModel` from integration routers**: The per-route `RequestModelGetter` callbacks and the inline catalog-lookup block in `GenericRouter.createHandler` are removed. The `HandlerStore` interface no longer requires `GetProvidersForModel`, and `Config.GetProvidersForModel` is deleted. The governance plugin no longer sets `BifrostContextKeyAvailableProviders`. - **Simplified `resolveModelAndProvider` and `resolveRealtimeTarget`**: These functions in the inference and WebSocket realtime handlers no longer do inline catalog lookups. The realtime handlers (`webrtc_realtime.go`, `realtime_client_secrets.go`) that couldn't go through `PreRequestHook` now call the exported `modelcatalogresolver.ResolveProviderFromCatalog` directly. - **WebSocket realtime empty-provider guard**: `wsrealtime.go`'s `handleUpgrade` now returns a clear WebSocket error when no routing layer could resolve a provider, mirroring the empty-provider validation in `handleRequest`/`handleStreamRequest`. - **Observability consolidation**: `EmitModelCatalogRoutingLog` is extracted into `lib/ctx.go` so all paths (normal HTTP, WebRTC, realtime client secrets) emit routing engine logs in the same shape. The `snapshotRealtimeMiddlewareValues` function no longer duplicates this logic. - **Governance plugin cleanup**: Removed the `BifrostContextKeyAvailableProviders` writes and the fallback-filtering logic from `extractAndParseFallbacks` that depended on them. - **`BifrostContextKeySkipModelCatalogProviderSelection` removed**: The context key and all references are deleted; the resolver plugin's position as the last hook makes the skip flag unnecessary. ## Type of change - [ ] Bug fix - [x] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [x] Providers/Integrations - [x] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./core/... go test ./transports/... go test ./plugins/governance/... go test ./plugins/modelcatalogresolver/... ``` Verify that requests with unprefixed model strings (e.g. `{"model":"gpt-4o"}` on `/v1/chat/completions`) still resolve to the correct provider when a model catalog is configured. Verify that requests with explicit prefixes (e.g. `{"model":"anthropic/claude-opus-4"}`) are unaffected. Verify that WebSocket realtime connections with unresolvable models receive a `400 invalid_request_error` frame instead of hanging. ## Breaking changes - [x] Yes - [ ] No `HandlerStore.GetProvidersForModel` is removed from the interface — any custom `HandlerStore` implementations must drop this method. `BifrostContextKeyAvailableProviders`, `BifrostContextKeyResolvedProvider`, and `BifrostContextKeySkipModelCatalogProviderSelection` context keys are removed from `schemas`; any code reading or writing these keys must be updated. `CheckAndSetDefaultProvider` is removed from `core/providers/utils`. ## Related issues ## Security considerations No auth, secrets, or PII changes. Provider resolution is now centralised in a single plugin rather than distributed across converters, reducing the surface area for routing bypasses. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable ## Summary by CodeRabbit ## Release Notes * **New Features** * Introduced model catalog resolver for intelligent provider selection based on model availability and integration type. * Added integration-aware provider preferences (e.g., Azure-OpenAI pairing). * **Bug Fixes** * Improved model-to-provider resolution to eliminate context-dependent provider defaults. * **Refactor** * Streamlined provider resolution logic across all AI provider integrations. * Simplified request routing configuration by removing model getter dependencies. --- core/providers/anthropic/responses.go | 2 +- core/providers/anthropic/text.go | 2 +- core/providers/bedrock/invoke.go | 8 +- core/providers/bedrock/rerank.go | 3 +- core/providers/bedrock/responses.go | 2 +- core/providers/bedrock/text.go | 13 +- core/providers/cohere/chat.go | 2 +- core/providers/cohere/count_tokens.go | 3 +- core/providers/cohere/embedding.go | 3 +- core/providers/cohere/rerank.go | 3 +- core/providers/gemini/embedding.go | 3 +- core/providers/gemini/images.go | 4 +- core/providers/gemini/responses.go | 2 +- core/providers/gemini/speech.go | 2 +- core/providers/gemini/transcription.go | 2 +- core/providers/gemini/videos.go | 2 +- core/providers/openai/chat.go | 2 +- core/providers/openai/embedding.go | 3 +- core/providers/openai/images.go | 6 +- core/providers/openai/responses.go | 21 +- core/providers/openai/speech.go | 3 +- core/providers/openai/text.go | 3 +- core/providers/openai/transcription.go | 2 +- core/providers/openai/videos.go | 28 +- core/providers/utils/utils.go | 30 -- core/providers/utils/utils_test.go | 24 -- core/providers/vertex/rerank.go | 5 +- core/schemas/bifrost.go | 67 ++-- plugins/governance/go.mod | 3 +- .../governance/httptransportprehook_test.go | 346 ------------------ plugins/governance/main.go | 4 - plugins/modelcatalogresolver/go.mod | 159 ++++++++ plugins/modelcatalogresolver/go.sum | 256 +++++++++++++ plugins/modelcatalogresolver/main.go | 191 ++++++++++ transports/bifrost-http/handlers/inference.go | 26 +- .../handlers/realtime_client_secrets.go | 11 +- .../bifrost-http/handlers/webrtc_realtime.go | 31 +- .../handlers/webrtc_realtime_test.go | 3 +- .../bifrost-http/handlers/wsrealtime.go | 83 ++--- .../bifrost-http/handlers/wsresponses_test.go | 4 - .../bifrost-http/integrations/anthropic.go | 20 +- .../bifrost-http/integrations/bedrock.go | 22 -- .../bifrost-http/integrations/bedrock_test.go | 4 - .../bifrost-http/integrations/cohere.go | 20 - transports/bifrost-http/integrations/genai.go | 35 -- .../bifrost-http/integrations/openai.go | 66 +--- .../bifrost-http/integrations/router.go | 87 +---- .../bifrost-http/integrations/router_test.go | 92 +---- transports/bifrost-http/integrations/utils.go | 10 - .../bifrost-http/integrations/utils_test.go | 53 --- transports/bifrost-http/lib/config.go | 24 -- transports/bifrost-http/lib/ctx.go | 35 +- transports/bifrost-http/server/plugins.go | 15 + transports/go.mod | 1 + transports/go.sum | 1 + 55 files changed, 797 insertions(+), 1055 deletions(-) delete mode 100644 plugins/governance/httptransportprehook_test.go create mode 100644 plugins/modelcatalogresolver/go.mod create mode 100644 plugins/modelcatalogresolver/go.sum create mode 100644 plugins/modelcatalogresolver/main.go diff --git a/core/providers/anthropic/responses.go b/core/providers/anthropic/responses.go index 85ddf8f8df..d9dfe2d2d9 100644 --- a/core/providers/anthropic/responses.go +++ b/core/providers/anthropic/responses.go @@ -2187,7 +2187,7 @@ func ToAnthropicResponsesStreamResponse(ctx *schemas.BifrostContext, bifrostResp // ToBifrostResponsesRequest converts an Anthropic message request to Bifrost format func (req *AnthropicMessageRequest) ToBifrostResponsesRequest(ctx *schemas.BifrostContext) *schemas.BifrostResponsesRequest { - provider, model := schemas.ParseModelString(req.Model, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Anthropic)) + provider, model := schemas.ParseModelString(req.Model, "") bifrostReq := &schemas.BifrostResponsesRequest{ Provider: provider, diff --git a/core/providers/anthropic/text.go b/core/providers/anthropic/text.go index 39a700499b..df6f488d7c 100644 --- a/core/providers/anthropic/text.go +++ b/core/providers/anthropic/text.go @@ -54,7 +54,7 @@ func (req *AnthropicTextRequest) ToBifrostTextCompletionRequest(ctx *schemas.Bif return nil } - provider, model := schemas.ParseModelString(req.Model, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Anthropic)) + provider, model := schemas.ParseModelString(req.Model, "") bifrostReq := &schemas.BifrostTextCompletionRequest{ Provider: provider, diff --git a/core/providers/bedrock/invoke.go b/core/providers/bedrock/invoke.go index 8227e8639a..65a8665af9 100644 --- a/core/providers/bedrock/invoke.go +++ b/core/providers/bedrock/invoke.go @@ -387,7 +387,7 @@ func (r *BedrockInvokeRequest) ToBifrostEmbeddingRequest(ctx *schemas.BifrostCon if unescaped, err := url.PathUnescape(r.ModelID); err == nil { modelID = unescaped } - provider, model := schemas.ParseModelString(modelID, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Bedrock)) + provider, model := schemas.ParseModelString(modelID, "") req := &schemas.BifrostEmbeddingRequest{ Provider: provider, Model: model, @@ -451,7 +451,7 @@ func (r *BedrockInvokeRequest) ToBifrostImageGenerationRequest(ctx *schemas.Bifr if unescaped, err := url.PathUnescape(r.ModelID); err == nil { modelID = unescaped } - provider, model := schemas.ParseModelString(modelID, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Bedrock)) + provider, model := schemas.ParseModelString(modelID, "") req := &schemas.BifrostImageGenerationRequest{ Provider: provider, Model: model, @@ -515,7 +515,7 @@ func (r *BedrockInvokeRequest) ToBifrostImageEditRequest(ctx *schemas.BifrostCon if unescaped, err := url.PathUnescape(r.ModelID); err == nil { modelID = unescaped } - provider, model := schemas.ParseModelString(modelID, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Bedrock)) + provider, model := schemas.ParseModelString(modelID, "") req := &schemas.BifrostImageEditRequest{ Provider: provider, Model: model, @@ -690,7 +690,7 @@ func (r *BedrockInvokeRequest) ToBifrostImageVariationRequest(ctx *schemas.Bifro if unescaped, err := url.PathUnescape(r.ModelID); err == nil { modelID = unescaped } - provider, model := schemas.ParseModelString(modelID, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Bedrock)) + provider, model := schemas.ParseModelString(modelID, "") req := &schemas.BifrostImageVariationRequest{ Provider: provider, Model: model, diff --git a/core/providers/bedrock/rerank.go b/core/providers/bedrock/rerank.go index 3ba99441c6..82bbd4554c 100644 --- a/core/providers/bedrock/rerank.go +++ b/core/providers/bedrock/rerank.go @@ -5,7 +5,6 @@ import ( "sort" "strings" - providerUtils "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" ) @@ -134,7 +133,7 @@ func (req *BedrockRerankRequest) ToBifrostRerankRequest(ctx *schemas.BifrostCont } modelARN := req.RerankingConfiguration.BedrockRerankingConfiguration.ModelConfiguration.ModelARN - provider, model := schemas.ParseModelString(modelARN, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Bedrock)) + provider, model := schemas.ParseModelString(modelARN, "") bifrostReq := &schemas.BifrostRerankRequest{ Provider: provider, diff --git a/core/providers/bedrock/responses.go b/core/providers/bedrock/responses.go index 905deb0a19..e081f6bacc 100644 --- a/core/providers/bedrock/responses.go +++ b/core/providers/bedrock/responses.go @@ -1850,7 +1850,7 @@ func (request *BedrockConverseRequest) ToBifrostResponsesRequest(ctx *schemas.Bi } // Extract provider from model ID (format: "bedrock/model-name") - provider, model := schemas.ParseModelString(request.ModelID, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Bedrock)) + provider, model := schemas.ParseModelString(request.ModelID, "") bifrostReq := &schemas.BifrostResponsesRequest{ Provider: provider, diff --git a/core/providers/bedrock/text.go b/core/providers/bedrock/text.go index d31d716ded..b019d2fee2 100644 --- a/core/providers/bedrock/text.go +++ b/core/providers/bedrock/text.go @@ -4,7 +4,6 @@ import ( "strings" "github.com/maximhq/bifrost/core/providers/anthropic" - "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" ) @@ -80,7 +79,7 @@ func (request *BedrockTextCompletionRequest) ToBifrostTextCompletionRequest(ctx prompt = strings.Join(parts, "\n\n") } - provider, model := schemas.ParseModelString(request.ModelID, utils.CheckAndSetDefaultProvider(ctx, schemas.Bedrock)) + provider, model := schemas.ParseModelString(request.ModelID, "") bifrostReq := &schemas.BifrostTextCompletionRequest{ Provider: provider, @@ -126,8 +125,7 @@ func (response *BedrockAnthropicTextResponse) ToBifrostTextCompletionResponse() FinishReason: &response.StopReason, }, }, - ExtraFields: schemas.BifrostResponseExtraFields{ - }, + ExtraFields: schemas.BifrostResponseExtraFields{}, } } @@ -149,10 +147,9 @@ func (response *BedrockMistralTextResponse) ToBifrostTextCompletionResponse() *s } return &schemas.BifrostTextCompletionResponse{ - Object: "text_completion", - Choices: choices, - ExtraFields: schemas.BifrostResponseExtraFields{ - }, + Object: "text_completion", + Choices: choices, + ExtraFields: schemas.BifrostResponseExtraFields{}, } } diff --git a/core/providers/cohere/chat.go b/core/providers/cohere/chat.go index 208d8fcaf1..9865516d91 100644 --- a/core/providers/cohere/chat.go +++ b/core/providers/cohere/chat.go @@ -250,7 +250,7 @@ func (req *CohereChatRequest) ToBifrostChatRequest(ctx *schemas.BifrostContext) return nil } - provider, model := schemas.ParseModelString(req.Model, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Cohere)) + provider, model := schemas.ParseModelString(req.Model, "") bifrostReq := &schemas.BifrostChatRequest{ Provider: provider, diff --git a/core/providers/cohere/count_tokens.go b/core/providers/cohere/count_tokens.go index 0a5e1b48e1..1ffa1a5005 100644 --- a/core/providers/cohere/count_tokens.go +++ b/core/providers/cohere/count_tokens.go @@ -5,7 +5,6 @@ import ( "strings" "unicode/utf8" - "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" ) @@ -15,7 +14,7 @@ func (req *CohereCountTokensRequest) ToBifrostResponsesRequest(ctx *schemas.Bifr return nil } - provider, model := schemas.ParseModelString(req.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.Cohere)) + provider, model := schemas.ParseModelString(req.Model, "") userRole := schemas.ResponsesInputMessageRoleUser return &schemas.BifrostResponsesRequest{ diff --git a/core/providers/cohere/embedding.go b/core/providers/cohere/embedding.go index a99ef14294..0f8976dcb5 100644 --- a/core/providers/cohere/embedding.go +++ b/core/providers/cohere/embedding.go @@ -1,7 +1,6 @@ package cohere import ( - "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" ) @@ -74,7 +73,7 @@ func (req *CohereEmbeddingRequest) ToBifrostEmbeddingRequest(ctx *schemas.Bifros return nil } - provider, model := schemas.ParseModelString(req.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.Cohere)) + provider, model := schemas.ParseModelString(req.Model, "") bifrostReq := &schemas.BifrostEmbeddingRequest{ Provider: provider, diff --git a/core/providers/cohere/rerank.go b/core/providers/cohere/rerank.go index b820e3796b..bf3f759c2d 100644 --- a/core/providers/cohere/rerank.go +++ b/core/providers/cohere/rerank.go @@ -4,7 +4,6 @@ import ( "sort" "github.com/bytedance/sonic" - "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" "gopkg.in/yaml.v3" ) @@ -43,7 +42,7 @@ func (req *CohereRerankRequest) ToBifrostRerankRequest(ctx *schemas.BifrostConte return nil } - provider, model := schemas.ParseModelString(req.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.Cohere)) + provider, model := schemas.ParseModelString(req.Model, "") bifrostReq := &schemas.BifrostRerankRequest{ Provider: provider, diff --git a/core/providers/gemini/embedding.go b/core/providers/gemini/embedding.go index 906b995c68..438f0dee65 100644 --- a/core/providers/gemini/embedding.go +++ b/core/providers/gemini/embedding.go @@ -1,7 +1,6 @@ package gemini import ( - "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" ) @@ -189,7 +188,7 @@ func (request *GeminiGenerationRequest) ToBifrostEmbeddingRequest(ctx *schemas.B return nil } - provider, model := schemas.ParseModelString(request.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.Gemini)) + provider, model := schemas.ParseModelString(request.Model, "") // Create the embedding request bifrostReq := &schemas.BifrostEmbeddingRequest{ diff --git a/core/providers/gemini/images.go b/core/providers/gemini/images.go index c94d33a327..7da0abb7db 100644 --- a/core/providers/gemini/images.go +++ b/core/providers/gemini/images.go @@ -20,7 +20,7 @@ func (request *GeminiGenerationRequest) ToBifrostImageGenerationRequest(ctx *sch // Parse provider from model string (e.g., "openai/gpt-image-1" -> provider="openai", model="gpt-image-1") // This allows cross-provider routing through the GenAI endpoint - provider, model := schemas.ParseModelString(request.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.Gemini)) + provider, model := schemas.ParseModelString(request.Model, "") bifrostReq := &schemas.BifrostImageGenerationRequest{ Provider: provider, @@ -114,7 +114,7 @@ func (request *GeminiGenerationRequest) ToBifrostImageEditRequest(ctx *schemas.B // Parse provider from model string (e.g., "openai/gpt-image-1" -> provider="openai", model="gpt-image-1") // This allows cross-provider routing through the GenAI endpoint - provider, model := schemas.ParseModelString(request.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.Gemini)) + provider, model := schemas.ParseModelString(request.Model, "") bifrostReq := &schemas.BifrostImageEditRequest{ Provider: provider, diff --git a/core/providers/gemini/responses.go b/core/providers/gemini/responses.go index 67b09eb1d9..8a9195e152 100644 --- a/core/providers/gemini/responses.go +++ b/core/providers/gemini/responses.go @@ -18,7 +18,7 @@ func (request *GeminiGenerationRequest) ToBifrostResponsesRequest(ctx *schemas.B return nil } - provider, model := schemas.ParseModelString(request.Model, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Gemini)) + provider, model := schemas.ParseModelString(request.Model, "") // Create the BifrostResponsesRequest bifrostReq := &schemas.BifrostResponsesRequest{ diff --git a/core/providers/gemini/speech.go b/core/providers/gemini/speech.go index d4683c250f..416c30239c 100644 --- a/core/providers/gemini/speech.go +++ b/core/providers/gemini/speech.go @@ -11,7 +11,7 @@ import ( // ToBifrostSpeechRequest converts a GeminiGenerationRequest to a BifrostSpeechRequest func (request *GeminiGenerationRequest) ToBifrostSpeechRequest(ctx *schemas.BifrostContext) *schemas.BifrostSpeechRequest { - provider, model := schemas.ParseModelString(request.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.Gemini)) + provider, model := schemas.ParseModelString(request.Model, "") bifrostReq := &schemas.BifrostSpeechRequest{ Provider: provider, diff --git a/core/providers/gemini/transcription.go b/core/providers/gemini/transcription.go index 0548a3f512..7388b077c3 100644 --- a/core/providers/gemini/transcription.go +++ b/core/providers/gemini/transcription.go @@ -10,7 +10,7 @@ import ( // ToBifrostTranscriptionRequest converts a GeminiGenerationRequest to a BifrostTranscriptionRequest func (request *GeminiGenerationRequest) ToBifrostTranscriptionRequest(ctx *schemas.BifrostContext) (*schemas.BifrostTranscriptionRequest, error) { - provider, model := schemas.ParseModelString(request.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.Gemini)) + provider, model := schemas.ParseModelString(request.Model, "") bifrostReq := &schemas.BifrostTranscriptionRequest{ Provider: provider, diff --git a/core/providers/gemini/videos.go b/core/providers/gemini/videos.go index 3f3c60802e..31b571b4d0 100644 --- a/core/providers/gemini/videos.go +++ b/core/providers/gemini/videos.go @@ -395,7 +395,7 @@ func (request *GeminiVideoGenerationRequest) ToBifrostVideoGenerationRequest(ctx // Use the first instance for the main input instance := request.Instances[0] - provider, model := schemas.ParseModelString(request.Model, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Gemini)) + provider, model := schemas.ParseModelString(request.Model, "") bifrostReq := &schemas.BifrostVideoGenerationRequest{ Provider: provider, diff --git a/core/providers/openai/chat.go b/core/providers/openai/chat.go index af520e7d9c..f198165c95 100644 --- a/core/providers/openai/chat.go +++ b/core/providers/openai/chat.go @@ -9,7 +9,7 @@ import ( // ToBifrostChatRequest converts an OpenAI chat request to Bifrost format func (req *OpenAIChatRequest) ToBifrostChatRequest(ctx *schemas.BifrostContext) *schemas.BifrostChatRequest { - provider, model := schemas.ParseModelString(req.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.OpenAI)) + provider, model := schemas.ParseModelString(req.Model, "") return &schemas.BifrostChatRequest{ Provider: provider, diff --git a/core/providers/openai/embedding.go b/core/providers/openai/embedding.go index fa243ac5b8..586d5f15eb 100644 --- a/core/providers/openai/embedding.go +++ b/core/providers/openai/embedding.go @@ -1,13 +1,12 @@ package openai import ( - "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" ) // ToBifrostEmbeddingRequest converts an OpenAI embedding request to Bifrost format func (request *OpenAIEmbeddingRequest) ToBifrostEmbeddingRequest(ctx *schemas.BifrostContext) *schemas.BifrostEmbeddingRequest { - provider, model := schemas.ParseModelString(request.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.OpenAI)) + provider, model := schemas.ParseModelString(request.Model, "") return &schemas.BifrostEmbeddingRequest{ Provider: provider, diff --git a/core/providers/openai/images.go b/core/providers/openai/images.go index 688df16831..799747fc80 100644 --- a/core/providers/openai/images.go +++ b/core/providers/openai/images.go @@ -56,7 +56,7 @@ func (request *OpenAIImageGenerationRequest) ToBifrostImageGenerationRequest(ctx return nil } - provider, model := schemas.ParseModelString(request.Model, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.OpenAI)) + provider, model := schemas.ParseModelString(request.Model, "") return &schemas.BifrostImageGenerationRequest{ Provider: provider, @@ -74,7 +74,7 @@ func (request *OpenAIImageEditRequest) ToBifrostImageEditRequest(ctx *schemas.Bi return nil } - provider, model := schemas.ParseModelString(request.Model, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.OpenAI)) + provider, model := schemas.ParseModelString(request.Model, "") return &schemas.BifrostImageEditRequest{ Provider: provider, @@ -90,7 +90,7 @@ func (request *OpenAIImageVariationRequest) ToBifrostImageVariationRequest(ctx * return nil } - provider, model := schemas.ParseModelString(request.Model, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.OpenAI)) + provider, model := schemas.ParseModelString(request.Model, "") return &schemas.BifrostImageVariationRequest{ Provider: provider, diff --git a/core/providers/openai/responses.go b/core/providers/openai/responses.go index a6216e0482..feaf4d1107 100644 --- a/core/providers/openai/responses.go +++ b/core/providers/openai/responses.go @@ -13,16 +13,7 @@ func (resp *OpenAIResponsesRequest) ToBifrostResponsesRequest(ctx *schemas.Bifro return nil } - defaultProvider := schemas.OpenAI - - // for requests coming from azure sdk without provider prefix, we need to set the default provider to azure - if ctx != nil { - if isAzureUser, ok := ctx.Value(schemas.BifrostContextKeyIsAzureUserAgent).(bool); ok && isAzureUser { - defaultProvider = schemas.Azure - } - } - - provider, model := schemas.ParseModelString(resp.Model, utils.CheckAndSetDefaultProvider(ctx, defaultProvider)) + provider, model := schemas.ParseModelString(resp.Model, "") input := resp.Input.OpenAIResponsesRequestInputArray if len(input) == 0 { @@ -463,16 +454,8 @@ func (r *OpenAICompactionRequest) ToBifrostCompactionRequest(ctx *schemas.Bifros if r == nil { return nil } - defaultProvider := schemas.OpenAI - - // for requests coming from azure sdk without provider prefix, we need to set the default provider to azure - if ctx != nil { - if isAzureUser, ok := ctx.Value(schemas.BifrostContextKeyIsAzureUserAgent).(bool); ok && isAzureUser { - defaultProvider = schemas.Azure - } - } - provider, model := schemas.ParseModelString(r.Model, utils.CheckAndSetDefaultProvider(ctx, defaultProvider)) + provider, model := schemas.ParseModelString(r.Model, "") input := r.Input.OpenAIResponsesRequestInputArray if len(input) == 0 && r.Input.OpenAIResponsesRequestInputStr != nil { input = []schemas.ResponsesMessage{ diff --git a/core/providers/openai/speech.go b/core/providers/openai/speech.go index 09c638fc5e..0a092e3e7d 100644 --- a/core/providers/openai/speech.go +++ b/core/providers/openai/speech.go @@ -1,13 +1,12 @@ package openai import ( - "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" ) // ToBifrostSpeechRequest converts an OpenAI speech request to Bifrost format func (request *OpenAISpeechRequest) ToBifrostSpeechRequest(ctx *schemas.BifrostContext) *schemas.BifrostSpeechRequest { - provider, model := schemas.ParseModelString(request.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.OpenAI)) + provider, model := schemas.ParseModelString(request.Model, "") return &schemas.BifrostSpeechRequest{ Provider: provider, diff --git a/core/providers/openai/text.go b/core/providers/openai/text.go index 07354a0263..e171050088 100644 --- a/core/providers/openai/text.go +++ b/core/providers/openai/text.go @@ -3,7 +3,6 @@ package openai import ( "maps" - "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" ) @@ -63,7 +62,7 @@ func (req *OpenAITextCompletionRequest) ToBifrostTextCompletionRequest(ctx *sche return nil } - provider, model := schemas.ParseModelString(req.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.OpenAI)) + provider, model := schemas.ParseModelString(req.Model, "") return &schemas.BifrostTextCompletionRequest{ Provider: provider, diff --git a/core/providers/openai/transcription.go b/core/providers/openai/transcription.go index 1bf419759a..cbfb130714 100644 --- a/core/providers/openai/transcription.go +++ b/core/providers/openai/transcription.go @@ -10,7 +10,7 @@ import ( // ToBifrostTranscriptionRequest converts an OpenAI transcription request to Bifrost format func (request *OpenAITranscriptionRequest) ToBifrostTranscriptionRequest(ctx *schemas.BifrostContext) *schemas.BifrostTranscriptionRequest { - provider, model := schemas.ParseModelString(request.Model, utils.CheckAndSetDefaultProvider(ctx, schemas.OpenAI)) + provider, model := schemas.ParseModelString(request.Model, "") return &schemas.BifrostTranscriptionRequest{ Provider: provider, diff --git a/core/providers/openai/videos.go b/core/providers/openai/videos.go index 512306b7c7..1b794b8ad3 100644 --- a/core/providers/openai/videos.go +++ b/core/providers/openai/videos.go @@ -6,7 +6,6 @@ import ( "mime/multipart" "net/http" - "github.com/maximhq/bifrost/core/providers/utils" providerUtils "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" ) @@ -101,16 +100,7 @@ func (req *OpenAIVideoGenerationRequest) ToBifrostVideoGenerationRequest(ctx *sc return nil } - defaultProvider := schemas.OpenAI - - // for requests coming from azure sdk without provider prefix, we need to set the default provider to azure - if ctx != nil { - if isAzureUser, ok := ctx.Value(schemas.BifrostContextKeyIsAzureUserAgent).(bool); ok && isAzureUser { - defaultProvider = schemas.Azure - } - } - - provider, model := schemas.ParseModelString(req.Model, utils.CheckAndSetDefaultProvider(ctx, defaultProvider)) + provider, model := schemas.ParseModelString(req.Model, "") input := &schemas.VideoGenerationInput{ Prompt: req.Prompt, @@ -132,30 +122,30 @@ func (req *OpenAIVideoGenerationRequest) ToBifrostVideoGenerationRequest(ctx *sc func parseVideoGenerationFormDataBodyFromRequest(writer *multipart.Writer, openaiReq *OpenAIVideoGenerationRequest, providerName schemas.ModelProvider) *schemas.BifrostError { // Add prompt field (required) if openaiReq.Prompt == "" { - return providerUtils.NewBifrostOperationError("prompt is required", nil) + return providerUtils.NewBifrostOperationError("prompt is required", nil) } if err := writer.WriteField("prompt", openaiReq.Prompt); err != nil { - return providerUtils.NewBifrostOperationError("failed to write prompt field", err) + return providerUtils.NewBifrostOperationError("failed to write prompt field", err) } // Add optional model field if openaiReq.Model != "" { if err := writer.WriteField("model", openaiReq.Model); err != nil { - return providerUtils.NewBifrostOperationError("failed to write model field", err) + return providerUtils.NewBifrostOperationError("failed to write model field", err) } } // Add optional seconds field if openaiReq.Seconds != nil { if err := writer.WriteField("seconds", *openaiReq.Seconds); err != nil { - return providerUtils.NewBifrostOperationError("failed to write seconds field", err) + return providerUtils.NewBifrostOperationError("failed to write seconds field", err) } } // Add optional size field if openaiReq.Size != "" { if err := writer.WriteField("size", openaiReq.Size); err != nil { - return providerUtils.NewBifrostOperationError("failed to write size field", err) + return providerUtils.NewBifrostOperationError("failed to write size field", err) } } @@ -196,16 +186,16 @@ func parseVideoGenerationFormDataBodyFromRequest(writer *multipart.Writer, opena "Content-Type": {mimeType}, }) if err != nil { - return providerUtils.NewBifrostOperationError("failed to create form part for input_reference", err) + return providerUtils.NewBifrostOperationError("failed to create form part for input_reference", err) } if _, err := part.Write(openaiReq.InputReference); err != nil { - return providerUtils.NewBifrostOperationError("failed to write input_reference file data", err) + return providerUtils.NewBifrostOperationError("failed to write input_reference file data", err) } } // Close the multipart writer if err := writer.Close(); err != nil { - return providerUtils.NewBifrostOperationError("failed to close multipart writer", err) + return providerUtils.NewBifrostOperationError("failed to close multipart writer", err) } return nil diff --git a/core/providers/utils/utils.go b/core/providers/utils/utils.go index f327412632..27f2ecb0d0 100644 --- a/core/providers/utils/utils.go +++ b/core/providers/utils/utils.go @@ -3056,36 +3056,6 @@ func completeDeferredSpan(ctx *schemas.BifrostContext, result *schemas.BifrostRe tracer.ClearDeferredSpan(traceID) } -// CheckAndSetDefaultProvider checks if the default provider should be used based on the context. -// It returns the default provider if it should be used, otherwise it returns an empty string. -// Checks if key selection is skipped, if a resolved provider was selected by routing, -// or if the available providers are set in the context and the default provider is in the list. -func CheckAndSetDefaultProvider(ctx *schemas.BifrostContext, defaultProvider schemas.ModelProvider) schemas.ModelProvider { - if ctx != nil { - if skip, ok := ctx.Value(schemas.BifrostContextKeySkipKeySelection).(bool); ok && skip { - return defaultProvider - } - if ctx.Value(schemas.BifrostContextKeyAvailableProviders) != nil { - availableProviders, ok := ctx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - if !ok || len(availableProviders) == 0 { - return "" - } - if resolvedProvider, ok := ctx.Value(schemas.BifrostContextKeyResolvedProvider).(schemas.ModelProvider); ok && slices.Contains(availableProviders, resolvedProvider) { - getLogger().Debug("[Provider] Using routing-resolved provider: %s (available: %v)", resolvedProvider, availableProviders) - return resolvedProvider - } - getLogger().Debug("[Provider] Available providers: %v, checking %s", availableProviders, defaultProvider) - if slices.Contains(availableProviders, defaultProvider) { - return defaultProvider - } - // Return the first available provider - return availableProviders[0] - } - return defaultProvider - } - return defaultProvider -} - // ModelMatchesDenylist reports whether any of the candidate model IDs matches // an entry in denylist, using both exact and base-model (SameBaseModel) matching. // Empty candidates are skipped. Returns false immediately if denylist is empty. diff --git a/core/providers/utils/utils_test.go b/core/providers/utils/utils_test.go index 66db016271..517da6bbaf 100644 --- a/core/providers/utils/utils_test.go +++ b/core/providers/utils/utils_test.go @@ -1878,27 +1878,3 @@ func TestExtractPassthroughProviderResponseHeaders(t *testing.T) { t.Fatalf("benign header x-request-id was dropped: %v", headers) } } - -// TestCheckAndSetDefaultProviderUsesResolvedProvider verifies routing-selected -// providers take precedence over the route default when still allowed. -func TestCheckAndSetDefaultProviderUsesResolvedProvider(t *testing.T) { - ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - ctx.SetValue(schemas.BifrostContextKeyAvailableProviders, []schemas.ModelProvider{schemas.Anthropic, schemas.Azure}) - ctx.SetValue(schemas.BifrostContextKeyResolvedProvider, schemas.Azure) - - if got := CheckAndSetDefaultProvider(ctx, schemas.Anthropic); got != schemas.Azure { - t.Fatalf("CheckAndSetDefaultProvider() = %s, want %s", got, schemas.Azure) - } -} - -// TestCheckAndSetDefaultProviderIgnoresDisallowedResolvedProvider verifies -// selected-provider context cannot bypass available-provider constraints. -func TestCheckAndSetDefaultProviderIgnoresDisallowedResolvedProvider(t *testing.T) { - ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - ctx.SetValue(schemas.BifrostContextKeyAvailableProviders, []schemas.ModelProvider{schemas.Anthropic}) - ctx.SetValue(schemas.BifrostContextKeyResolvedProvider, schemas.Azure) - - if got := CheckAndSetDefaultProvider(ctx, schemas.Anthropic); got != schemas.Anthropic { - t.Fatalf("CheckAndSetDefaultProvider() = %s, want %s", got, schemas.Anthropic) - } -} diff --git a/core/providers/vertex/rerank.go b/core/providers/vertex/rerank.go index 257a1f8def..af2e68d25c 100644 --- a/core/providers/vertex/rerank.go +++ b/core/providers/vertex/rerank.go @@ -7,7 +7,6 @@ import ( "strings" "github.com/bytedance/sonic" - providerUtils "github.com/maximhq/bifrost/core/providers/utils" "github.com/maximhq/bifrost/core/schemas" ) @@ -157,9 +156,9 @@ func (req *VertexRankRequest) ToBifrostRerankRequest(ctx *schemas.BifrostContext var provider schemas.ModelProvider var model string if req.Model != nil { - provider, model = schemas.ParseModelString(*req.Model, providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Vertex)) + provider, model = schemas.ParseModelString(*req.Model, schemas.Vertex) } else { - provider = providerUtils.CheckAndSetDefaultProvider(ctx, schemas.Vertex) + provider = schemas.Vertex } bifrostReq := &schemas.BifrostRerankRequest{ diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index e9f8c1c7e5..ce42cb2990 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -264,8 +264,6 @@ const ( BifrostContextKeyGovernanceRateLimitIDs BifrostContextKey = "bifrost-governance-rate-limit-ids" // []string (rate limit IDs applicable to this request - set by governance plugin) BifrostContextKeyPromptsPluginName BifrostContextKey = "prompts-plugin-name" // string (name of the prompts plugin to use - set by bifrost - DO NOT SET THIS MANUALLY)) BifrostContextKeyIsEnterprise BifrostContextKey = "is-enterprise" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) - BifrostContextKeyAvailableProviders BifrostContextKey = "available-providers" // []ModelProvider (set by bifrost - DO NOT SET THIS MANUALLY)) - BifrostContextKeyResolvedProvider BifrostContextKey = "bifrost-resolved-provider" // ModelProvider (set by routing - DO NOT SET THIS MANUALLY)) BifrostContextKeyStoreRawRequestResponse BifrostContextKey = "bifrost-store-raw-request-response" // bool (per-request override — read by bifrost.go, never overwritten) BifrostContextKeyCaptureRawRequest BifrostContextKey = "bifrost-capture-raw-request" // bool (set by bifrost - DO NOT SET THIS MANUALLY) — true when providers should capture raw request bytes BifrostContextKeyCaptureRawResponse BifrostContextKey = "bifrost-capture-raw-response" // bool (set by bifrost - DO NOT SET THIS MANUALLY) — true when providers should capture raw response bytes @@ -304,39 +302,38 @@ const ( BifrostContextKeyIsAzureUserAgent BifrostContextKey = "bifrost-is-azure-user-agent" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) - whether the request is an Azure user agent (only used in gateway) BifrostContextKeyUserRoleID BifrostContextKey = "bifrost-user-role-id" BifrostContextKeyVideoOutputRequested BifrostContextKey = "bifrost-video-output-requested" - BifrostContextKeyValidateKeys BifrostContextKey = "bifrost-validate-keys" // bool (triggers additional key validation during provider add/update) - BifrostContextKeyProviderResponseHeaders BifrostContextKey = "bifrost-provider-response-headers" // map[string]string (set by provider handlers for response header forwarding) - BifrostContextKeyMCPAddedTools BifrostContextKey = "bifrost-mcp-added-tools" // []string (set by bifrost - DO NOT SET THIS MANUALLY)) - list of tools added to the request by MCP, all the tool are in the format "clientName-toolName" - BifrostContextKeyLargePayloadMode BifrostContextKey = "bifrost-large-payload-mode" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) indicates large payload streaming mode is active - BifrostContextKeyLargePayloadReader BifrostContextKey = "bifrost-large-payload-reader" // io.Reader (set by bifrost - DO NOT SET THIS MANUALLY)) upstream reader for large payloads - BifrostContextKeyLargePayloadContentLength BifrostContextKey = "bifrost-large-payload-content-length" // int (set by bifrost - DO NOT SET THIS MANUALLY)) content length for large payloads - BifrostContextKeyLargePayloadContentType BifrostContextKey = "bifrost-large-payload-content-type" // string (set by enterprise - DO NOT SET THIS MANUALLY)) original content type for large payload passthrough - BifrostContextKeyLargePayloadMetadata BifrostContextKey = "bifrost-large-payload-metadata" // *LargePayloadMetadata (set by bifrost - DO NOT SET THIS MANUALLY)) routing metadata for large payloads - BifrostContextKeyLargePayloadRequestThreshold BifrostContextKey = "bifrost-large-payload-request-threshold" // int64 (set by enterprise - DO NOT SET THIS MANUALLY)) request threshold used by transport heuristics - BifrostContextKeyLargeResponseMode BifrostContextKey = "bifrost-large-response-mode" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) indicates large response streaming mode is active - BifrostContextKeyLargePayloadRequestPreview BifrostContextKey = "bifrost-large-payload-request-preview" // string (set by bifrost - DO NOT SET THIS MANUALLY)) truncated request body preview for logging - BifrostContextKeyLargePayloadResponsePreview BifrostContextKey = "bifrost-large-payload-response-preview" // string (set by bifrost - DO NOT SET THIS MANUALLY)) truncated response body preview for logging - BifrostContextKeyLargeResponseReader BifrostContextKey = "bifrost-large-response-reader" // io.ReadCloser (set by bifrost - DO NOT SET THIS MANUALLY)) upstream reader for large responses - BifrostContextKeyLargeResponseContentLength BifrostContextKey = "bifrost-large-response-content-length" // int (set by bifrost - DO NOT SET THIS MANUALLY)) content length for large responses - BifrostContextKeyLargeResponseContentType BifrostContextKey = "bifrost-large-response-content-type" // string (set by bifrost - DO NOT SET THIS MANUALLY)) upstream content type for large responses - BifrostContextKeyLargeResponseContentDisposition BifrostContextKey = "bifrost-large-response-content-disposition" // string (set by bifrost - DO NOT SET THIS MANUALLY)) downstream content disposition for large responses - BifrostContextKeyLargeResponseThreshold BifrostContextKey = "bifrost-large-response-threshold" // int64 (set by enterprise - DO NOT SET THIS MANUALLY)) threshold for response streaming - BifrostContextKeyLargePayloadPrefetchSize BifrostContextKey = "bifrost-large-payload-prefetch-size" // int (set by enterprise - DO NOT SET THIS MANUALLY)) prefetch buffer size for metadata extraction from large responses - BifrostContextKeyDeferredUsage BifrostContextKey = "bifrost-deferred-usage" // chan *BifrostLLMUsage (set by provider Phase B — delivers usage after response streaming completes) - BifrostContextKeyDeferredLargePayloadMetadata BifrostContextKey = "bifrost-deferred-large-payload-metadata" // <-chan *LargePayloadMetadata (set by enterprise Phase B request — delivers metadata after body streaming) - BifrostContextKeySSEReaderFactory BifrostContextKey = "bifrost-sse-reader-factory" // *providerUtils.SSEReaderFactory (set by enterprise — replaces default bufio.Scanner SSE readers with streaming readers) - BifrostContextKeySessionID BifrostContextKey = "bifrost-session-id" // string session ID for the request (session stickiness) - BifrostContextKeySessionTTL BifrostContextKey = "bifrost-session-ttl" // time.Duration session TTL for the request (session stickiness) - BifrostContextKeyMCPExtraHeaders BifrostContextKey = "bifrost-mcp-extra-headers" // map[string][]string (these headers are forwarded only to the MCP while tool execution if they are in the allowlist of the MCP client) - BifrostContextKeyMCPLogID BifrostContextKey = "bifrost-mcp-log-id" // string (unique UUID for each MCP tool log entry - set per goroutine by agent executor - DO NOT SET THIS MANUALLY) - BifrostContextKeyCompatConvertTextToChat BifrostContextKey = "bifrost-compat-convert-text-to-chat" // bool (per-request override from x-bf-compat header) - BifrostContextKeyCompatConvertChatToResponses BifrostContextKey = "bifrost-compat-convert-chat-to-responses" // bool (per-request override from x-bf-compat header) - BifrostContextKeyCompatShouldDropParams BifrostContextKey = "bifrost-compat-should-drop-params" // bool (per-request override from x-bf-compat header) - BifrostContextKeyCompatShouldConvertParams BifrostContextKey = "bifrost-compat-should-convert-params" // bool (per-request override from x-bf-compat header) - BifrostContextKeySupportsAssistantPrefill BifrostContextKey = "bifrost-supports-assistant-prefill" // bool (set by compat plugin) - if model supports assistant prefill - BifrostContextKeyAttemptTrail BifrostContextKey = "bifrost-attempt-trail" // []KeyAttemptRecord (set by bifrost - DO NOT SET THIS MANUALLY) - per-attempt key selection history - BifrostContextKeyDimensions BifrostContextKey = "bifrost-dimensions" // map[string]string (set by HTTP transport from x-bf-dim-* headers) BifrostContextKeyDimensions holds per-request key/value dimensions supplied via x-bf-dim- request headers. These dimensions are forwarded to internal logs (as metadata) - BifrostContextKeySkipModelCatalogProviderSelection BifrostContextKey = "bifrost-skip-model-catalog-provider-selection" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) - skip model catalog provider selection + BifrostContextKeyValidateKeys BifrostContextKey = "bifrost-validate-keys" // bool (triggers additional key validation during provider add/update) + BifrostContextKeyProviderResponseHeaders BifrostContextKey = "bifrost-provider-response-headers" // map[string]string (set by provider handlers for response header forwarding) + BifrostContextKeyMCPAddedTools BifrostContextKey = "bifrost-mcp-added-tools" // []string (set by bifrost - DO NOT SET THIS MANUALLY)) - list of tools added to the request by MCP, all the tool are in the format "clientName-toolName" + BifrostContextKeyLargePayloadMode BifrostContextKey = "bifrost-large-payload-mode" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) indicates large payload streaming mode is active + BifrostContextKeyLargePayloadReader BifrostContextKey = "bifrost-large-payload-reader" // io.Reader (set by bifrost - DO NOT SET THIS MANUALLY)) upstream reader for large payloads + BifrostContextKeyLargePayloadContentLength BifrostContextKey = "bifrost-large-payload-content-length" // int (set by bifrost - DO NOT SET THIS MANUALLY)) content length for large payloads + BifrostContextKeyLargePayloadContentType BifrostContextKey = "bifrost-large-payload-content-type" // string (set by enterprise - DO NOT SET THIS MANUALLY)) original content type for large payload passthrough + BifrostContextKeyLargePayloadMetadata BifrostContextKey = "bifrost-large-payload-metadata" // *LargePayloadMetadata (set by bifrost - DO NOT SET THIS MANUALLY)) routing metadata for large payloads + BifrostContextKeyLargePayloadRequestThreshold BifrostContextKey = "bifrost-large-payload-request-threshold" // int64 (set by enterprise - DO NOT SET THIS MANUALLY)) request threshold used by transport heuristics + BifrostContextKeyLargeResponseMode BifrostContextKey = "bifrost-large-response-mode" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) indicates large response streaming mode is active + BifrostContextKeyLargePayloadRequestPreview BifrostContextKey = "bifrost-large-payload-request-preview" // string (set by bifrost - DO NOT SET THIS MANUALLY)) truncated request body preview for logging + BifrostContextKeyLargePayloadResponsePreview BifrostContextKey = "bifrost-large-payload-response-preview" // string (set by bifrost - DO NOT SET THIS MANUALLY)) truncated response body preview for logging + BifrostContextKeyLargeResponseReader BifrostContextKey = "bifrost-large-response-reader" // io.ReadCloser (set by bifrost - DO NOT SET THIS MANUALLY)) upstream reader for large responses + BifrostContextKeyLargeResponseContentLength BifrostContextKey = "bifrost-large-response-content-length" // int (set by bifrost - DO NOT SET THIS MANUALLY)) content length for large responses + BifrostContextKeyLargeResponseContentType BifrostContextKey = "bifrost-large-response-content-type" // string (set by bifrost - DO NOT SET THIS MANUALLY)) upstream content type for large responses + BifrostContextKeyLargeResponseContentDisposition BifrostContextKey = "bifrost-large-response-content-disposition" // string (set by bifrost - DO NOT SET THIS MANUALLY)) downstream content disposition for large responses + BifrostContextKeyLargeResponseThreshold BifrostContextKey = "bifrost-large-response-threshold" // int64 (set by enterprise - DO NOT SET THIS MANUALLY)) threshold for response streaming + BifrostContextKeyLargePayloadPrefetchSize BifrostContextKey = "bifrost-large-payload-prefetch-size" // int (set by enterprise - DO NOT SET THIS MANUALLY)) prefetch buffer size for metadata extraction from large responses + BifrostContextKeyDeferredUsage BifrostContextKey = "bifrost-deferred-usage" // chan *BifrostLLMUsage (set by provider Phase B — delivers usage after response streaming completes) + BifrostContextKeyDeferredLargePayloadMetadata BifrostContextKey = "bifrost-deferred-large-payload-metadata" // <-chan *LargePayloadMetadata (set by enterprise Phase B request — delivers metadata after body streaming) + BifrostContextKeySSEReaderFactory BifrostContextKey = "bifrost-sse-reader-factory" // *providerUtils.SSEReaderFactory (set by enterprise — replaces default bufio.Scanner SSE readers with streaming readers) + BifrostContextKeySessionID BifrostContextKey = "bifrost-session-id" // string session ID for the request (session stickiness) + BifrostContextKeySessionTTL BifrostContextKey = "bifrost-session-ttl" // time.Duration session TTL for the request (session stickiness) + BifrostContextKeyMCPExtraHeaders BifrostContextKey = "bifrost-mcp-extra-headers" // map[string][]string (these headers are forwarded only to the MCP while tool execution if they are in the allowlist of the MCP client) + BifrostContextKeyMCPLogID BifrostContextKey = "bifrost-mcp-log-id" // string (unique UUID for each MCP tool log entry - set per goroutine by agent executor - DO NOT SET THIS MANUALLY) + BifrostContextKeyCompatConvertTextToChat BifrostContextKey = "bifrost-compat-convert-text-to-chat" // bool (per-request override from x-bf-compat header) + BifrostContextKeyCompatConvertChatToResponses BifrostContextKey = "bifrost-compat-convert-chat-to-responses" // bool (per-request override from x-bf-compat header) + BifrostContextKeyCompatShouldDropParams BifrostContextKey = "bifrost-compat-should-drop-params" // bool (per-request override from x-bf-compat header) + BifrostContextKeyCompatShouldConvertParams BifrostContextKey = "bifrost-compat-should-convert-params" // bool (per-request override from x-bf-compat header) + BifrostContextKeySupportsAssistantPrefill BifrostContextKey = "bifrost-supports-assistant-prefill" // bool (set by compat plugin) - if model supports assistant prefill + BifrostContextKeyAttemptTrail BifrostContextKey = "bifrost-attempt-trail" // []KeyAttemptRecord (set by bifrost - DO NOT SET THIS MANUALLY) - per-attempt key selection history + BifrostContextKeyDimensions BifrostContextKey = "bifrost-dimensions" // map[string]string (set by HTTP transport from x-bf-dim-* headers) BifrostContextKeyDimensions holds per-request key/value dimensions supplied via x-bf-dim- request headers. These dimensions are forwarded to internal logs (as metadata) IsAPIKeyAuthContextKey BifrostContextKey = "is_api_key_auth" IsLocalAdminContextKey BifrostContextKey = "is_local_admin" // bool (set by auth middleware when password-based auth succeeds - local admin user bypasses RBAC) BifrostContextKeyPassthroughOverridesPresent BifrostContextKey = "passthrough_overrides_present" // bool (set by HTTP transport) - passthrough raw request requested diff --git a/plugins/governance/go.mod b/plugins/governance/go.mod index 44b0a3effc..45772b1f6b 100644 --- a/plugins/governance/go.mod +++ b/plugins/governance/go.mod @@ -5,7 +5,6 @@ go 1.26.4 require gorm.io/gorm v1.31.1 require ( - github.com/bytedance/sonic v1.15.1 github.com/google/cel-go v0.28.1 github.com/google/uuid v1.6.0 github.com/maximhq/bifrost/core v1.5.18 @@ -55,7 +54,7 @@ require ( github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/buger/jsonparser v1.1.2 // indirect github.com/bytedance/gopkg v0.1.3 // indirect - github.com/bytedance/sonic v1.15.0 // indirect + github.com/bytedance/sonic v1.15.1 // indirect github.com/bytedance/sonic/loader v0.5.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/cloudwego/base64x v0.1.6 // indirect diff --git a/plugins/governance/httptransportprehook_test.go b/plugins/governance/httptransportprehook_test.go deleted file mode 100644 index 7e3f1db452..0000000000 --- a/plugins/governance/httptransportprehook_test.go +++ /dev/null @@ -1,346 +0,0 @@ -package governance - -import ( - "context" - "testing" - - "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/stretchr/testify/require" -) - -func TestHTTPTransportPreHook_ModelOnlyVirtualKeySetsAvailableProviders(t *testing.T) { - logger := NewMockLogger() - - openAIConfig := buildProviderConfig("openai", []string{"gpt-4o"}) - openAIConfig.Weight = nil - anthropicConfig := buildProviderConfig("anthropic", []string{"claude-3-5-sonnet"}) - anthropicConfig.Weight = nil - - virtualKey := buildVirtualKeyWithProviders( - "vk-constraint", - "sk-bf-constraint-test", - "provider-constraint-vk", - []configstoreTables.TableVirtualKeyProviderConfig{ - openAIConfig, - anthropicConfig, - }, - ) - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - }, nil) - require.NoError(t, err) - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/v1/chat/completions" - req.Headers["Authorization"] = "Bearer sk-bf-constraint-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"Hello!"}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - allowedProviders, ok := bfCtx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - require.True(t, ok, "provider constraint should be set") - require.Equal(t, []schemas.ModelProvider{schemas.OpenAI}, allowedProviders) -} - -func TestHTTPTransportPreHook_ModelOnlyVirtualKeySetsEmptyAvailableProvidersWhenNoProviderAllowsModel(t *testing.T) { - logger := NewMockLogger() - - virtualKey := buildVirtualKeyWithProviders( - "vk-empty-constraint", - "sk-bf-empty-constraint-test", - "empty-provider-constraint-vk", - []configstoreTables.TableVirtualKeyProviderConfig{ - buildProviderConfig("openai", []string{"gpt-4o"}), - }, - ) - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - }, nil) - require.NoError(t, err) - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, nil, nil, nil) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/v1/chat/completions" - req.Headers["Authorization"] = "Bearer sk-bf-empty-constraint-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"model":"claude-3-5-sonnet","messages":[{"role":"user","content":"Hello!"}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - allowedProviders, ok := bfCtx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - require.True(t, ok, "provider constraint should be set") - require.Empty(t, allowedProviders) -} - -// TestHTTPTransportPreHook_WildcardKeepsCatalogOpaqueProvider_VLLM verifies that a VK with a -// wildcard ("*") allow-list on a catalog-opaque provider (here vLLM, whose self-hosted models -// are never in the bundled catalog) keeps that provider in BifrostContextKeyAvailableProviders -// for a bare, uncatalogued model. Before the fix, loadBalanceProvider gates the provider on the -// catalog (GetProvidersForModel is empty), drops it, and publishes an empty provider set — -// dead-ending the request (issue #4122 / #3282). -func TestHTTPTransportPreHook_WildcardKeepsCatalogOpaqueProvider_VLLM(t *testing.T) { - logger := NewMockLogger() - - // Catalog knows a first-party model but has NO model list for vLLM (self-hosted). - mc := modelcatalog.NewTestCatalog(map[string]string{"openai/gpt-4o": "gpt-4o"}) - mc.UpsertModelDataForProvider(schemas.OpenAI, - &schemas.BifrostListModelsResponse{Data: []schemas.Model{{ID: "openai/gpt-4o"}}}, nil) - - // inMemoryStore must be non-nil so loadBalanceProvider takes the catalog branch. - inMem := &mockInMemoryStore{ - configuredProviders: map[schemas.ModelProvider]configstore.ProviderConfig{ - schemas.VLLM: {}, // native keyless: no CustomProviderConfig - }, - } - - virtualKey := buildVirtualKeyWithProviders( - "vk-vllm", - "sk-bf-vllm-test", - "vllm-vk", - []configstoreTables.TableVirtualKeyProviderConfig{ - buildProviderConfig("vllm", []string{"*"}), - }, - ) - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - }, mc) - require.NoError(t, err) - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, mc, nil, inMem) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/v1/chat/completions" - req.Headers["Authorization"] = "Bearer sk-bf-vllm-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"model":"my-self-hosted-llama","messages":[{"role":"user","content":"Hi"}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - allowedProviders, ok := bfCtx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - require.True(t, ok, "available providers should be set") - // PRE-PATCH: catalog has no vLLM models -> provider excluded -> [] -> FAILS. - // POST-PATCH: wildcard + catalog-opaque -> kept -> [vllm]. - require.Equal(t, []schemas.ModelProvider{schemas.VLLM}, allowedProviders) -} - -// TestHTTPTransportPreHook_MixedOpaqueAndCatalogProvider_GPT4o shows what lands in -// BifrostContextKeyAvailableProviders when a VK has catalog-known providers (openai, anthropic, -// vertex) AND a catalog-opaque vLLM (no list-models) — all under wildcard allow-lists — and the -// model is gpt-4o. Only openai (which serves gpt-4o per the catalog) and vLLM (a wildcard -// catch-all) should be available; anthropic and vertex are catalog-known but do not serve gpt-4o. -func TestHTTPTransportPreHook_MixedOpaqueAndCatalogProvider_GPT4o(t *testing.T) { - logger := NewMockLogger() - - // Catalog knows openai/gpt-4o, anthropic/claude-3-5-sonnet, vertex/gemini-1.5-pro. - // It has NO model list for vLLM (opaque). - mc := modelcatalog.NewTestCatalog(map[string]string{"openai/gpt-4o": "gpt-4o"}) - mc.UpsertModelDataForProvider(schemas.OpenAI, - &schemas.BifrostListModelsResponse{Data: []schemas.Model{{ID: "openai/gpt-4o"}}}, nil) - mc.UpsertModelDataForProvider(schemas.Anthropic, - &schemas.BifrostListModelsResponse{Data: []schemas.Model{{ID: "anthropic/claude-3-5-sonnet"}}}, nil) - mc.UpsertModelDataForProvider(schemas.Vertex, - &schemas.BifrostListModelsResponse{Data: []schemas.Model{{ID: "vertex/gemini-1.5-pro"}}}, nil) - - inMem := &mockInMemoryStore{ - configuredProviders: map[schemas.ModelProvider]configstore.ProviderConfig{ - schemas.OpenAI: {}, - schemas.Anthropic: {}, - schemas.Vertex: {}, - schemas.VLLM: {}, // opaque: no CustomProviderConfig, no catalog models - }, - } - - virtualKey := buildVirtualKeyWithProviders( - "vk-mixed", - "sk-bf-mixed-test", - "mixed-providers-vk", - []configstoreTables.TableVirtualKeyProviderConfig{ - buildProviderConfig("openai", []string{"*"}), - buildProviderConfig("anthropic", []string{"*"}), - buildProviderConfig("vertex", []string{"*"}), - buildProviderConfig("vllm", []string{"*"}), - }, - ) - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - }, mc) - require.NoError(t, err) - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, mc, nil, inMem) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/v1/chat/completions" - req.Headers["Authorization"] = "Bearer sk-bf-mixed-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"Hi"}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - allowedProviders, ok := bfCtx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - require.True(t, ok, "available providers should be set") - t.Logf("AvailableProviders for gpt-4o (VK = openai + vllm-opaque, both wildcard): %v", allowedProviders) - - // Both compete: openai matches the catalog for gpt-4o; vLLM is a wildcard catch-all. - require.ElementsMatch(t, []schemas.ModelProvider{schemas.OpenAI, schemas.VLLM}, allowedProviders) -} - -// TestHTTPTransportPreHook_VKExcludesUnlistedProviderEvenIfItServesModel shows that VK scoping -// wins: even when the catalog says BOTH openai and vertex serve gpt-4o, a VK granting access to -// only openai + vLLM yields exactly [openai, vllm] — vertex is never a candidate because it is -// not in the VK's provider configs. -func TestHTTPTransportPreHook_VKExcludesUnlistedProviderEvenIfItServesModel(t *testing.T) { - logger := NewMockLogger() - - // Catalog: BOTH openai and vertex serve gpt-4o. vLLM has no catalog models (opaque). - mc := modelcatalog.NewTestCatalog(map[string]string{"openai/gpt-4o": "gpt-4o"}) - mc.UpsertModelDataForProvider(schemas.OpenAI, - &schemas.BifrostListModelsResponse{Data: []schemas.Model{{ID: "openai/gpt-4o"}}}, nil) - mc.UpsertModelDataForProvider(schemas.Vertex, - &schemas.BifrostListModelsResponse{Data: []schemas.Model{{ID: "vertex/gpt-4o"}}}, nil) - - inMem := &mockInMemoryStore{ - configuredProviders: map[schemas.ModelProvider]configstore.ProviderConfig{ - schemas.OpenAI: {}, - schemas.Vertex: {}, - schemas.VLLM: {}, // opaque - }, - } - - // VK grants access to ONLY openai and vLLM — NOT vertex, even though vertex serves gpt-4o. - virtualKey := buildVirtualKeyWithProviders( - "vk-scoped", - "sk-bf-scoped-test", - "openai-vllm-only-vk", - []configstoreTables.TableVirtualKeyProviderConfig{ - buildProviderConfig("openai", []string{"*"}), - buildProviderConfig("vllm", []string{"*"}), - }, - ) - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - }, mc) - require.NoError(t, err) - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, mc, nil, inMem) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/v1/chat/completions" - req.Headers["Authorization"] = "Bearer sk-bf-scoped-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"Hi"}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - allowedProviders, ok := bfCtx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - require.True(t, ok, "available providers should be set") - t.Logf("AvailableProviders for gpt-4o (catalog: openai+vertex serve it; VK = openai + vllm only): %v", allowedProviders) - - // Vertex serves gpt-4o per the catalog but is NOT in the VK, so it must be absent. - require.ElementsMatch(t, []schemas.ModelProvider{schemas.OpenAI, schemas.VLLM}, allowedProviders) -} - -// TestHTTPTransportPreHook_WildcardOpaqueProviderRespectsBlacklist guards the ordering in -// loadBalanceProvider: the blacklist pre-pass must exclude a provider before the wildcard + -// catalog-opaque shortcut applies, so a blacklisted model on an opaque provider is dropped -// from BifrostContextKeyAvailableProviders even under a ["*"] allow-list. -func TestHTTPTransportPreHook_WildcardOpaqueProviderRespectsBlacklist(t *testing.T) { - logger := NewMockLogger() - - mc := modelcatalog.NewTestCatalog(nil) // no vLLM models -> opaque - - inMem := &mockInMemoryStore{ - configuredProviders: map[schemas.ModelProvider]configstore.ProviderConfig{ - schemas.VLLM: {}, - }, - } - - vllmConfig := buildProviderConfig("vllm", []string{"*"}) - vllmConfig.BlacklistedModels = schemas.BlackList{"my-self-hosted-llama"} - - virtualKey := buildVirtualKeyWithProviders( - "vk-vllm-bl", - "sk-bf-vllm-bl-test", - "vllm-bl-vk", - []configstoreTables.TableVirtualKeyProviderConfig{vllmConfig}, - ) - store, err := NewLocalGovernanceStore(context.Background(), logger, nil, &configstore.GovernanceConfig{ - VirtualKeys: []configstoreTables.TableVirtualKey{*virtualKey}, - }, mc) - require.NoError(t, err) - - plugin, err := InitFromStore(context.Background(), &Config{IsVkMandatory: boolPtr(false)}, logger, store, nil, mc, nil, inMem) - require.NoError(t, err) - defer func() { - require.NoError(t, plugin.Cleanup()) - }() - - req := schemas.AcquireHTTPRequest() - defer schemas.ReleaseHTTPRequest(req) - req.Method = "POST" - req.Path = "/v1/chat/completions" - req.Headers["Authorization"] = "Bearer sk-bf-vllm-bl-test" - req.Headers["Content-Type"] = "application/json" - req.Body = []byte(`{"model":"my-self-hosted-llama","messages":[{"role":"user","content":"Hi"}]}`) - - bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) - resp, err := plugin.HTTPTransportPreHook(bfCtx, req) - require.NoError(t, err) - require.Nil(t, resp) - - allowedProviders, ok := bfCtx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - require.True(t, ok, "available providers should be set") - // Blacklisted model is excluded even though the provider is catalog-opaque under ["*"]. - require.Empty(t, allowedProviders) -} diff --git a/plugins/governance/main.go b/plugins/governance/main.go index fb47d59d0b..53953ae32b 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -430,7 +430,6 @@ func (p *GovernancePlugin) loadBalanceProvider(ctx *schemas.BifrostContext, req // Get provider configs for this virtual key providerConfigs := virtualKey.ProviderConfigs if len(providerConfigs) == 0 { - ctx.SetValue(schemas.BifrostContextKeyAvailableProviders, []schemas.ModelProvider{}) ctx.AppendRoutingEngineLog(schemas.RoutingEngineGovernance, schemas.LogLevelWarn, fmt.Sprintf("No provider configs on virtual key %s for model %s, skipping load balancing", virtualKey.Name, modelStr)) // No provider configs, continue without modification return nil @@ -494,12 +493,9 @@ func (p *GovernancePlugin) loadBalanceProvider(ctx *schemas.BifrostContext, req } var allowedProviders []string - allowedModelProviders := make([]schemas.ModelProvider, 0, len(allowedProviderConfigs)) for _, pc := range allowedProviderConfigs { allowedProviders = append(allowedProviders, pc.Provider) - allowedModelProviders = append(allowedModelProviders, schemas.ModelProvider(pc.Provider)) } - ctx.SetValue(schemas.BifrostContextKeyAvailableProviders, allowedModelProviders) p.logger.Debug("[Governance] Allowed providers after filtering: %v", allowedProviders) ctx.AppendRoutingEngineLog(schemas.RoutingEngineGovernance, schemas.LogLevelInfo, fmt.Sprintf("Allowed providers after filtering: %v", allowedProviders)) diff --git a/plugins/modelcatalogresolver/go.mod b/plugins/modelcatalogresolver/go.mod new file mode 100644 index 0000000000..fe0b864412 --- /dev/null +++ b/plugins/modelcatalogresolver/go.mod @@ -0,0 +1,159 @@ +module github.com/maximhq/bifrost/plugins/modelcatalogresolver + +go 1.26.3 + +require ( + github.com/maximhq/bifrost/core v1.5.15 + github.com/maximhq/bifrost/framework v1.3.15 +) + +require ( + cel.dev/expr v0.25.1 // indirect + cloud.google.com/go v0.123.0 // indirect + cloud.google.com/go/auth v0.20.0 // indirect + cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect + cloud.google.com/go/compute/metadata v0.9.0 // indirect + cloud.google.com/go/iam v1.5.3 // indirect + cloud.google.com/go/monitoring v1.24.3 // indirect + cloud.google.com/go/storage v1.61.3 // indirect + github.com/Azure/azure-sdk-for-go/sdk/azcore v1.20.0 // indirect + github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.1 // indirect + github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.2 // indirect + github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0 // indirect + github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.31.0 // indirect + github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.55.0 // indirect + github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.55.0 // indirect + github.com/andybalholm/brotli v1.2.1 // indirect + github.com/apapsch/go-jsonmerge/v2 v2.0.0 // indirect + github.com/aws/aws-sdk-go-v2 v1.41.7 // indirect + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.10 // indirect + github.com/aws/aws-sdk-go-v2/config v1.32.11 // indirect + github.com/aws/aws-sdk-go-v2/credentials v1.19.14 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.23 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.23 // indirect + github.com/aws/aws-sdk-go-v2/internal/ini v1.8.5 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.22 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.13 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.21 // indirect + github.com/aws/aws-sdk-go-v2/service/s3 v1.97.3 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.0.9 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.30.15 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.19 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.41.10 // indirect + github.com/aws/smithy-go v1.25.1 // indirect + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/buger/jsonparser v1.1.2 // indirect + github.com/bytedance/gopkg v0.1.3 // indirect + github.com/bytedance/sonic v1.15.1 // indirect + github.com/bytedance/sonic/loader v0.5.1 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/cloudwego/base64x v0.1.6 // indirect + github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 // indirect + github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/envoyproxy/go-control-plane/envoy v1.37.0 // indirect + github.com/envoyproxy/protoc-gen-validate v1.3.3 // indirect + github.com/felixge/httpsnoop v1.0.4 // indirect + github.com/go-jose/go-jose/v4 v4.1.4 // indirect + github.com/go-logr/logr v1.4.3 // indirect + github.com/go-logr/stdr v1.2.2 // indirect + github.com/go-openapi/analysis v0.24.2 // indirect + github.com/go-openapi/errors v0.22.5 // indirect + github.com/go-openapi/jsonpointer v0.22.4 // indirect + github.com/go-openapi/jsonreference v0.21.4 // indirect + github.com/go-openapi/loads v0.23.2 // indirect + github.com/go-openapi/runtime v0.29.2 // indirect + github.com/go-openapi/spec v0.22.2 // indirect + github.com/go-openapi/strfmt v0.25.0 // indirect + github.com/go-openapi/swag v0.25.4 // indirect + github.com/go-openapi/swag/cmdutils v0.25.4 // indirect + github.com/go-openapi/swag/conv v0.25.4 // indirect + github.com/go-openapi/swag/fileutils v0.25.4 // indirect + github.com/go-openapi/swag/jsonname v0.25.4 // indirect + github.com/go-openapi/swag/jsonutils v0.25.4 // indirect + github.com/go-openapi/swag/loading v0.25.4 // indirect + github.com/go-openapi/swag/mangling v0.25.4 // indirect + github.com/go-openapi/swag/netutils v0.25.4 // indirect + github.com/go-openapi/swag/stringutils v0.25.4 // indirect + github.com/go-openapi/swag/typeutils v0.25.4 // indirect + github.com/go-openapi/swag/yamlutils v0.25.4 // indirect + github.com/go-openapi/validate v0.25.1 // indirect + github.com/go-viper/mapstructure/v2 v2.5.0 // indirect + github.com/golang-jwt/jwt/v5 v5.3.1 // indirect + github.com/google/s2a-go v0.1.9 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/googleapis/enterprise-certificate-proxy v0.3.16 // indirect + github.com/googleapis/gax-go/v2 v2.22.0 // indirect + github.com/invopop/jsonschema v0.13.0 // indirect + github.com/jackc/pgpassfile v1.0.0 // indirect + github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/jackc/pgx/v5 v5.9.2 // indirect + github.com/jackc/puddle/v2 v2.2.2 // indirect + github.com/jinzhu/inflection v1.0.0 // indirect + github.com/jinzhu/now v1.1.5 // indirect + github.com/klauspost/compress v1.18.6 // indirect + github.com/klauspost/cpuid/v2 v2.3.0 // indirect + github.com/kylelemons/godebug v1.1.0 // indirect + github.com/mailru/easyjson v0.9.1 // indirect + github.com/mark3labs/mcp-go v0.43.2 // indirect + github.com/mattn/go-colorable v0.1.14 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-sqlite3 v1.14.32 // indirect + github.com/oapi-codegen/runtime v1.1.1 // indirect + github.com/oklog/ulid v1.3.1 // indirect + github.com/pinecone-io/go-pinecone/v5 v5.3.0 // indirect + github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect + github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect + github.com/qdrant/go-client v1.16.2 // indirect + github.com/redis/go-redis/v9 v9.17.2 // indirect + github.com/rs/zerolog v1.34.0 // indirect + github.com/spf13/cast v1.10.0 // indirect + github.com/spiffe/go-spiffe/v2 v2.6.0 // indirect + github.com/stretchr/testify v1.11.1 // indirect + github.com/tidwall/gjson v1.18.0 // indirect + github.com/tidwall/match v1.1.1 // indirect + github.com/tidwall/pretty v1.2.0 // indirect + github.com/tidwall/sjson v1.2.5 // indirect + github.com/twitchyliquid64/golang-asm v0.15.1 // indirect + github.com/valyala/bytebufferpool v1.0.0 // indirect + github.com/valyala/fasthttp v1.71.0 // indirect + github.com/weaviate/weaviate v1.36.5 // indirect + github.com/weaviate/weaviate-go-client/v5 v5.7.1 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect + go.mongodb.org/mongo-driver v1.17.6 // indirect + go.opentelemetry.io/auto/sdk v1.2.1 // indirect + go.opentelemetry.io/contrib/detectors/gcp v1.42.0 // indirect + go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.67.0 // indirect + go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.67.0 // indirect + go.opentelemetry.io/otel v1.43.0 // indirect + go.opentelemetry.io/otel/metric v1.43.0 // indirect + go.opentelemetry.io/otel/sdk v1.43.0 // indirect + go.opentelemetry.io/otel/sdk/metric v1.43.0 // indirect + go.opentelemetry.io/otel/trace v1.43.0 // indirect + go.starlark.net v0.0.0-20260102030733-3fee463870c9 // indirect + go.yaml.in/yaml/v3 v3.0.4 // indirect + golang.org/x/arch v0.23.0 // indirect + golang.org/x/crypto v0.52.0 // indirect + golang.org/x/net v0.55.0 // indirect + golang.org/x/oauth2 v0.36.0 // indirect + golang.org/x/sync v0.20.0 // indirect + golang.org/x/sys v0.45.0 // indirect + golang.org/x/text v0.37.0 // indirect + golang.org/x/time v0.15.0 // indirect + google.golang.org/api v0.282.0 // indirect + google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20260401024825-9d38bb4040a9 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260523011958-0a33c5d7ca68 // indirect + google.golang.org/grpc v1.81.1 // indirect + google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect + gorm.io/driver/postgres v1.6.0 // indirect + gorm.io/driver/sqlite v1.6.0 // indirect + gorm.io/gorm v1.31.1 // indirect +) diff --git a/plugins/modelcatalogresolver/go.sum b/plugins/modelcatalogresolver/go.sum new file mode 100644 index 0000000000..599492593b --- /dev/null +++ b/plugins/modelcatalogresolver/go.sum @@ -0,0 +1,256 @@ +cel.dev/expr v0.25.1 h1:1KrZg61W6TWSxuNZ37Xy49ps13NUovb66QLprthtwi4= +cloud.google.com/go v0.123.0 h1:2NAUJwPR47q+E35uaJeYoNhuNEM9kM8SjgRgdeOJUSE= +cloud.google.com/go v0.123.0/go.mod h1:xBoMV08QcqUGuPW65Qfm1o9Y4zKZBpGS+7bImXLTAZU= +cloud.google.com/go/auth v0.20.0 h1:kXTssoVb4azsVDoUiF8KvxAqrsQcQtB53DcSgta74CA= +cloud.google.com/go/auth/oauth2adapt v0.2.8 h1:keo8NaayQZ6wimpNSmW5OPc283g65QNIiLpZnkHRbnc= +cloud.google.com/go/auth/oauth2adapt v0.2.8/go.mod h1:XQ9y31RkqZCcwJWNSx2Xvric3RrU88hAYYbjDWYDL+c= +cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs= +cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10= +cloud.google.com/go/iam v1.5.3 h1:+vMINPiDF2ognBJ97ABAYYwRgsaqxPbQDlMnbHMjolc= +cloud.google.com/go/logging v1.13.2 h1:qqlHCBvieJT9Cdq4QqYx1KPadCQ2noD4FK02eNqHAjA= +cloud.google.com/go/longrunning v0.8.0 h1:LiKK77J3bx5gDLi4SMViHixjD2ohlkwBi+mKA7EhfW8= +cloud.google.com/go/monitoring v1.24.3 h1:dde+gMNc0UhPZD1Azu6at2e79bfdztVDS5lvhOdsgaE= +cloud.google.com/go/storage v1.61.3 h1:VS//ZfBuPGDvakfD9xyPW1RGF1Vy3BWUoVZXgW1KMOg= +cloud.google.com/go/trace v1.11.7 h1:kDNDX8JkaAG3R2nq1lIdkb7FCSi1rCmsEtKVsty7p+U= +github.com/Azure/azure-sdk-for-go/sdk/azcore v1.20.0 h1:JXg2dwJUmPB9JmtVmdEB16APJ7jurfbY5jnfXpJoRMc= +github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.1 h1:Hk5QBxZQC1jb2Fwj6mpzme37xbCDdNTxU7O9eb5+LB4= +github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.13.1/go.mod h1:IYus9qsFobWIc2YVwe/WPjcnyCkPKtnHAqUYeebc8z0= +github.com/Azure/azure-sdk-for-go/sdk/azidentity/cache v0.3.2 h1:yz1bePFlP5Vws5+8ez6T3HWXPmwOK7Yvq8QxDBD3SKY= +github.com/Azure/azure-sdk-for-go/sdk/internal v1.11.2 h1:9iefClla7iYpfYWdzPCRDozdmndjTm8DXdpCzPajMgA= +github.com/AzureAD/microsoft-authentication-extensions-for-go/cache v0.1.1 h1:WJTmL004Abzc5wDB5VtZG2PJk5ndYDgVacGqfirKxjM= +github.com/AzureAD/microsoft-authentication-library-for-go v1.6.0 h1:XRzhVemXdgvJqCH0sFfrBUTnUJSBrBf7++ypk+twtRs= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.31.0 h1:DHa2U07rk8syqvCge0QIGMCE1WxGj9njT44GH7zNJLQ= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.55.0 h1:UnDZ/zFfG1JhH/DqxIZYU/1CUAlTUScoXD/LcM2Ykk8= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0.55.0 h1:7t/qx5Ost0s0wbA/VDrByOooURhp+ikYwv20i9Y07TQ= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.55.0 h1:0s6TxfCu2KHkkZPnBfsQ2y5qia0jl3MMrmBhu3nCOYk= +github.com/RaveNoX/go-jsoncommentstrip v1.0.0/go.mod h1:78ihd09MekBnJnxpICcwzCMzGrKSKYe4AqU6PDYYpjk= +github.com/andybalholm/brotli v1.2.1 h1:R+f5xP285VArJDRgowrfb9DqL18yVK0gKAW/F+eTWro= +github.com/apapsch/go-jsonmerge/v2 v2.0.0 h1:axGnT1gRIfimI7gJifB699GoE/oq+F2MU7Dml6nw9rQ= +github.com/apapsch/go-jsonmerge/v2 v2.0.0/go.mod h1:lvDnEdqiQrp0O42VQGgmlKpxL1AP2+08jFMw88y4klk= +github.com/aws/aws-sdk-go-v2 v1.41.7 h1:DWpAJt66FmnnaRIOT/8ASTucrvuDPZASqhhLey6tLY8= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.10 h1:gx1AwW1Iyk9Z9dD9F4akX5gnN3QZwUB20GGKH/I+Rho= +github.com/aws/aws-sdk-go-v2/config v1.32.11 h1:ftxI5sgz8jZkckuUHXfC/wMUc8u3fG1vQS0plr2F2Zs= +github.com/aws/aws-sdk-go-v2/credentials v1.19.14 h1:n+UcGWAIZHkXzYt87uMFBv/l8THYELoX6gVcUvgl6fI= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 h1:NUS3K4BTDArQqNu2ih7yeDLaS3bmHD0YndtA6UP884g= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.23 h1:GpT/TrnBYuE5gan2cZbTtvP+JlHsutdmlV2YfEyNde0= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.23 h1:bpd8vxhlQi2r1hiueOw02f/duEPTMK59Q4QMAoTTtTo= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.5 h1:clHU5fm//kWS1C2HgtgWxfQbFbx4b6rx+5jzhgX9HrI= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.22 h1:rWyie/PxDRIdhNf4DzRk0lvjVOqFJuNnO8WwaIRVxzQ= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 h1:5EniKhLZe4xzL7a+fU3C2tfUN4nWIqlLesfrjkuPFTY= +github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.13 h1:JRaIgADQS/U6uXDqlPiefP32yXTda7Kqfx+LgspooZM= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 h1:c31//R3xgIJMSC8S6hEVq+38DcvUlgFY0FM6mSI5oto= +github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.21 h1:ZlvrNcHSFFWURB8avufQq9gFsheUgjVD9536obIknfM= +github.com/aws/aws-sdk-go-v2/service/s3 v1.97.3 h1:HwxWTbTrIHm5qY+CAEur0s/figc3qwvLWsNkF4RPToo= +github.com/aws/aws-sdk-go-v2/service/signin v1.0.9 h1:QKZH0S178gCmFEgst8hN0mCX1KxLgHBKKY/CLqwP8lg= +github.com/aws/aws-sdk-go-v2/service/sso v1.30.15 h1:lFd1+ZSEYJZYvv9d6kXzhkZu07si3f+GQ1AaYwa2LUM= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.19 h1:dzztQ1YmfPrxdrOiuZRMF6fuOwWlWpD2StNLTceKpys= +github.com/aws/aws-sdk-go-v2/service/sts v1.41.10 h1:p8ogvvLugcR/zLBXTXrTkj0RYBUdErbMnAFFp12Lm/U= +github.com/aws/smithy-go v1.25.1 h1:J8ERsGSU7d+aCmdQur5Txg6bVoYelvQJgtZehD12GkI= +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/bmatcuk/doublestar v1.1.1/go.mod h1:UD6OnuiIn0yFxxA2le/rnRU1G4RaI4UvFv1sNto9p6w= +github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/buger/jsonparser v1.1.2 h1:frqHqw7otoVbk5M8LlE/L7HTnIq2v9RX6EJ48i9AxJk= +github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M= +github.com/bytedance/sonic v1.15.1 h1:nJD5PmM0vY7J8CT6MxoqbVAAMhkSmV2HgRAUrrpLoOw= +github.com/bytedance/sonic/loader v0.5.1 h1:Ygpfa9zwRCCKSlrp5bBP/b/Xzc3VxsAW+5NIYXrOOpI= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M= +github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2 h1:aBangftG7EVZoUb69Os8IaYg++6uMOdKK83QtkkvJik= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/envoyproxy/go-control-plane v0.14.0 h1:hbG2kr4RuFj222B6+7T83thSPqLjwBIfQawTkC++2HA= +github.com/envoyproxy/go-control-plane/envoy v1.37.0 h1:u3riX6BoYRfF4Dr7dwSOroNfdSbEPe9Yyl09/B6wBrQ= +github.com/envoyproxy/go-control-plane/ratelimit v0.1.0 h1:/G9QYbddjL25KvtKTv3an9lx6VBE2cnb8wp1vEGNYGI= +github.com/envoyproxy/protoc-gen-validate v1.3.3 h1:MVQghNeW+LZcmXe7SY1V36Z+WFMDjpqGAGacLe2T0ds= +github.com/fasthttp/websocket v1.5.12 h1:e4RGPpWW2HTbL3zV0Y/t7g0ub294LkiuXXUuTOUInlE= +github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= +github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= +github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= +github.com/go-jose/go-jose/v4 v4.1.4 h1:moDMcTHmvE6Groj34emNPLs/qtYXRVcd6S7NHbHz3kA= +github.com/go-jose/go-jose/v4 v4.1.4/go.mod h1:x4oUasVrzR7071A4TnHLGSPpNOm2a21K9Kf04k1rs08= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= +github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= +github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= +github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/go-openapi/analysis v0.24.2 h1:6p7WXEuKy1llDgOH8FooVeO+Uq2za9qoAOq4ZN08B50= +github.com/go-openapi/errors v0.22.5 h1:Yfv4O/PRYpNF3BNmVkEizcHb3uLVVsrDt3LNdgAKRY4= +github.com/go-openapi/jsonpointer v0.22.4 h1:dZtK82WlNpVLDW2jlA1YCiVJFVqkED1MegOUy9kR5T4= +github.com/go-openapi/jsonreference v0.21.4 h1:24qaE2y9bx/q3uRK/qN+TDwbok1NhbSmGjjySRCHtC8= +github.com/go-openapi/loads v0.23.2 h1:rJXAcP7g1+lWyBHC7iTY+WAF0rprtM+pm8Jxv1uQJp4= +github.com/go-openapi/runtime v0.29.2 h1:UmwSGWNmWQqKm1c2MGgXVpC2FTGwPDQeUsBMufc5Yj0= +github.com/go-openapi/spec v0.22.2 h1:KEU4Fb+Lp1qg0V4MxrSCPv403ZjBl8Lx1a83gIPU8Qc= +github.com/go-openapi/strfmt v0.25.0 h1:7R0RX7mbKLa9EYCTHRcCuIPcaqlyQiWNPTXwClK0saQ= +github.com/go-openapi/swag v0.25.4 h1:OyUPUFYDPDBMkqyxOTkqDYFnrhuhi9NR6QVUvIochMU= +github.com/go-openapi/swag/cmdutils v0.25.4 h1:8rYhB5n6WawR192/BfUu2iVlxqVR9aRgGJP6WaBoW+4= +github.com/go-openapi/swag/conv v0.25.4 h1:/Dd7p0LZXczgUcC/Ikm1+YqVzkEeCc9LnOWjfkpkfe4= +github.com/go-openapi/swag/fileutils v0.25.4 h1:2oI0XNW5y6UWZTC7vAxC8hmsK/tOkWXHJQH4lKjqw+Y= +github.com/go-openapi/swag/jsonname v0.25.4 h1:bZH0+MsS03MbnwBXYhuTttMOqk+5KcQ9869Vye1bNHI= +github.com/go-openapi/swag/jsonutils v0.25.4 h1:VSchfbGhD4UTf4vCdR2F4TLBdLwHyUDTd1/q4i+jGZA= +github.com/go-openapi/swag/jsonutils/fixtures_test v0.25.4 h1:IACsSvBhiNJwlDix7wq39SS2Fh7lUOCJRmx/4SN4sVo= +github.com/go-openapi/swag/loading v0.25.4 h1:jN4MvLj0X6yhCDduRsxDDw1aHe+ZWoLjW+9ZQWIKn2s= +github.com/go-openapi/swag/mangling v0.25.4 h1:2b9kBJk9JvPgxr36V23FxJLdwBrpijI26Bx5JH4Hp48= +github.com/go-openapi/swag/netutils v0.25.4 h1:Gqe6K71bGRb3ZQLusdI8p/y1KLgV4M/k+/HzVSqT8H0= +github.com/go-openapi/swag/stringutils v0.25.4 h1:O6dU1Rd8bej4HPA3/CLPciNBBDwZj9HiEpdVsb8B5A8= +github.com/go-openapi/swag/typeutils v0.25.4 h1:1/fbZOUN472NTc39zpa+YGHn3jzHWhv42wAJSN91wRw= +github.com/go-openapi/swag/yamlutils v0.25.4 h1:6jdaeSItEUb7ioS9lFoCZ65Cne1/RZtPBZ9A56h92Sw= +github.com/go-openapi/testify/enable/yaml/v2 v2.0.2 h1:0+Y41Pz1NkbTHz8NngxTuAXxEodtNSI1WG1c/m5Akw4= +github.com/go-openapi/testify/v2 v2.0.2 h1:X999g3jeLcoY8qctY/c/Z8iBHTbwLz7R2WXd6Ub6wls= +github.com/go-openapi/validate v0.25.1 h1:sSACUI6Jcnbo5IWqbYHgjibrhhmt3vR6lCzKZnmAgBw= +github.com/go-viper/mapstructure/v2 v2.5.0 h1:vM5IJoUAy3d7zRSVtIwQgBj7BiWtMPfmPEgAXnvj1Ro= +github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY= +github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/martian/v3 v3.3.3 h1:DIhPTQrbPkgs2yJYdXU/eNACCG5DVQjySNRNlflZ9Fc= +github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= +github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/googleapis/enterprise-certificate-proxy v0.3.16 h1:F/VPrx0YPBdksZJQdCAp0WUsqnNmZpUZszzfYt0M5Dw= +github.com/googleapis/gax-go/v2 v2.22.0 h1:PjIWBpgGIVKGoCXuiCoP64altEJCj3/Ei+kSU5vlZD4= +github.com/hajimehoshi/go-mp3 v0.3.4 h1:NUP7pBYH8OguP4diaTZ9wJbUbk3tC0KlfzsEpWmYj68= +github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= +github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= +github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= +github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM= +github.com/jackc/pgx/v5 v5.9.2 h1:3ZhOzMWnR4yJ+RW1XImIPsD1aNSz4T4fyP7zlQb56hw= +github.com/jackc/pgx/v5 v5.9.2/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4= +github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo= +github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4= +github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= +github.com/jinzhu/inflection v1.0.0/go.mod h1:h+uFLlag+Qp1Va5pdKtLDYj+kHp5pxUVkryuEj+Srlc= +github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ= +github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8= +github.com/juju/gnuflag v0.0.0-20171113085948-2ce1bb71843d/go.mod h1:2PavIy+JPciBPrBUjwbNvtwB6RQlve+hkpll6QSNmOE= +github.com/keybase/go-keychain v0.0.1 h1:way+bWYa6lDppZoZcgMbYsvC7GxljxrskdNInRtuthU= +github.com/klauspost/compress v1.18.6 h1:2jupLlAwFm95+YDR+NwD2MEfFO9d4z4Prjl1XXDjuao= +github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y= +github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= +github.com/mailru/easyjson v0.9.1 h1:LbtsOm5WAswyWbvTEOqhypdPeZzHavpZx96/n553mR8= +github.com/mark3labs/mcp-go v0.43.2 h1:21PUSlWWiSbUPQwXIJ5WKlETixpFpq+WBpbMGDSVy/I= +github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-sqlite3 v1.14.32 h1:JD12Ag3oLy1zQA+BNn74xRgaBbdhbNIDYvQUEuuErjs= +github.com/maximhq/bifrost/core v1.5.15 h1:iXvDufyZd7willDmbVzFfzCdy/2NsKEI1S8Iv9LjCSM= +github.com/maximhq/bifrost/core v1.5.15/go.mod h1:f6QHCvvzCQziMZ4JCNZP/GdZSeD50hww0vt7Uwl7lYY= +github.com/maximhq/bifrost/framework v1.3.15 h1:Lf/0S5bmD6i4NU+GdhAdrDZltD8RSvXLuR1vv4eh30I= +github.com/maximhq/bifrost/framework v1.3.15/go.mod h1:FlqWzdsFwal2XIG1Ousk/P76zjjHN/MmV5XR3YqD2+8= +github.com/oapi-codegen/runtime v1.1.1 h1:EXLHh0DXIJnWhdRPN2w4MXAzFyE4CskzhNLUmtpMYro= +github.com/oklog/ulid v1.3.1 h1:EGfNDEx6MqHz8B3uNV6QAib1UR2Lm97sHi3ocA6ESJ4= +github.com/oklog/ulid v1.3.1/go.mod h1:CirwcVhetQ6Lv90oh/F+FBtV6XMibvdAFo93nm5qn4U= +github.com/pinecone-io/go-pinecone/v5 v5.3.0 h1:0YQlEtmXGWK/I8ztkOVM6PuBYgFJZhjSdb0ddU+bHPE= +github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ= +github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo= +github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/qdrant/go-client v1.16.2 h1:UUMJJfvXTByhwhH1DwWdbkhZ2cTdvSqVkXSIfBrVWSg= +github.com/redis/go-redis/v9 v9.17.2 h1:P2EGsA4qVIM3Pp+aPocCJ7DguDHhqrXNhVcEp4ViluI= +github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= +github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0= +github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY= +github.com/savsgio/gotils v0.0.0-20250408102913-196191ec6287 h1:qIQ0tWF9vxGtkJa24bR+2i53WBCz1nW/Pc47oVYauC4= +github.com/spf13/cast v1.10.0 h1:h2x0u2shc1QuLHfxi+cTJvs30+ZAHOGRic8uyGTDWxY= +github.com/spf13/cast v1.10.0/go.mod h1:jNfB8QC9IA6ZuY2ZjDp0KtFO2LZZlg4S/7bzP6qqeHo= +github.com/spiffe/go-spiffe/v2 v2.6.0 h1:l+DolpxNWYgruGQVV0xsfeya3CsC7m8iBzDnMpsbLuo= +github.com/spiffe/go-spiffe/v2 v2.6.0/go.mod h1:gm2SeUoMZEtpnzPNs2Csc0D/gX33k1xIx7lEzqblHEs= +github.com/spkg/bom v0.0.0-20160624110644-59b7046e48ad/go.mod h1:qLr4V1qq6nMqFKkMo8ZTx3f+BZEkzsRUY10Xsm2mwU0= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/objx v0.5.3 h1:jmXUvGomnU1o3W/V5h2VEradbpJDwGrzugQQvL0POH4= +github.com/stretchr/objx v0.5.3/go.mod h1:rDQraq+vQZU7Fde9LOZLr8Tax6zZvy4kuNKF+QYS+U0= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= +github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= +github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= +github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs= +github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= +github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= +github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= +github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= +github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw= +github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= +github.com/valyala/fasthttp v1.71.0 h1:tepR7H+Guh9VUqxxcPggYi8R3lGUu2Rsdh+z7/FCY3k= +github.com/weaviate/weaviate v1.36.5 h1:lCiuEfQ08+5wK0DkTCUBb6ayNep9QpBH6JJhmZaRfzk= +github.com/weaviate/weaviate-go-client/v5 v5.7.1 h1:vEMxh486QqRqWaq58UEe/TiTbGbo9T5x7ZPFd5QENvQ= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= +go.mongodb.org/mongo-driver v1.17.6 h1:87JUG1wZfWsr6rIz3ZmpH90rL5tea7O3IHuSwHUpsss= +go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= +go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= +go.opentelemetry.io/contrib/detectors/gcp v1.42.0 h1:kpt2PEJuOuqYkPcktfJqWWDjTEd/FNgrxcniL7kQrXQ= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.67.0 h1:yI1/OhfEPy7J9eoa6Sj051C7n5dvpj0QX8g4sRchg04= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.67.0 h1:OyrsyzuttWTSur2qN/Lm0m2a8yqyIjUVBZcxFPuXq2o= +go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I= +go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.40.0 h1:ZrPRak/kS4xI3AVXy8F7pipuDXmDsrO8Lg+yQjBLjw0= +go.opentelemetry.io/otel/metric v1.43.0 h1:d7638QeInOnuwOONPp4JAOGfbCEpYb+K6DVWvdxGzgM= +go.opentelemetry.io/otel/sdk v1.43.0 h1:pi5mE86i5rTeLXqoF/hhiBtUNcrAGHLKQdhg4h4V9Dg= +go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfCGLEo89fDkw= +go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A= +go.starlark.net v0.0.0-20260102030733-3fee463870c9 h1:nV1OyvU+0CYrp5eKfQ3rD03TpFYYhH08z31NK1HmtTk= +go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= +go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= +golang.org/x/arch v0.23.0 h1:lKF64A2jF6Zd8L0knGltUnegD62JMFBiCPBmQpToHhg= +golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988= +golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc= +golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= +golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= +golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= +golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= +golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= +golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= +gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= +google.golang.org/api v0.282.0 h1:WmJiSVqUnKqJCpJOx7YADbXaC+9DDsnGSfllFSj7R2I= +google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 h1:XzmzkmB14QhVhgnawEVsOn6OFsnpyxNPRY9QV01dNB0= +google.golang.org/genproto/googleapis/api v0.0.0-20260401024825-9d38bb4040a9 h1:VPWxll4HlMw1Vs/qXtN7BvhZqsS9cdAittCNvVENElA= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260523011958-0a33c5d7ca68 h1:PvEgGJf9C/1u5CHkInMg7UFYYUoiaQmW2LbtH0pjB78= +google.golang.org/grpc v1.81.1 h1:VnnIIZ88UzOOKLukQi+ImGz8O1Wdp8nAGGnvOfEIWQQ= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gorm.io/driver/postgres v1.6.0 h1:2dxzU8xJ+ivvqTRph34QX+WrRaJlmfyPqXmoGVjMBa4= +gorm.io/driver/postgres v1.6.0/go.mod h1:vUw0mrGgrTK+uPHEhAdV4sfFELrByKVGnaVRkXDhtWo= +gorm.io/driver/sqlite v1.6.0 h1:WHRRrIiulaPiPFmDcod6prc4l2VGVWHz80KspNsxSfQ= +gorm.io/driver/sqlite v1.6.0/go.mod h1:AO9V1qIQddBESngQUKWL9yoH93HIeA1X6V633rBwyT8= +gorm.io/gorm v1.31.1 h1:7CA8FTFz/gRfgqgpeKIBcervUn3xSyPUmr6B2WXJ7kg= +gorm.io/gorm v1.31.1/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs= diff --git a/plugins/modelcatalogresolver/main.go b/plugins/modelcatalogresolver/main.go new file mode 100644 index 0000000000..104c6e9a77 --- /dev/null +++ b/plugins/modelcatalogresolver/main.go @@ -0,0 +1,191 @@ +// Package modelcatalogresolver provides a built-in PreRequestHook plugin that resolves +// the default provider for an unprefixed model via the model catalog. It is the single +// owner of "if no provider specified, look up which providers serve this model" — the +// transport handlers, integrations router, and realtime handlers no longer do this +// inline. Governance/LB plugins run before this resolver; it only fires as a final +// fallback when no earlier routing plugin picked a provider. +package modelcatalogresolver + +import ( + "fmt" + "slices" + "strings" + + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/modelcatalog" +) + +const PluginName = "model-catalog-resolver" + +// integrationTypeToDefaultProvider maps the integration-type ctx value (set by +// transports/bifrost-http/integrations/router.go on integration routes) to the +// integration's canonical provider. When the catalog returns multiple providers +// for an unprefixed model, the resolver prefers the integration's canonical +// provider if it's in the candidate list. +var integrationTypeToDefaultProvider = map[string]schemas.ModelProvider{ + "openai": schemas.OpenAI, + "anthropic": schemas.Anthropic, + "genai": schemas.Gemini, + "bedrock": schemas.Bedrock, + "cohere": schemas.Cohere, +} + +// Plugin resolves the default provider for unprefixed model strings using the model catalog. +type Plugin struct { + catalog *modelcatalog.ModelCatalog + logger schemas.Logger +} + +// Init returns a new resolver plugin. The catalog is required; if nil, the plugin returns +// an error rather than silently no-op'ing — a nil catalog at boot is a misconfiguration. +func Init(catalog *modelcatalog.ModelCatalog, logger schemas.Logger) (*Plugin, error) { + if catalog == nil { + return nil, fmt.Errorf("model-catalog-resolver: catalog is required") + } + return &Plugin{catalog: catalog, logger: logger}, nil +} + +// GetName implements schemas.BasePlugin. +func (p *Plugin) GetName() string { return PluginName } + +// Cleanup implements schemas.BasePlugin. +func (p *Plugin) Cleanup() error { return nil } + +// PreRequestHook fills in req.Provider from the model catalog when no provider was specified. +// Skips passthrough requests and requests that already have a provider set (e.g., from a model +// string like "openai/gpt-5", or from an earlier routing plugin — governance, LB). +// +// When the catalog returns multiple providers for an unprefixed model, the resolver prefers the +// integration's canonical provider (looked up from BifrostContextKeyIntegrationType set by the +// integration router) if it's in the candidate list. Otherwise it picks the first candidate. +// +// If the catalog returns zero providers, the resolver leaves req.Provider empty — the +// empty-provider validation in handleRequest/handleStreamRequest then returns a clear error. +func (p *Plugin) PreRequestHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) error { + if req.RequestType == schemas.PassthroughRequest || req.RequestType == schemas.PassthroughStreamRequest { + return nil + } + provider, model, existingFallbacks := req.GetRequestFields() + if provider != "" || model == "" { + return nil + } + + selected, candidates := ResolveProviderFromCatalog(ctx, p.catalog, model) + if selected == "" { + return nil + } + req.SetProvider(selected) + + candidateStrs := make([]string, len(candidates)) + for i, prov := range candidates { + candidateStrs[i] = string(prov) + } + ctx.AppendRoutingEngineLog(schemas.RoutingEngineModelCatalog, schemas.LogLevelInfo, fmt.Sprintf( + "No provider specified for model %s, found %d options in model catalog: [%s], selected: %s", + model, len(candidates), strings.Join(candidateStrs, ", "), selected, + )) + + // Populate fallbacks from the remaining catalog candidates so the request gets + // cross-provider resilience automatically — matches the governance and load + // balancing plugins, which both promote unselected candidates to fallbacks + // when the caller didn't configure any. Only fires when the caller passed + // none; an explicit fallback list (even an empty one set deliberately) is + // always respected. Model refinement is not needed here: GetProvidersForModel + // only returns providers that already serve this exact model string. + if len(existingFallbacks) == 0 && len(candidates) > 1 { + fallbacks := make([]schemas.Fallback, 0, len(candidates)-1) + for _, prov := range candidates { + if prov == selected { + continue + } + fallbacks = append(fallbacks, schemas.Fallback{Provider: prov, Model: model}) + } + if len(fallbacks) > 0 { + req.SetFallbacks(fallbacks) + fallbackStrs := make([]string, len(fallbacks)) + for i, fb := range fallbacks { + fallbackStrs[i] = string(fb.Provider) + } + ctx.AppendRoutingEngineLog(schemas.RoutingEngineModelCatalog, schemas.LogLevelInfo, fmt.Sprintf( + "Added %d catalog fallback provider(s) for model %s: [%s]", + len(fallbacks), model, strings.Join(fallbackStrs, ", "), + )) + } + } + + schemas.AppendToContextList(ctx, schemas.BifrostContextKeyRoutingEnginesUsed, schemas.RoutingEngineModelCatalog) + return nil +} + +// ResolveProviderFromCatalog performs the deterministic, integration-aware provider pick +// that PreRequestHook does, exposed for transport paths that can't run through +// PreRequestHook (realtime client_secrets, WebRTC). Returns the selected provider plus the +// full ordered candidate list. Returns ("", nil) when the catalog has no match for the model. +// +// The integration hint (BifrostContextKeyIntegrationType, when present and mapped) biases +// the pick toward the integration's canonical provider if it is in the candidate set; +// otherwise selection falls back to the alphabetically-first candidate for determinism. +// +// For requests routed through the openai integration whose user-agent identifies an Azure +// OpenAI SDK (BifrostContextKeyIsAzureUserAgent), schemas.Azure is preferred over +// schemas.OpenAI when Azure is in the candidate list — the openai-format converters no +// longer apply this default inline. +func ResolveProviderFromCatalog(ctx *schemas.BifrostContext, catalog *modelcatalog.ModelCatalog, model string) (schemas.ModelProvider, []schemas.ModelProvider) { + if catalog == nil || model == "" { + return "", nil + } + providers := catalog.GetProvidersForModel(model) + if len(providers) == 0 { + return "", nil + } + + // GetProvidersForModel iterates a Go map; the returned order is not stable. + // Sort alphabetically so the fallback pick (providers[0]) is deterministic across + // restarts and across processes — critical when no IntegrationType hint is set. + slices.SortFunc(providers, func(a, b schemas.ModelProvider) int { + return strings.Compare(string(a), string(b)) + }) + + selected := providers[0] + var integrationType string + var isAzureUser bool + if ctx != nil { + integrationType, _ = ctx.Value(schemas.BifrostContextKeyIntegrationType).(string) + isAzureUser, _ = ctx.Value(schemas.BifrostContextKeyIsAzureUserAgent).(bool) + } + if integrationType != "" { + if integrationDefault, mapped := integrationTypeToDefaultProvider[integrationType]; mapped && integrationDefault != "" { + preferred := integrationDefault + if integrationType == "openai" && isAzureUser { + preferred = schemas.Azure + } + if slices.Contains(providers, preferred) { + selected = preferred + } + + // For Anthropic-type routes, raw request body passthrough is only valid for + // providers that speak the Anthropic Messages API natively. When the model + // catalog falls back to a provider that doesn't (e.g. Bedrock), clear the + // flag so the provider performs its own format conversion. + if integrationType == "anthropic" && + selected != schemas.Anthropic && + selected != schemas.Vertex && + selected != schemas.Azure { + ctx.SetValue(schemas.BifrostContextKeyUseRawRequestBody, false) + ctx.SetValue(schemas.BifrostContextKeySendBackRawResponse, false) + ctx.SetValue(schemas.BifrostContextKeyPassthroughOverridesPresent, false) + } + } + } + return selected, providers +} + +// PreLLMHook implements schemas.LLMPlugin (no-op). +func (p *Plugin) PreLLMHook(_ *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { + return req, nil, nil +} + +// PostLLMHook implements schemas.LLMPlugin (no-op). +func (p *Plugin) PostLLMHook(_ *schemas.BifrostContext, resp *schemas.BifrostResponse, bifrostErr *schemas.BifrostError) (*schemas.BifrostResponse, *schemas.BifrostError, error) { + return resp, bifrostErr, nil +} diff --git a/transports/bifrost-http/handlers/inference.go b/transports/bifrost-http/handlers/inference.go index ed840445d9..18088cd112 100644 --- a/transports/bifrost-http/handlers/inference.go +++ b/transports/bifrost-http/handlers/inference.go @@ -56,27 +56,13 @@ func NewInferenceHandler(client *bifrost.Bifrost, config *lib.Config) *Completio } } -// resolveModelAndProvider parses the model string, validates it, and resolves -// the provider via model catalog when no provider prefix is present. Stores -// resolution metadata on the fasthttp context for ConvertToBifrostContext to -// emit the routing engine log. -func resolveModelAndProvider(ctx *fasthttp.RequestCtx, config *lib.Config, model string) (schemas.ModelProvider, string, error) { +// resolveModelAndProvider parses the model string. An empty provider is allowed here — +// the ModelCatalogResolver built-in PreRequestHook plugin fills it in as the last routing +// layer when no other routing plugin (governance routing rules, governance VK LB, enterprise +// LB) picked one. The empty-provider validation in handleRequest/handleStreamRequest catches +// the case where catalog resolution also fails. +func resolveModelAndProvider(_ *fasthttp.RequestCtx, _ *lib.Config, model string) (schemas.ModelProvider, string, error) { provider, modelName := schemas.ParseModelString(model, "") - if modelName == "" { - return "", "", fmt.Errorf("model is required") - } - if provider == "" { - providers := config.GetProvidersForModel(modelName) - if len(providers) == 0 { - return "", "", fmt.Errorf("provider is required in model field (format: provider/model) — no providers found for model %q in model catalog to auto-resolve", modelName) - } - ctx.SetUserValue(lib.FastHTTPUserValueModelCatalogResolution, &lib.ModelCatalogResolution{ - Model: modelName, - ResolvedProvider: providers[0], - AllProviders: providers, - }) - provider = providers[0] - } return provider, modelName, nil } diff --git a/transports/bifrost-http/handlers/realtime_client_secrets.go b/transports/bifrost-http/handlers/realtime_client_secrets.go index 6b8f680e15..4b500ce0bf 100644 --- a/transports/bifrost-http/handlers/realtime_client_secrets.go +++ b/transports/bifrost-http/handlers/realtime_client_secrets.go @@ -11,6 +11,7 @@ import ( bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/plugins/governance" + "github.com/maximhq/bifrost/plugins/modelcatalogresolver" "github.com/maximhq/bifrost/transports/bifrost-http/integrations" "github.com/maximhq/bifrost/transports/bifrost-http/lib" "github.com/valyala/fasthttp" @@ -231,14 +232,14 @@ func resolveRealtimeClientSecretTarget(ctx *fasthttp.RequestCtx, config *lib.Con providerKey, model := schemas.ParseModelString(rawModel, defaultProvider) // Model catalog auto-resolution for bare model names on /v1 client secret routes if defaultProvider == "" && providerKey == "" && model != "" { - providers := config.GetProvidersForModel(model) - if len(providers) > 0 { + selected, candidates := modelcatalogresolver.ResolveProviderFromCatalog(nil, config.ModelCatalog, model) + if selected != "" { ctx.SetUserValue(lib.FastHTTPUserValueModelCatalogResolution, &lib.ModelCatalogResolution{ Model: model, - ResolvedProvider: providers[0], - AllProviders: providers, + ResolvedProvider: selected, + AllProviders: candidates, }) - providerKey = providers[0] + providerKey = selected } } if defaultProvider == "" && providerKey == "" { diff --git a/transports/bifrost-http/handlers/webrtc_realtime.go b/transports/bifrost-http/handlers/webrtc_realtime.go index 6119ee4a44..3f88d88655 100644 --- a/transports/bifrost-http/handlers/webrtc_realtime.go +++ b/transports/bifrost-http/handlers/webrtc_realtime.go @@ -15,6 +15,7 @@ import ( bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/providers/openai" "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/plugins/modelcatalogresolver" "github.com/maximhq/bifrost/transports/bifrost-http/integrations" "github.com/maximhq/bifrost/transports/bifrost-http/lib" bfws "github.com/maximhq/bifrost/transports/bifrost-http/websocket" @@ -167,14 +168,14 @@ func parseCallsWebRTCRequest(ctx *fasthttp.RequestCtx, config *lib.Config) (stri providerKey, model := schemas.ParseModelString(rawModel, realtimeDefaultProviderForPath(path)) // Model catalog auto-resolution for bare model names on base /v1 routes if providerKey == "" && strings.TrimSpace(model) != "" { - providers := config.GetProvidersForModel(model) - if len(providers) > 0 { + selected, candidates := modelcatalogresolver.ResolveProviderFromCatalog(nil, config.ModelCatalog, model) + if selected != "" { ctx.SetUserValue(lib.FastHTTPUserValueModelCatalogResolution, &lib.ModelCatalogResolution{ Model: model, - ResolvedProvider: providers[0], - AllProviders: providers, + ResolvedProvider: selected, + AllProviders: candidates, }) - providerKey = providers[0] + providerKey = selected } } if providerKey == "" || strings.TrimSpace(model) == "" { @@ -199,14 +200,14 @@ func (h *WebRTCRealtimeHandler) handleLegacyRequest(ctx *fasthttp.RequestCtx, de providerKey, model := schemas.ParseModelString(rawModel, defaultProvider) // Model catalog auto-resolution for bare model names on base /v1 routes if providerKey == "" && strings.TrimSpace(model) != "" { - providers := h.config.GetProvidersForModel(model) - if len(providers) > 0 { + selected, candidates := modelcatalogresolver.ResolveProviderFromCatalog(nil, h.config.ModelCatalog, model) + if selected != "" { ctx.SetUserValue(lib.FastHTTPUserValueModelCatalogResolution, &lib.ModelCatalogResolution{ Model: model, - ResolvedProvider: providers[0], - AllProviders: providers, + ResolvedProvider: selected, + AllProviders: candidates, }) - providerKey = providers[0] + providerKey = selected } } if providerKey == "" || model == "" { @@ -1236,14 +1237,14 @@ func resolveRealtimeSDPTarget(ctx *fasthttp.RequestCtx, config *lib.Config, path providerKey, model := schemas.ParseModelString(strings.TrimSpace(rawModel), realtimeDefaultProviderForPath(path)) // Model catalog auto-resolution for bare model names in session body if providerKey == "" && strings.TrimSpace(model) != "" { - providers := config.GetProvidersForModel(model) - if len(providers) > 0 { + selected, candidates := modelcatalogresolver.ResolveProviderFromCatalog(nil, config.ModelCatalog, model) + if selected != "" { ctx.SetUserValue(lib.FastHTTPUserValueModelCatalogResolution, &lib.ModelCatalogResolution{ Model: model, - ResolvedProvider: providers[0], - AllProviders: providers, + ResolvedProvider: selected, + AllProviders: candidates, }) - providerKey = providers[0] + providerKey = selected } } if providerKey == "" || strings.TrimSpace(model) == "" { diff --git a/transports/bifrost-http/handlers/webrtc_realtime_test.go b/transports/bifrost-http/handlers/webrtc_realtime_test.go index 9fb8a1321a..d4b09b3f33 100644 --- a/transports/bifrost-http/handlers/webrtc_realtime_test.go +++ b/transports/bifrost-http/handlers/webrtc_realtime_test.go @@ -18,8 +18,7 @@ type testHandlerStore struct { kv *kvstore.Store } -func (s testHandlerStore) GetHeaderMatcher() *lib.HeaderMatcher { return nil } -func (s testHandlerStore) GetProvidersForModel(model string) []schemas.ModelProvider { return nil } +func (s testHandlerStore) GetHeaderMatcher() *lib.HeaderMatcher { return nil } func (s testHandlerStore) GetStreamChunkInterceptor() lib.StreamChunkInterceptor { return nil } diff --git a/transports/bifrost-http/handlers/wsrealtime.go b/transports/bifrost-http/handlers/wsrealtime.go index 856ded4746..761bc269ef 100644 --- a/transports/bifrost-http/handlers/wsrealtime.go +++ b/transports/bifrost-http/handlers/wsrealtime.go @@ -141,11 +141,26 @@ func (h *WSRealtimeHandler) handleUpgrade(ctx *fasthttp.RequestCtx) { }, } h.client.RunPreRequestHooks(preReqCtx, preReq) - if routedProvider, routedModel, _ := preReq.GetRequestFields(); routedProvider != "" { - providerKey = routedProvider - if routedModel != "" { - model = routedModel + routedProvider, routedModel, _ := preReq.GetRequestFields() + if routedProvider == "" { + // Mirror the empty-provider check in core handleRequest. No routing layer + // (governance routing rules / LB / modelcatalogresolver) could pick a provider + // for this model — caller's input is unresolvable. + upgrader := h.websocketUpgrader("") + upgradeErr := upgrader.Upgrade(ctx, func(conn *ws.Conn) { + defer conn.Close() + clientConn := newRealtimeClientConn(conn) + clientConn.writeRealtimeError(newRealtimeWireBifrostError(400, "invalid_request_error", fmt.Sprintf("no provider could be resolved for model %q (set as provider/model or configure the model catalog)", model))) + }) + if upgradeErr != nil { + logger.Warn("websocket upgrade failed for %s: %v", path, upgradeErr) } + preReqCancel() + return + } + providerKey = routedProvider + if routedModel != "" { + model = routedModel } // Mirror ctx values back to fasthttp user values so snapshotRealtimeMiddlewareValues // (called below) picks them up — same mechanism TransportInterceptorMiddleware uses. @@ -579,7 +594,7 @@ func (h *WSRealtimeHandler) relayRealtimeProviderToClient( } } -func resolveRealtimeTarget(ctx *fasthttp.RequestCtx, config *lib.Config, path, modelParam, deploymentParam string) (schemas.ModelProvider, string, error) { +func resolveRealtimeTarget(_ *fasthttp.RequestCtx, _ *lib.Config, path, modelParam, deploymentParam string) (schemas.ModelProvider, string, error) { defaultProvider := realtimeDefaultProviderForPath(path) var rawParam string @@ -597,22 +612,9 @@ func resolveRealtimeTarget(ctx *fasthttp.RequestCtx, config *lib.Config, path, m return "", "", errRealtimeModelFormat } - // Model catalog auto-resolution: when no provider prefix is present and the - // path doesn't imply a default provider, look up the model catalog — same - // logic as resolveModelAndProvider in inference.go. - if provider == "" { - providers := config.GetProvidersForModel(model) - if len(providers) == 0 { - return "", "", errRealtimeModelFormat - } - ctx.SetUserValue(lib.FastHTTPUserValueModelCatalogResolution, &lib.ModelCatalogResolution{ - Model: model, - ResolvedProvider: providers[0], - AllProviders: providers, - }) - provider = providers[0] - } - + // Provider may be empty here when no path-default applies and the model has + // no explicit prefix. The modelcatalogresolver PreRequestHook will fill it in + // (or surface a clear error if no provider matches) — no inline lookup needed. return provider, model, nil } @@ -834,13 +836,11 @@ var realtimeMiddlewareKeys = []any{ // snapshotRealtimeMiddlewareValues reads governance/routing values from the fasthttp // context's UserValue store. TransportInterceptorMiddleware copies them there as -// individual key-value pairs (not inside a BifrostContext). -// -// It also processes FastHTTPUserValueModelCatalogResolution, which is set by -// resolveRealtimeTarget when a bare model name is auto-resolved via the model -// catalog. ConvertToBifrostContext normally handles this for regular inference, -// but WebSocket handlers use createBifrostContextFromAuth instead, so we do the -// same log/engine enrichment here. +// individual key-value pairs (not inside a BifrostContext). Routing engine logs +// emitted by PreRequestHook (governance routing rules, LB, modelcatalogresolver) +// are surfaced through the same mechanism — the hooks write them onto preReqCtx +// and handleUpgrade mirrors that ctx's user values onto the fasthttp ctx before +// this function is called. func snapshotRealtimeMiddlewareValues(ctx *fasthttp.RequestCtx) map[any]any { result := make(map[any]any) for _, key := range realtimeMiddlewareKeys { @@ -848,33 +848,6 @@ func snapshotRealtimeMiddlewareValues(ctx *fasthttp.RequestCtx) map[any]any { result[key] = value } } - - // Model catalog auto-resolution: replicate the routing engine log that - // ConvertToBifrostContext would normally emit (see lib/ctx.go). - if res, ok := ctx.UserValue(lib.FastHTTPUserValueModelCatalogResolution).(*lib.ModelCatalogResolution); ok && res != nil { - providerStrs := make([]string, len(res.AllProviders)) - for i, p := range res.AllProviders { - providerStrs[i] = string(p) - } - logEntry := schemas.RoutingEngineLogEntry{ - Engine: schemas.RoutingEngineModelCatalog, - Level: schemas.LogLevelInfo, - Message: fmt.Sprintf("No provider specified for model %s, found %d options in model catalog: [%s], selecting first: %s", res.Model, len(res.AllProviders), strings.Join(providerStrs, ", "), res.ResolvedProvider), - Timestamp: time.Now().UnixMilli(), - } - // Merge with any existing routing engine logs from governance middleware. - if existing, ok := result[schemas.BifrostContextKeyRoutingEngineLogs].([]schemas.RoutingEngineLogEntry); ok { - result[schemas.BifrostContextKeyRoutingEngineLogs] = append(existing, logEntry) - } else { - result[schemas.BifrostContextKeyRoutingEngineLogs] = []schemas.RoutingEngineLogEntry{logEntry} - } - if existing, ok := result[schemas.BifrostContextKeyRoutingEnginesUsed].([]string); ok { - result[schemas.BifrostContextKeyRoutingEnginesUsed] = append(existing, schemas.RoutingEngineModelCatalog) - } else { - result[schemas.BifrostContextKeyRoutingEnginesUsed] = []string{schemas.RoutingEngineModelCatalog} - } - } - if len(result) == 0 { return nil } diff --git a/transports/bifrost-http/handlers/wsresponses_test.go b/transports/bifrost-http/handlers/wsresponses_test.go index 424061ab65..0cb2c7a0d7 100644 --- a/transports/bifrost-http/handlers/wsresponses_test.go +++ b/transports/bifrost-http/handlers/wsresponses_test.go @@ -24,10 +24,6 @@ func (s testWSHandlerStore) GetHeaderMatcher() *lib.HeaderMatcher { return s.matcher } -func (s testWSHandlerStore) GetProvidersForModel(model string) []schemas.ModelProvider { - return nil -} - func (s testWSHandlerStore) GetStreamChunkInterceptor() lib.StreamChunkInterceptor { return nil } diff --git a/transports/bifrost-http/integrations/anthropic.go b/transports/bifrost-http/integrations/anthropic.go index 87743291bd..627767d35f 100644 --- a/transports/bifrost-http/integrations/anthropic.go +++ b/transports/bifrost-http/integrations/anthropic.go @@ -23,18 +23,6 @@ type AnthropicRouter struct { *GenericRouter } -// anthropicModelGetter extracts the model field from any Anthropic integration request type. -// It is called after body parsing, so req is fully populated. -func anthropicModelGetter(_ *fasthttp.RequestCtx, req interface{}) (string, error) { - switch r := req.(type) { - case *anthropic.AnthropicTextRequest: - return r.Model, nil - case *anthropic.AnthropicMessageRequest: - return r.Model, nil - } - return "", nil -} - // createAnthropicCompleteRouteConfig creates a route configuration for the `/v1/complete` endpoint. func createAnthropicCompleteRouteConfig(pathPrefix string) RouteConfig { return RouteConfig{ @@ -47,7 +35,6 @@ func createAnthropicCompleteRouteConfig(pathPrefix string) RouteConfig { GetRequestTypeInstance: func(ctx context.Context) interface{} { return &anthropic.AnthropicTextRequest{} }, - GetRequestModel: anthropicModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if anthropicReq, ok := req.(*anthropic.AnthropicTextRequest); ok { return &schemas.BifrostRequest{ @@ -88,7 +75,6 @@ func createAnthropicMessagesRouteConfig(pathPrefix string, logger schemas.Logger GetRequestTypeInstance: func(ctx context.Context) interface{} { return &anthropic.AnthropicMessageRequest{} }, - GetRequestModel: anthropicModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if anthropicReq, ok := req.(*anthropic.AnthropicMessageRequest); ok { bifrostReq := anthropicReq.ToBifrostResponsesRequest(ctx) @@ -320,17 +306,14 @@ func checkAnthropicPassthrough(ctx *fasthttp.RequestCtx, bifrostCtx *schemas.Bif switch r := req.(type) { case *anthropic.AnthropicTextRequest: provider, model = schemas.ParseModelString(r.Model, "") - // Check if model parameter explicitly has `anthropic/` prefix + // Strip the explicit `anthropic/` prefix so downstream code sees the bare model. if provider == schemas.Anthropic { - bifrostCtx.SetValue(schemas.BifrostContextKeySkipModelCatalogProviderSelection, true) r.Model = model } case *anthropic.AnthropicMessageRequest: provider, model = schemas.ParseModelString(r.Model, "") - // Check if model parameter explicitly has `anthropic/` prefix if provider == schemas.Anthropic { - bifrostCtx.SetValue(schemas.BifrostContextKeySkipModelCatalogProviderSelection, true) r.Model = model } } @@ -425,7 +408,6 @@ func CreateAnthropicCountTokensRouteConfigs(pathPrefix string, handlerStore lib. GetRequestTypeInstance: func(ctx context.Context) interface{} { return &anthropic.AnthropicMessageRequest{} }, - GetRequestModel: anthropicModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if anthropicReq, ok := req.(*anthropic.AnthropicMessageRequest); ok { bifrostReq := anthropicReq.ToBifrostResponsesRequest(ctx) diff --git a/transports/bifrost-http/integrations/bedrock.go b/transports/bifrost-http/integrations/bedrock.go index 00da5932c8..efe2045261 100644 --- a/transports/bifrost-http/integrations/bedrock.go +++ b/transports/bifrost-http/integrations/bedrock.go @@ -20,23 +20,6 @@ type BedrockRouter struct { *GenericRouter } -// bedrockModelGetter extracts the model ID from any Bedrock integration request type. -// It is called after PreCallback, so req.ModelID is populated from the URL path param. -func bedrockModelGetter(_ *fasthttp.RequestCtx, req interface{}) (string, error) { - switch r := req.(type) { - case *bedrock.BedrockConverseRequest: - return r.ModelID, nil - case *bedrock.BedrockInvokeRequest: - return r.ModelID, nil - case *bedrock.BedrockCountTokensRequest: - if r.Input.Converse != nil { - return r.Input.Converse.ModelID, nil - } - return "", nil - } - return "", nil -} - // S3 context keys for storing request parameters const ( @@ -58,7 +41,6 @@ func createBedrockConverseRouteConfig(pathPrefix string, handlerStore lib.Handle GetHTTPRequestType: func(ctx *fasthttp.RequestCtx) schemas.RequestType { return schemas.ResponsesRequest }, - GetRequestModel: bedrockModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if bedrockReq, ok := req.(*bedrock.BedrockConverseRequest); ok { bifrostReq, err := bedrockReq.ToBifrostResponsesRequest(ctx) @@ -94,7 +76,6 @@ func createBedrockConverseStreamRouteConfig(pathPrefix string, handlerStore lib. GetRequestTypeInstance: func(ctx context.Context) interface{} { return &bedrock.BedrockConverseRequest{} }, - GetRequestModel: bedrockModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if bedrockReq, ok := req.(*bedrock.BedrockConverseRequest); ok { // Mark as streaming request @@ -145,7 +126,6 @@ func createBedrockInvokeWithResponseStreamRouteConfig(pathPrefix string, handler GetRequestTypeInstance: func(ctx context.Context) interface{} { return &bedrock.BedrockInvokeRequest{} }, - GetRequestModel: bedrockModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if invokeReq, ok := req.(*bedrock.BedrockInvokeRequest); ok { requestType, _ := ctx.Value(schemas.BifrostContextKeyHTTPRequestType).(schemas.RequestType) @@ -220,7 +200,6 @@ func createBedrockInvokeRouteConfig(pathPrefix string, handlerStore lib.HandlerS GetRequestTypeInstance: func(ctx context.Context) interface{} { return &bedrock.BedrockInvokeRequest{} }, - GetRequestModel: bedrockModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { invokeReq, ok := req.(*bedrock.BedrockInvokeRequest) if !ok { @@ -337,7 +316,6 @@ func createBedrockCountTokensRouteConfig(pathPrefix string, handlerStore lib.Han GetHTTPRequestType: func(ctx *fasthttp.RequestCtx) schemas.RequestType { return schemas.CountTokensRequest }, - GetRequestModel: bedrockModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if countTokensReq, ok := req.(*bedrock.BedrockCountTokensRequest); ok { if countTokensReq.Input.Converse == nil { diff --git a/transports/bifrost-http/integrations/bedrock_test.go b/transports/bifrost-http/integrations/bedrock_test.go index 42de6b3bc8..6b038c3758 100644 --- a/transports/bifrost-http/integrations/bedrock_test.go +++ b/transports/bifrost-http/integrations/bedrock_test.go @@ -25,10 +25,6 @@ func (m *mockHandlerStore) GetHeaderMatcher() *lib.HeaderMatcher { return m.headerMatcher } -func (m *mockHandlerStore) GetProvidersForModel(model string) []schemas.ModelProvider { - return m.availableProviders -} - func (m *mockHandlerStore) GetStreamChunkInterceptor() lib.StreamChunkInterceptor { return nil } diff --git a/transports/bifrost-http/integrations/cohere.go b/transports/bifrost-http/integrations/cohere.go index cf6b7ceaca..37aad1c1a8 100644 --- a/transports/bifrost-http/integrations/cohere.go +++ b/transports/bifrost-http/integrations/cohere.go @@ -69,22 +69,6 @@ func NewCohereRouter(client *bifrost.Bifrost, handlerStore lib.HandlerStore, log } } -// cohereModelGetter extracts the model field from any Cohere integration request type. -// It is called after body parsing, so req is fully populated. -func cohereModelGetter(_ *fasthttp.RequestCtx, req interface{}) (string, error) { - switch r := req.(type) { - case *cohere.CohereChatRequest: - return r.Model, nil - case *cohere.CohereEmbeddingRequest: - return r.Model, nil - case *cohere.CohereRerankRequest: - return r.Model, nil - case *cohere.CohereCountTokensRequest: - return r.Model, nil - } - return "", nil -} - // CreateCohereRouteConfigs creates route configurations for Cohere API endpoints. func CreateCohereRouteConfigs(pathPrefix string) []RouteConfig { var routes []RouteConfig @@ -101,7 +85,6 @@ func CreateCohereRouteConfigs(pathPrefix string) []RouteConfig { GetRequestTypeInstance: func(ctx context.Context) interface{} { return &cohere.CohereChatRequest{} }, - GetRequestModel: cohereModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if cohereReq, ok := req.(*cohere.CohereChatRequest); ok { return &schemas.BifrostRequest{ @@ -148,7 +131,6 @@ func CreateCohereRouteConfigs(pathPrefix string) []RouteConfig { GetRequestTypeInstance: func(ctx context.Context) interface{} { return &cohere.CohereEmbeddingRequest{} }, - GetRequestModel: cohereModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if cohereReq, ok := req.(*cohere.CohereEmbeddingRequest); ok { return &schemas.BifrostRequest{ @@ -182,7 +164,6 @@ func CreateCohereRouteConfigs(pathPrefix string) []RouteConfig { GetRequestTypeInstance: func(ctx context.Context) interface{} { return &cohere.CohereRerankRequest{} }, - GetRequestModel: cohereModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if cohereReq, ok := req.(*cohere.CohereRerankRequest); ok { return &schemas.BifrostRequest{ @@ -216,7 +197,6 @@ func CreateCohereRouteConfigs(pathPrefix string) []RouteConfig { GetRequestTypeInstance: func(ctx context.Context) interface{} { return &cohere.CohereCountTokensRequest{} }, - GetRequestModel: cohereModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if cohereReq, ok := req.(*cohere.CohereCountTokensRequest); ok { return &schemas.BifrostRequest{ diff --git a/transports/bifrost-http/integrations/genai.go b/transports/bifrost-http/integrations/genai.go index bc24f3fe27..5ac45fd198 100644 --- a/transports/bifrost-http/integrations/genai.go +++ b/transports/bifrost-http/integrations/genai.go @@ -38,33 +38,6 @@ type GenAIRouter struct { *GenericRouter } -// genAIModelGetter extracts the model name for GenAI routes. -// For request types populated by extractAndSetModelAndRequestType (the PreCallback), -// the model is already clean on the struct. For BifrostVideoRetrieveRequest (which has -// no model field), the provider-scoped model is extracted from the operation_id suffix -// (format: "op123:openai/gpt-4o") since the route pins the provider via operation_id. -func genAIModelGetter(ctx *fasthttp.RequestCtx, req interface{}) (string, error) { - switch r := req.(type) { - case *gemini.GeminiGenerationRequest: - return r.Model, nil - case *gemini.GeminiEmbeddingRequest: - return r.Model, nil - case *gemini.GeminiVideoGenerationRequest: - return r.Model, nil - case *gemini.GeminiBatchCreateRequest: - return r.Model, nil - case *schemas.BifrostVideoRetrieveRequest: - // operation_id encodes the full model string: "op123:gpt-4o" or "op123:openai/gpt-4o". - operationID, _ := ctx.UserValue("operation_id").(string) - parts := strings.Split(operationID, ":") - if len(parts) >= 2 && parts[len(parts)-1] != "" { - return parts[len(parts)-1], nil - } - return "", nil - } - return "", nil -} - // CreateGenAIRouteConfigs creates a route configurations for GenAI endpoints. func CreateGenAIRouteConfigs(pathPrefix string) []RouteConfig { var routes []RouteConfig @@ -81,7 +54,6 @@ func CreateGenAIRouteConfigs(pathPrefix string) []RouteConfig { GetRequestTypeInstance: func(ctx context.Context) interface{} { return &schemas.BifrostVideoRetrieveRequest{} }, - GetRequestModel: genAIModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if videoRetrieveReq, ok := req.(*schemas.BifrostVideoRetrieveRequest); ok { return &schemas.BifrostRequest{ @@ -120,7 +92,6 @@ func CreateGenAIRouteConfigs(pathPrefix string) []RouteConfig { } return &gemini.GeminiGenerationRequest{} }, - GetRequestModel: genAIModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if geminiReq, ok := req.(*gemini.GeminiGenerationRequest); ok { if geminiReq.IsCountTokens { @@ -826,12 +797,6 @@ func createGenAIRerankRouteConfig(pathPrefix string) RouteConfig { GetRequestTypeInstance: func(ctx context.Context) interface{} { return &vertex.VertexRankRequest{} }, - GetRequestModel: func(_ *fasthttp.RequestCtx, req interface{}) (string, error) { - if r, ok := req.(*vertex.VertexRankRequest); ok && r.Model != nil { - return *r.Model, nil - } - return "", nil - }, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if vertexReq, ok := req.(*vertex.VertexRankRequest); ok { return &schemas.BifrostRequest{ diff --git a/transports/bifrost-http/integrations/openai.go b/transports/bifrost-http/integrations/openai.go index 7b2bfa3d9c..0657a40372 100644 --- a/transports/bifrost-http/integrations/openai.go +++ b/transports/bifrost-http/integrations/openai.go @@ -269,43 +269,14 @@ func AzureEndpointPreHook(handlerStore lib.HandlerStore) func(ctx *fasthttp.Requ } } -// openAIModelGetter extracts the model field from any OpenAI integration request type. -// It is called after body parsing and PreCallback, so req is fully populated. -func openAIModelGetter(_ *fasthttp.RequestCtx, req interface{}) (string, error) { - switch r := req.(type) { - case *openai.OpenAIChatRequest: - return r.Model, nil - case *openai.OpenAITextCompletionRequest: - return r.Model, nil - case *openai.OpenAIEmbeddingRequest: - return r.Model, nil - case *openai.OpenAIResponsesRequest: - return r.Model, nil - case *openai.OpenAISpeechRequest: - return r.Model, nil - case *openai.OpenAITranscriptionRequest: - return r.Model, nil - case *openai.OpenAIImageGenerationRequest: - return r.Model, nil - case *openai.OpenAIImageEditRequest: - return r.Model, nil - case *openai.OpenAIImageVariationRequest: - return r.Model, nil - case *openai.OpenAIVideoGenerationRequest: - return r.Model, nil - } - return "", nil -} - // CreateOpenAIRouteConfigs creates route configurations for OpenAI endpoints. func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) []RouteConfig { var routes []RouteConfig routes = append(routes, RouteConfig{ - Type: RouteConfigTypeOpenAI, - Path: pathPrefix + "/openai/deployments/{deploymentPath:*}", - Method: "POST", - GetRequestModel: openAIModelGetter, + Type: RouteConfigTypeOpenAI, + Path: pathPrefix + "/openai/deployments/{deploymentPath:*}", + Method: "POST", GetHTTPRequestType: func(ctx *fasthttp.RequestCtx) schemas.RequestType { deploymentPathVal, ok := ctx.UserValue("deploymentPath").(string) if !ok { @@ -578,7 +549,6 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAIChatRequest{} }, - GetRequestModel: openAIModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if openaiReq, ok := req.(*openai.OpenAIChatRequest); ok { br := &schemas.BifrostRequest{ @@ -675,7 +645,6 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAITextCompletionRequest{} }, - GetRequestModel: openAIModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if openaiReq, ok := req.(*openai.OpenAITextCompletionRequest); ok { return &schemas.BifrostRequest{ @@ -727,7 +696,6 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAIResponsesRequest{} }, - GetRequestModel: openAIModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if openaiReq, ok := req.(*openai.OpenAIResponsesRequest); ok { return &schemas.BifrostRequest{ @@ -810,7 +778,6 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAIResponsesRequest{} }, - GetRequestModel: openAIModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if openaiReq, ok := req.(*openai.OpenAIResponsesRequest); ok { return &schemas.BifrostRequest{ @@ -840,9 +807,9 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) "/openai/responses/compact", } { routes = append(routes, RouteConfig{ - Type: RouteConfigTypeOpenAI, - Path: pathPrefix + path, - Method: "POST", + Type: RouteConfigTypeOpenAI, + Path: pathPrefix + path, + Method: "POST", PreCallback: func(ctx *fasthttp.RequestCtx, bifrostCtx *schemas.BifrostContext, req interface{}) error { hydrateOpenAIRequestFromLargePayloadMetadata(ctx, bifrostCtx, req) schemas.ExtractAndSetUserAgentFromHeaders(extractHeadersFromRequest(ctx), bifrostCtx) @@ -857,12 +824,6 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAICompactionRequest{} }, - GetRequestModel: func(_ *fasthttp.RequestCtx, req interface{}) (string, error) { - if r, ok := req.(*openai.OpenAICompactionRequest); ok { - return r.Model, nil - } - return "", nil - }, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if r, ok := req.(*openai.OpenAICompactionRequest); ok { return &schemas.BifrostRequest{ @@ -901,7 +862,6 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAIEmbeddingRequest{} }, - GetRequestModel: openAIModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if embeddingReq, ok := req.(*openai.OpenAIEmbeddingRequest); ok { return &schemas.BifrostRequest{ @@ -940,7 +900,6 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAISpeechRequest{} }, - GetRequestModel: openAIModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if speechReq, ok := req.(*openai.OpenAISpeechRequest); ok { return &schemas.BifrostRequest{ @@ -984,8 +943,7 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAITranscriptionRequest{} }, - GetRequestModel: openAIModelGetter, - RequestParser: parseTranscriptionMultipartRequest, // Handle multipart form parsing + RequestParser: parseTranscriptionMultipartRequest, // Handle multipart form parsing RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if transcriptionReq, ok := req.(*openai.OpenAITranscriptionRequest); ok { return &schemas.BifrostRequest{ @@ -1040,7 +998,6 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAIImageGenerationRequest{} }, - GetRequestModel: openAIModelGetter, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if imageGenReq, ok := req.(*openai.OpenAIImageGenerationRequest); ok { return &schemas.BifrostRequest{ @@ -1091,8 +1048,7 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAIImageEditRequest{} }, - GetRequestModel: openAIModelGetter, - RequestParser: parseOpenAIImageEditMultipartRequest, // Handle multipart form parsing + RequestParser: parseOpenAIImageEditMultipartRequest, // Handle multipart form parsing RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if imageEditReq, ok := req.(*openai.OpenAIImageEditRequest); ok { return &schemas.BifrostRequest{ @@ -1142,8 +1098,7 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAIImageVariationRequest{} }, - GetRequestModel: openAIModelGetter, - RequestParser: parseOpenAIImageVariationMultipartRequest, + RequestParser: parseOpenAIImageVariationMultipartRequest, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if imageVariationReq, ok := req.(*openai.OpenAIImageVariationRequest); ok { return &schemas.BifrostRequest{ @@ -1195,8 +1150,7 @@ func CreateOpenAIRouteConfigs(pathPrefix string, handlerStore lib.HandlerStore) GetRequestTypeInstance: func(ctx context.Context) interface{} { return &openai.OpenAIVideoGenerationRequest{} }, - GetRequestModel: openAIModelGetter, - RequestParser: parseOpenAIVideoGenerationMultipartRequest, + RequestParser: parseOpenAIVideoGenerationMultipartRequest, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { if videoGenerationReq, ok := req.(*openai.OpenAIVideoGenerationRequest); ok { return &schemas.BifrostRequest{ diff --git a/transports/bifrost-http/integrations/router.go b/transports/bifrost-http/integrations/router.go index 6ba2fc8b5d..a07f56cf6c 100644 --- a/transports/bifrost-http/integrations/router.go +++ b/transports/bifrost-http/integrations/router.go @@ -55,7 +55,6 @@ import ( "io" "mime" "mime/multipart" - "slices" "strconv" "strings" @@ -397,10 +396,6 @@ type PostRequestCallback func(ctx *fasthttp.RequestCtx, req interface{}, resp in // returns a schemas.RequestType indicating the HTTP request type derived from the context. type HTTPRequestTypeGetter func(ctx *fasthttp.RequestCtx) schemas.RequestType -// RequestModelGetter is a function type that accepts only a *fasthttp.RequestCtx and -// returns a string indicating the model derived from the context. -type RequestModelGetter func(ctx *fasthttp.RequestCtx, req interface{}) (string, error) - // ShortCircuit is a function that determines if the request should be short-circuited. type ShortCircuit func(ctx *fasthttp.RequestCtx, bifrostCtx *schemas.BifrostContext, req interface{}) (bool, error) @@ -444,14 +439,6 @@ const ( RouteConfigTypeCohere RouteConfigType = "cohere" ) -var RouteConfigTypeToProvider = map[RouteConfigType]schemas.ModelProvider{ - RouteConfigTypeOpenAI: schemas.OpenAI, - RouteConfigTypeAnthropic: schemas.Anthropic, - RouteConfigTypeGenAI: schemas.Gemini, - RouteConfigTypeBedrock: schemas.Bedrock, - RouteConfigTypeCohere: schemas.Cohere, -} - // RouteConfig defines the configuration for a single route in an integration. // It specifies the path, method, and handlers for request/response conversion. type RouteConfig struct { @@ -459,7 +446,6 @@ type RouteConfig struct { Path string // HTTP path pattern (e.g., "/openai/v1/chat/completions") Method string // HTTP method (POST, GET, PUT, DELETE) GetHTTPRequestType HTTPRequestTypeGetter // Function to get the HTTP request type from the context (SHOULD NOT BE NIL) - GetRequestModel RequestModelGetter // Function to get the model from the context (SHOULD NOT BE NIL) GetRequestTypeInstance func(ctx context.Context) interface{} // Factory function to create request instance (SHOULD NOT BE NIL) RequestParser RequestParser // Optional: custom request parsing (e.g., multipart/form-data) RequestConverter RequestConverter // Function to convert request to BifrostRequest (for inference requests) @@ -690,7 +676,9 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle } }() - // Set integration type to context + // Set integration type to context. Used by the ModelCatalogResolver built-in + // PreRequestHook (last routing layer) to prefer this integration's canonical + // provider when the model is unprefixed and the catalog returns multiple options. bifrostCtx.SetValue(schemas.BifrostContextKeyIntegrationType, string(config.Type)) // Async retrieve: check x-bf-async-id header early (before body parsing) @@ -781,75 +769,6 @@ func (g *GenericRouter) createHandler(config RouteConfig) fasthttp.RequestHandle } } - // Set available providers to context - if config.GetRequestModel != nil { - model, err := config.GetRequestModel(ctx, req) - if err != nil { - g.sendError(ctx, bifrostCtx, config.ErrorConverter, newBifrostError(err, "failed to get model from context")) - return - } - extractedProvider, extractedModel := schemas.ParseModelString(model, "") - // Skip model-catalog when governance already made a routing decision. - // Governance uses dot-notation aliases (e.g. "anthropic.claude-sonnet-4-6") which - // ParseModelString cannot extract a provider from (it only handles slash separators), - // causing a spurious model-catalog lookup that can override governance's selection. - skipModelCatalogProviderSelection, _ := bifrostCtx.Value(schemas.BifrostContextKeySkipModelCatalogProviderSelection).(bool) - if extractedProvider == "" && !skipModelCatalogProviderSelection { - availableProviders := g.handlerStore.GetProvidersForModel(extractedModel) - existingProviders, hasExistingProviders := bifrostCtx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - if hasExistingProviders { - if len(existingProviders) == 0 { - availableProviders = []schemas.ModelProvider{} - } else if len(availableProviders) == 0 { - availableProviders = existingProviders - } else { - availableProviders = slices.DeleteFunc(availableProviders, func(provider schemas.ModelProvider) bool { - return !slices.Contains(existingProviders, provider) - }) - } - } - availableProvidersStrs := make([]string, len(availableProviders)) - for i, p := range availableProviders { - availableProvidersStrs[i] = string(p) - } - bifrostCtx.AppendRoutingEngineLog(schemas.RoutingEngineModelCatalog, schemas.LogLevelInfo, fmt.Sprintf( - "No provider specified for model %s, found %d options in model catalog: [%s]", - extractedModel, len(availableProviders), strings.Join(availableProvidersStrs, ", "), - )) - if len(availableProviders) > 0 { - if slices.Contains(availableProviders, RouteConfigTypeToProvider[config.Type]) { - availableProviders = []schemas.ModelProvider{RouteConfigTypeToProvider[config.Type]} - bifrostCtx.AppendRoutingEngineLog(schemas.RoutingEngineModelCatalog, schemas.LogLevelInfo, fmt.Sprintf( - "Integration route default provider %s is found in the available providers list, selecting it", - RouteConfigTypeToProvider[config.Type], - )) - } else { - bifrostCtx.AppendRoutingEngineLog(schemas.RoutingEngineModelCatalog, schemas.LogLevelInfo, fmt.Sprintf( - "Integration route default provider %s is not found in the available providers list, selecting first: %s", - RouteConfigTypeToProvider[config.Type], availableProviders[0], - )) - // For Anthropic-type routes, raw request body passthrough is only valid for - // providers that speak the Anthropic Messages API natively. When the model - // catalog falls back to a provider that doesn't (e.g. Bedrock), clear the - // flag so the provider performs its own format conversion. - firstProvider := availableProviders[0] - if config.Type == RouteConfigTypeAnthropic && - firstProvider != schemas.Anthropic && - firstProvider != schemas.Vertex && - firstProvider != schemas.Azure { - bifrostCtx.SetValue(schemas.BifrostContextKeyUseRawRequestBody, false) - bifrostCtx.SetValue(schemas.BifrostContextKeySendBackRawResponse, false) - bifrostCtx.SetValue(schemas.BifrostContextKeyPassthroughOverridesPresent, false) - } - } - bifrostCtx.SetValue(schemas.BifrostContextKeyAvailableProviders, availableProviders) - } else if hasExistingProviders { - bifrostCtx.SetValue(schemas.BifrostContextKeyAvailableProviders, []schemas.ModelProvider{}) - } - schemas.AppendToContextList(bifrostCtx, schemas.BifrostContextKeyRoutingEnginesUsed, schemas.RoutingEngineModelCatalog) - } - } - // Handle batch requests if BatchRequestConverter is set // GenAI has two cases: (1) Dedicated batch routes (list/retrieve) have only BatchRequestConverter — always use batch path. // (2) The models path has both BatchRequestConverter and RequestConverter — use batch path only for batch create. diff --git a/transports/bifrost-http/integrations/router_test.go b/transports/bifrost-http/integrations/router_test.go index f2bef00469..7acbc15318 100644 --- a/transports/bifrost-http/integrations/router_test.go +++ b/transports/bifrost-http/integrations/router_test.go @@ -377,95 +377,6 @@ func TestOpenAIChatStructuredOutputRequestParserAndConverter(t *testing.T) { assert.Contains(t, responseFormat, "json_schema") } -func TestCreateHandler_AnthropicRouteConstrainsCatalogProvidersWhenAvailableProvidersSet(t *testing.T) { - handlerStore := &mockHandlerStore{ - availableProviders: []schemas.ModelProvider{ - schemas.Anthropic, - schemas.Azure, - schemas.Bedrock, - schemas.Vertex, - }, - } - - var capturedProviders []schemas.ModelProvider - route := RouteConfig{ - Type: RouteConfigTypeAnthropic, - Path: "/v1/messages", - Method: fasthttp.MethodPost, - GetHTTPRequestType: func(ctx *fasthttp.RequestCtx) schemas.RequestType { - return schemas.ResponsesRequest - }, - GetRequestTypeInstance: func(ctx context.Context) interface{} { - return &anthropic.AnthropicMessageRequest{} - }, - GetRequestModel: anthropicModelGetter, - PreCallback: checkAnthropicPassthrough, - RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { - capturedProviders, _ = ctx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - return nil, fmt.Errorf("stop before bifrost execution") - }, - ErrorConverter: func(ctx *schemas.BifrostContext, err *schemas.BifrostError) interface{} { - return err - }, - } - - router := NewGenericRouter(nil, handlerStore, nil, nil, nil) - ctx := &fasthttp.RequestCtx{} - ctx.Request.Header.SetMethod(fasthttp.MethodPost) - ctx.SetUserValue(schemas.BifrostContextKeyAvailableProviders, []schemas.ModelProvider{ - schemas.Azure, - schemas.OpenAI, - schemas.Ollama, - }) - ctx.Request.SetBodyString(`{"model":"claude-opus-4-8","max_tokens":1024,"messages":[{"role":"user","content":"hi"}]}`) - - router.createHandler(route)(ctx) - - require.Equal(t, fasthttp.StatusInternalServerError, ctx.Response.StatusCode()) - require.Equal(t, []schemas.ModelProvider{schemas.Azure}, capturedProviders) -} - -func TestCreateHandler_AnthropicRouteKeepsCatalogProvidersWhenAvailableProvidersUnset(t *testing.T) { - handlerStore := &mockHandlerStore{ - availableProviders: []schemas.ModelProvider{ - schemas.Bedrock, - schemas.Vertex, - }, - } - - var capturedProviders []schemas.ModelProvider - route := RouteConfig{ - Type: RouteConfigTypeAnthropic, - Path: "/v1/messages", - Method: fasthttp.MethodPost, - GetHTTPRequestType: func(ctx *fasthttp.RequestCtx) schemas.RequestType { - return schemas.ResponsesRequest - }, - GetRequestTypeInstance: func(ctx context.Context) interface{} { - return &anthropic.AnthropicMessageRequest{} - }, - GetRequestModel: anthropicModelGetter, - PreCallback: checkAnthropicPassthrough, - RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { - capturedProviders, _ = ctx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - return nil, fmt.Errorf("stop before bifrost execution") - }, - ErrorConverter: func(ctx *schemas.BifrostContext, err *schemas.BifrostError) interface{} { - return err - }, - } - - router := NewGenericRouter(nil, handlerStore, nil, nil, nil) - ctx := &fasthttp.RequestCtx{} - ctx.Request.Header.SetMethod(fasthttp.MethodPost) - ctx.Request.SetBodyString(`{"model":"claude-opus-4-8","max_tokens":1024,"messages":[{"role":"user","content":"hi"}]}`) - - router.createHandler(route)(ctx) - - require.Equal(t, fasthttp.StatusInternalServerError, ctx.Response.StatusCode()) - require.Equal(t, []schemas.ModelProvider{schemas.Bedrock, schemas.Vertex}, capturedProviders) -} - func TestCreateHandler_AnthropicRouteClears_UseRawRequestBody_WhenCatalogSelectsBedrock(t *testing.T) { handlerStore := &mockHandlerStore{ availableProviders: []schemas.ModelProvider{schemas.Bedrock}, @@ -484,8 +395,7 @@ func TestCreateHandler_AnthropicRouteClears_UseRawRequestBody_WhenCatalogSelects GetRequestTypeInstance: func(ctx context.Context) interface{} { return &anthropic.AnthropicMessageRequest{} }, - GetRequestModel: anthropicModelGetter, - PreCallback: checkAnthropicPassthrough, + PreCallback: checkAnthropicPassthrough, RequestConverter: func(ctx *schemas.BifrostContext, req interface{}) (*schemas.BifrostRequest, error) { capturedUseRaw = ctx.Value(schemas.BifrostContextKeyUseRawRequestBody) capturedSendRawResponse = ctx.Value(schemas.BifrostContextKeySendBackRawResponse) diff --git a/transports/bifrost-http/integrations/utils.go b/transports/bifrost-http/integrations/utils.go index 29be18477c..e4c54590b6 100644 --- a/transports/bifrost-http/integrations/utils.go +++ b/transports/bifrost-http/integrations/utils.go @@ -5,7 +5,6 @@ import ( "fmt" "net/url" "reflect" - "slices" "strconv" "strings" @@ -349,11 +348,6 @@ func (g *GenericRouter) extractAndParseFallbacks(ctx *schemas.BifrostContext, re } provider, _, _ := bifrostReq.GetRequestFields() - var availableProviders []schemas.ModelProvider - var hasAvailableProviders bool - if ctx != nil { - availableProviders, hasAvailableProviders = ctx.Value(schemas.BifrostContextKeyAvailableProviders).([]schemas.ModelProvider) - } // Parse fallbacks from strings to Fallback structs parsedFallbacks := make([]schemas.Fallback, 0, len(fallbacks)) @@ -364,9 +358,6 @@ func (g *GenericRouter) extractAndParseFallbacks(ctx *schemas.BifrostContext, re // Use ParseModelString to extract provider and model provider, model := schemas.ParseModelString(fallbackStr, provider) - if hasAvailableProviders && !slices.Contains(availableProviders, provider) { - continue - } parsedFallback := schemas.Fallback{ Provider: provider, @@ -376,7 +367,6 @@ func (g *GenericRouter) extractAndParseFallbacks(ctx *schemas.BifrostContext, re } if len(parsedFallbacks) == 0 { - bifrostReq.SetFallbacks(nil) return nil // No valid fallbacks found } diff --git a/transports/bifrost-http/integrations/utils_test.go b/transports/bifrost-http/integrations/utils_test.go index e1e1dd09d3..50fe7e8216 100644 --- a/transports/bifrost-http/integrations/utils_test.go +++ b/transports/bifrost-http/integrations/utils_test.go @@ -69,59 +69,6 @@ func TestExtractAndParseFallbacks_GeminiGenerationRequest(t *testing.T) { assert.Equal(t, "gemini-3-flash-preview", bifrostReq.ResponsesRequest.Fallbacks[0].Model) } -func TestExtractAndParseFallbacks_FiltersByAvailableProviders(t *testing.T) { - router := newTestGenericRouter() - geminiReq := &gemini.GeminiGenerationRequest{ - Model: "gemini/gemini-3-flash-preview", - Fallbacks: []string{ - "azure/claude-opus-4-8", - "bedrock/claude-opus-4-8", - "vertex/claude-opus-4-8", - }, - } - bifrostReq := &schemas.BifrostRequest{ - ResponsesRequest: &schemas.BifrostResponsesRequest{ - Provider: schemas.Gemini, - Model: "gemini-3-flash-preview", - }, - } - ctx := newTestBifrostContext() - ctx.SetValue(schemas.BifrostContextKeyAvailableProviders, []schemas.ModelProvider{schemas.Azure}) - - err := router.extractAndParseFallbacks(ctx, geminiReq, bifrostReq) - - require.NoError(t, err) - require.NotNil(t, bifrostReq.ResponsesRequest) - require.Len(t, bifrostReq.ResponsesRequest.Fallbacks, 1) - assert.Equal(t, schemas.Azure, bifrostReq.ResponsesRequest.Fallbacks[0].Provider) - assert.Equal(t, "claude-opus-4-8", bifrostReq.ResponsesRequest.Fallbacks[0].Model) -} - -func TestExtractAndParseFallbacks_ClearsDisallowedPreparsedFallbacks(t *testing.T) { - router := newTestGenericRouter() - geminiReq := &gemini.GeminiGenerationRequest{ - Model: "gemini/gemini-3-flash-preview", - Fallbacks: []string{"bedrock/claude-opus-4-8"}, - } - bifrostReq := &schemas.BifrostRequest{ - ResponsesRequest: &schemas.BifrostResponsesRequest{ - Provider: schemas.Gemini, - Model: "gemini-3-flash-preview", - Fallbacks: []schemas.Fallback{ - {Provider: schemas.Bedrock, Model: "claude-opus-4-8"}, - }, - }, - } - ctx := newTestBifrostContext() - ctx.SetValue(schemas.BifrostContextKeyAvailableProviders, []schemas.ModelProvider{schemas.Azure}) - - err := router.extractAndParseFallbacks(ctx, geminiReq, bifrostReq) - - require.NoError(t, err) - require.NotNil(t, bifrostReq.ResponsesRequest) - require.Empty(t, bifrostReq.ResponsesRequest.Fallbacks) -} - // TestSendStreamError_PropagatesProviderStatusCode verifies that sendStreamError // sets the HTTP status code from the provider's BifrostError.StatusCode field. // All three providers (OpenAI, Anthropic, Bedrock) return actual HTTP error codes diff --git a/transports/bifrost-http/lib/config.go b/transports/bifrost-http/lib/config.go index c9a05f5d97..685a70eba1 100644 --- a/transports/bifrost-http/lib/config.go +++ b/transports/bifrost-http/lib/config.go @@ -65,8 +65,6 @@ type StreamChunkInterceptor interface { type HandlerStore interface { // GetHeaderMatcher returns the precompiled header matcher for header filtering GetHeaderMatcher() *HeaderMatcher - // GetProvidersForModel returns the list of providers that can serve a given model. - GetProvidersForModel(model string) []schemas.ModelProvider // GetStreamChunkInterceptor returns the interceptor for streaming chunks. // Returns nil if no plugins are loaded or streaming interception is not needed. GetStreamChunkInterceptor() StreamChunkInterceptor @@ -4373,28 +4371,6 @@ func (c *Config) GetAllowOnAllVirtualKeysClients() map[string]string { return result } -// GetProvidersForModel returns the list of providers for a given model, sorted -// deterministically so callers picking providers[0] always get the same result. -func (c *Config) GetProvidersForModel(model string) []schemas.ModelProvider { - if c.ModelCatalog == nil { - return []schemas.ModelProvider{} - } - providersInCatalog := c.ModelCatalog.GetProvidersForModel(model) - // Filter out the providers which are not present in the configured provider list for the client - c.Mu.RLock() - defer c.Mu.RUnlock() - allowedProviders := make([]schemas.ModelProvider, 0, len(providersInCatalog)) - for configuredProvider := range c.Providers { - if slices.Contains(providersInCatalog, configuredProvider) { - allowedProviders = append(allowedProviders, configuredProvider) - } - } - slices.SortFunc(allowedProviders, func(a, b schemas.ModelProvider) int { - return strings.Compare(string(a), string(b)) - }) - return allowedProviders -} - // GetPluginOrder returns the names of all base plugins in their sorted placement order. // This method is lock-free and safe for concurrent access from hot paths. // Do not modify the returned slice; it is a shared snapshot and must be treated read-only. diff --git a/transports/bifrost-http/lib/ctx.go b/transports/bifrost-http/lib/ctx.go index 79e1d50f7f..c86e023301 100644 --- a/transports/bifrost-http/lib/ctx.go +++ b/transports/bifrost-http/lib/ctx.go @@ -33,8 +33,9 @@ const ( // It is used by transport middleware to avoid re-buffering response bodies for post-hooks. FastHTTPUserValueLargeResponseMode = "__bifrost_large_response_mode" // FastHTTPUserValueModelCatalogResolution stores model catalog resolution metadata - // set by prepare*Request functions when a provider was auto-resolved. Picked up - // centrally in ConvertToBifrostContext to add the routing engine log. + // set by prepare*Request functions (and inline realtime catalog lookups) when a + // provider was auto-resolved. Picked up centrally in ConvertToBifrostContext to + // add the routing engine log via EmitModelCatalogRoutingLog. FastHTTPUserValueModelCatalogResolution = "__bifrost_model_catalog_resolution" ) @@ -46,6 +47,26 @@ type ModelCatalogResolution struct { AllProviders []schemas.ModelProvider } +// EmitModelCatalogRoutingLog appends a RoutingEngineModelCatalog log entry and +// engines-used marker to bifrostCtx for an inline catalog resolution. Used by +// ConvertToBifrostContext (normal HTTP path) and by realtime handlers that +// bypass it (WebRTC, realtime client_secrets) so all paths emit observability +// in the same shape regardless of which routing layer did the lookup. +func EmitModelCatalogRoutingLog(bifrostCtx *schemas.BifrostContext, res *ModelCatalogResolution) { + if bifrostCtx == nil || res == nil { + return + } + providerStrs := make([]string, len(res.AllProviders)) + for i, p := range res.AllProviders { + providerStrs[i] = string(p) + } + bifrostCtx.AppendRoutingEngineLog(schemas.RoutingEngineModelCatalog, schemas.LogLevelInfo, fmt.Sprintf( + "No provider specified for model %s, found %d options in model catalog: [%s], selected: %s", + res.Model, len(res.AllProviders), strings.Join(providerStrs, ", "), res.ResolvedProvider, + )) + schemas.AppendToContextList(bifrostCtx, schemas.BifrostContextKeyRoutingEnginesUsed, schemas.RoutingEngineModelCatalog) +} + // ParseSessionIDFromBaggage extracts the session-id baggage member value. // It supports simple W3C baggage parsing sufficient for log grouping. func ParseSessionIDFromBaggage(header string) string { @@ -208,15 +229,7 @@ func ConvertToBifrostContext(ctx *fasthttp.RequestCtx, store HandlerStore) (*sch // it stores the resolution info on the fasthttp context. Emit the routing // engine log and mark the engine as used centrally here. if res, ok := ctx.UserValue(FastHTTPUserValueModelCatalogResolution).(*ModelCatalogResolution); ok && res != nil { - providerStrs := make([]string, len(res.AllProviders)) - for i, p := range res.AllProviders { - providerStrs[i] = string(p) - } - bifrostCtx.AppendRoutingEngineLog(schemas.RoutingEngineModelCatalog, schemas.LogLevelInfo, fmt.Sprintf( - "No provider specified for model %s, found %d options in model catalog: [%s], selecting first: %s", - res.Model, len(res.AllProviders), strings.Join(providerStrs, ", "), res.ResolvedProvider, - )) - schemas.AppendToContextList(bifrostCtx, schemas.BifrostContextKeyRoutingEnginesUsed, schemas.RoutingEngineModelCatalog) + EmitModelCatalogRoutingLog(bifrostCtx, res) } // Initialize tags map for collecting maxim tags diff --git a/transports/bifrost-http/server/plugins.go b/transports/bifrost-http/server/plugins.go index a868964e4e..f79d3a51c9 100644 --- a/transports/bifrost-http/server/plugins.go +++ b/transports/bifrost-http/server/plugins.go @@ -10,6 +10,7 @@ import ( "github.com/maximhq/bifrost/plugins/governance" "github.com/maximhq/bifrost/plugins/logging" "github.com/maximhq/bifrost/plugins/maxim" + "github.com/maximhq/bifrost/plugins/modelcatalogresolver" "github.com/maximhq/bifrost/plugins/otel" "github.com/maximhq/bifrost/plugins/prompts" "github.com/maximhq/bifrost/plugins/semanticcache" @@ -120,6 +121,9 @@ func loadBuiltinPlugin(ctx context.Context, name string, pluginConfig any, bifro } return compat.Init(*compatConfig, logger, bifrostConfig.ModelCatalog) + case modelcatalogresolver.PluginName: + return modelcatalogresolver.Init(bifrostConfig.ModelCatalog, logger) + default: return nil, fmt.Errorf("unknown built-in plugin: %s", name) } @@ -252,6 +256,17 @@ func (s *BifrostHTTPServer) loadBuiltinPlugins(ctx context.Context) error { } s.Config.SetPluginOrderInfo(maxim.PluginName, builtinPlacement, schemas.Ptr(8)) + // 9. ModelCatalogResolver (last routing layer — fills req.Provider from catalog only when + // no earlier routing plugin (governance routing rules, governance VK LB, enterprise LB) + // already set one. CEL rules can still match on provider == "" because this runs last. + // Requires a model catalog; only register when one is configured. + if s.Config.ModelCatalog != nil { + s.registerPluginWithStatus(ctx, modelcatalogresolver.PluginName, nil, nil, false) + } else { + s.markPluginDisabled(modelcatalogresolver.PluginName) + } + s.Config.SetPluginOrderInfo(modelcatalogresolver.PluginName, builtinPlacement, schemas.Ptr(9)) + return nil } diff --git a/transports/go.mod b/transports/go.mod index 1cd2004995..3cbfddeb22 100644 --- a/transports/go.mod +++ b/transports/go.mod @@ -18,6 +18,7 @@ require ( github.com/maximhq/bifrost/plugins/governance v1.5.18 github.com/maximhq/bifrost/plugins/logging v1.5.18 github.com/maximhq/bifrost/plugins/maxim v1.6.18 + github.com/maximhq/bifrost/plugins/modelcatalogresolver v0.0.0-20260531215024-856c9963e662 github.com/maximhq/bifrost/plugins/otel v1.2.18 github.com/maximhq/bifrost/plugins/prompts v1.0.18 github.com/maximhq/bifrost/plugins/semanticcache v1.5.18 diff --git a/transports/go.sum b/transports/go.sum index a8f0c0f333..cd1cbc8b82 100644 --- a/transports/go.sum +++ b/transports/go.sum @@ -273,6 +273,7 @@ github.com/maximhq/bifrost/plugins/maxim v1.6.18 h1:0EfHmwBLbmrG9hwofdU41x5x+SCw github.com/maximhq/bifrost/plugins/maxim v1.6.18/go.mod h1:L5gE+GCGWLiSi1UljV6ZV5v5sf0YhxV1j+EvHQNhqts= github.com/maximhq/bifrost/plugins/mocker v1.5.18 h1:4HCqMfTcxzjO2nicAxWzIKyyfOKwiGAJqT3lVX32oKI= github.com/maximhq/bifrost/plugins/mocker v1.5.18/go.mod h1:zr9x3vsPDYmdOPnbQlqC3a6TEFLnwcCa2hb1Z9zq5SE= +github.com/maximhq/bifrost/plugins/modelcatalogresolver v0.0.0-20260531215024-856c9963e662 h1:RMa9QlP7IIPjCJ6w+XTI+68HzV303XiPFBWVf1risPQ= github.com/maximhq/bifrost/plugins/otel v1.2.18 h1:dBrB0P9RCpJ71p+z+JBNSJAF/ZskqC7GkvFa+TlFnu0= github.com/maximhq/bifrost/plugins/otel v1.2.18/go.mod h1:G/wM8Ks+tv6QkRd0QXK/lwveKYTLNGTCsykLLVy3xTc= github.com/maximhq/bifrost/plugins/prompts v1.0.18 h1:BBtD2h4nQvZ2ewKurE3mu+I9l5VO94ZpfitYceaT4r0= From 88ac5700b8b11722f62e2b7d77fa0bfc4537c35f Mon Sep 17 00:00:00 2001 From: Pratham Mishra <99235987+Pratham-Mishra04@users.noreply.github.com> Date: Tue, 9 Jun 2026 15:50:44 +0530 Subject: [PATCH 005/108] docs: add `PreRequestHook` routing phase to plugin lifecycle and sequencing docs (#4178) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Documents the `PreRequestHook` interface introduced in v1.6.x — a new per-request routing phase that fires exactly once before any provider call, distinct from `PreLLMHook` which runs per provider attempt. This clarifies where routing decisions (provider, model, fallbacks) should be made and how they propagate through the fallback chain. ## Changes - Added `PreRequestHook` to the plugin lifecycle state diagram, sequence diagrams, and execution order descriptions across the architecture and getting-started docs, making clear it runs once per request while `PreLLMHook`/`PostLLMHook` run per attempt - Added a routing layer order table in `sequencing.mdx` documenting the built-in plugin execution order within `PreRequestHook`: governance (order 4) → enterprise load balancer → model-catalog-resolver (order 9, final fallback) - Added a full `PreRequestHook` reference section in `writing-go-plugin.mdx` with a comparison table against `PreLLMHook`, a routing example using `SetProvider`/`SetModel`, and notes on the two routing observability helpers (`AppendRoutingEngineLog`, `AppendToContextList`) - Updated `provider-routing.mdx` to reflect that all three routing layers (governance, enterprise LB Level 1, model-catalog-resolver) now execute inside the `PreRequestHook` phase rather than across separate middleware stages, and updated the flowcharts and execution order lists accordingly - Clarified that `model-catalog-resolver` now prefers the integration's canonical provider (OpenAI/Anthropic/GenAI/Bedrock/Cohere) when the request arrived via an integration route, rather than always selecting the first catalog candidate - Updated the log message example from `selecting first:` to `selected:` to match the new resolver behavior ## Type of change - [ ] Bug fix - [ ] Feature - [ ] Refactor - [x] Documentation - [ ] Chore/CI ## Affected areas - [ ] Core (Go) - [ ] Transports (HTTP) - [ ] Providers/Integrations - [x] Plugins - [ ] UI (React) - [x] Docs ## How to test Review the rendered docs for accuracy against the v1.6.x plugin interface. Verify that: - The `PreRequestHook` signature (`func PreRequestHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) error`) matches the SDK - The execution order table in `sequencing.mdx` matches the registered plugin orders in the codebase - The flowcharts in `provider-routing.mdx` correctly reflect that all Level 1 routing now happens in `PreRequestHook` ## Breaking changes - [ ] Yes - [x] No ## Related issues Documents the `PreRequestHook` routing phase introduced alongside the model-catalog-resolver and routing engine changes in v1.6.x. ## Security considerations None. Documentation-only change. ## Checklist - [x] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [x] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable ## Summary by CodeRabbit * **Documentation** * Updated plugin lifecycle documentation to clarify `PreRequestHook` execution (v1.6.x+) as a once-per-request routing phase separate from per-attempt hook phases. * Enhanced plugin sequencing and execution order documentation with clearer per-request vs. per-attempt semantics. * Expanded provider routing documentation with updated default resolution order and governance/load-balancing interaction details. * Added comprehensive `PreRequestHook` guidance for plugin developers with routing examples and helpers. --- docs/architecture/core/plugins.mdx | 32 +++++-- docs/plugins/getting-started.mdx | 16 ++-- docs/plugins/sequencing.mdx | 16 +++- docs/plugins/writing-go-plugin.mdx | 69 +++++++++++++++ docs/providers/provider-routing.mdx | 126 +++++++++++++++------------- 5 files changed, 188 insertions(+), 71 deletions(-) diff --git a/docs/architecture/core/plugins.mdx b/docs/architecture/core/plugins.mdx index 4a901272ee..479a4ad5ed 100644 --- a/docs/architecture/core/plugins.mdx +++ b/docs/architecture/core/plugins.mdx @@ -68,8 +68,12 @@ Every plugin goes through a well-defined lifecycle that ensures proper resource stateDiagram-v2 [*] --> PluginInit: Plugin Creation PluginInit --> Registered: Add to BifrostConfig - Registered --> PreHookCall: Request Received + Registered --> PreRequestHookCall: Request Received (once per request) + PreRequestHookCall --> RouteDecided: Provider/Model resolved + PreRequestHookCall --> RouteDecided: Return Error (logged, non-blocking) + + RouteDecided --> PreHookCall: Per-attempt phase PreHookCall --> ModifyRequest: Normal Flow PreHookCall --> ShortCircuitResponse: Return Response PreHookCall --> ShortCircuitError: Return Error @@ -87,7 +91,7 @@ stateDiagram-v2 FallbackCheck --> TryFallback: AllowFallbacks=true/nil FallbackCheck --> ResponseReady: AllowFallbacks=false - TryFallback --> PreHookCall: Next Provider + TryFallback --> PreHookCall: Next Provider (PreRequestHook NOT re-run) ModifyResponse --> ResponseReady: Modified RecoverError --> ResponseReady: Recovered @@ -148,6 +152,12 @@ sequenceDiagram participant Provider Client->>Bifrost: Request + Note over Bifrost,Plugin2: PreRequestHook phase (once per request, before any fan-out) + Bifrost->>Plugin1: PreRequestHook(request) + Plugin1-->>Bifrost: routed request + Bifrost->>Plugin2: PreRequestHook(request) + Plugin2-->>Bifrost: routed request + Note over Bifrost,Plugin2: PreLLMHook phase (per provider attempt) Bifrost->>Plugin1: PreLLMHook(request) Plugin1-->>Bifrost: modified request Bifrost->>Plugin2: PreLLMHook(request) @@ -163,9 +173,19 @@ sequenceDiagram **Execution Order:** -1. **PreHooks:** Execute in registration order (1 → 2 → N) -2. **Provider Call:** If no short-circuit occurred -3. **PostHooks:** Execute in reverse order (N → 2 → 1) +1. **PreRequestHooks** (per-request, registration order 1 → 2 → N): the **routing phase**. Plugins decide which provider/model the request goes to. Mutations to `req.Provider`/`req.Model`/`req.Fallbacks` commit to the shared request and are observed by every subsequent phase and every fallback attempt. There is no short-circuit. Plugin errors are non-blocking — logged as warnings and the pipeline continues to the next plugin. After all PreRequestHooks have run, the core validates `req.Provider`: an unresolved provider returns a 400 to the caller. +2. **PreLLMHooks** (per attempt, registration order 1 → 2 → N): pre-call transforms — caching, validation, content modification. May short-circuit with a synthetic response. +3. **Provider Call:** if no short-circuit occurred. +4. **PostLLMHooks** (per attempt, reverse order N → 2 → 1): response transforms — error recovery, logging, observability. + +**Per-request vs per-attempt:** `PreRequestHook` runs **exactly once** at the top of `handleRequest` / `handleStreamRequest`, before any provider call. `PreLLMHook` and `PostLLMHook` run **once per provider attempt** — so if the primary call fails and a fallback fires, `PreLLMHook` and `PostLLMHook` run again on the fallback, but `PreRequestHook` does **not**. This is what makes `PreRequestHook` the right place for routing decisions: the decision is committed once and applies uniformly to the primary attempt and every fallback. + + +**When to use which hook:** +- **PreRequestHook** → routing decisions (governance rules, load balancing, model-catalog provider resolution). Mutations to `req.Provider`/`req.Model`/`req.Fallbacks` stick. +- **PreLLMHook** → per-attempt transforms (semantic-cache lookups, request validation, content rewrites). Mutations to provider/model are intentionally no-ops here. +- **PostLLMHook** → per-attempt response handling (caching writes, logging, error recovery). + #### **Short-Circuit Response Flow (Cache Hit)** @@ -178,6 +198,7 @@ sequenceDiagram participant Provider Client->>Bifrost: Request + Note over Bifrost,Cache: PreRequestHook phase (routing decided) Bifrost->>Auth: PreLLMHook(request) Auth-->>Bifrost: modified request Bifrost->>Cache: PreLLMHook(request) @@ -203,6 +224,7 @@ sequenceDiagram participant Provider Client->>Bifrost: Stream Request + Note over Bifrost,Plugin2: PreRequestHook phase (routing decided) Bifrost->>Plugin1: PreLLMHook(request) Plugin1-->>Bifrost: modified request Bifrost->>Plugin2: PreLLMHook(request) diff --git a/docs/plugins/getting-started.mdx b/docs/plugins/getting-started.mdx index c84b5fba42..6eb025fdf9 100644 --- a/docs/plugins/getting-started.mdx +++ b/docs/plugins/getting-started.mdx @@ -54,8 +54,9 @@ This generates a `.so` file that exports specific functions matching Bifrost's p - `GetName() string` - Return the plugin name - `HTTPTransportPreHook()` - Intercept HTTP requests before they enter Bifrost core (HTTP transport only) - `HTTPTransportPostHook()` - Intercept HTTP responses after they exit Bifrost core (HTTP transport only) - - `PreLLMHook()` - Intercept requests before they reach providers - - `PostLLMHook()` - Process responses after provider calls + - `PreRequestHook()` v1.6.x+ - Once-per-request routing phase: decide provider/model/fallbacks + - `PreLLMHook()` - Intercept requests before they reach providers (runs per provider attempt) + - `PostLLMHook()` - Process responses after provider calls (runs per provider attempt) - `Cleanup() error` - Clean up resources on shutdown @@ -83,7 +84,7 @@ This means if you're running Bifrost on Linux AMD64, you must build your plugin 1. **Load** - Bifrost loads the `.so` file using Go's `plugin.Open()` 2. **Initialize** - Calls `Init()` with configuration from `config.json` -3. **Hook Execution** - Calls `PreLLMHook()` and `PostLLMHook()` for each request +3. **Hook Execution** - Calls `PreRequestHook()`, `PreLLMHook()` and `PostLLMHook()` for each request 4. **Cleanup** - Calls `Cleanup()` when Bifrost shuts down Plugins execute in a specific order: @@ -91,10 +92,11 @@ Plugins execute in a specific order: 1. `HTTPTransportPreHook` - Intercept HTTP requests (HTTP transport only) - 2. `PreLLMHook`/`PreMCPHook` - Executes in registration order, can short-circuit requests - 3. Provider call (if not short-circuited) - 4. `PostLLMHook`/`PostMCPHook` - Executes in reverse order of PreHooks - 5. `HTTPTransportPostHook` - Intercept HTTP responses (HTTP transport only, reverse order) + 2. `PreRequestHook` v1.6.x+ - **Once per request**, before any provider call. Routing decisions (provider/model/fallbacks) happen here and propagate to every attempt. + 3. `PreLLMHook`/`PreMCPHook` - Per provider attempt, registration order, can short-circuit requests + 4. Provider call (if not short-circuited) + 5. `PostLLMHook`/`PostMCPHook` - Per provider attempt, reverse order of PreHooks + 6. `HTTPTransportPostHook` - Intercept HTTP responses (HTTP transport only, reverse order) 1. `TransportInterceptor` - Modifies raw HTTP requests (HTTP transport only) diff --git a/docs/plugins/sequencing.mdx b/docs/plugins/sequencing.mdx index e261028b28..0355c5dc3f 100644 --- a/docs/plugins/sequencing.mdx +++ b/docs/plugins/sequencing.mdx @@ -32,6 +32,20 @@ graph LR Post-hooks execute in **reverse order** of pre-hooks (LIFO pattern). This means a `pre_builtin` plugin's `PreLLMHook` runs first, but its `PostLLMHook` runs last - ensuring proper cleanup and state unwinding. +### Routing layer order (PreRequestHook) + +`PreRequestHook` is the per-request **routing phase**. All routing-capable plugins fire here in registration order, and each one sees the routing decisions of those that ran before it. Built-in routing plugins are sequenced as follows within the `builtin` group: + +| Order | Plugin | Role | +|-------|--------|------| +| 4 | governance | Routing rules (CEL) + VK-scoped weighted load balancing | +| Higher | adaptive-loadbalancer (Enterprise) | Performance-based provider selection across the model catalog | +| 9 (last) | model-catalog-resolver | **Final fallback** — fills in `req.Provider` from the model catalog for unprefixed models when no earlier plugin picked one | + +The resolver runs last so CEL routing rules can match on `provider == ""` (the unresolved state) and earlier plugins always get the canonical bare model. After `PreRequestHook` returns, the core validates that `req.Provider` is non-empty — an unresolvable request returns a 400 with a clear error. + +Custom routing plugins can slot into this chain via `placement` + `order` like any other plugin. Place them in `pre_builtin` to override governance, or `post_builtin` to act as a custom fallback after all built-ins. + ### Ordering within a group Within each placement group, plugins are sorted by their `order` value (lower executes earlier). Plugins with the same order preserve their registration order. @@ -185,6 +199,6 @@ When in doubt, use the default `post_builtin` placement. Most custom plugins - l ## Next steps -- **[Writing a Go plugin](./writing-go-plugin)** - Build your first custom plugin with `PreLLMHook` and `PostLLMHook` +- **[Writing a Go plugin](./writing-go-plugin)** - Build your first custom plugin with `PreRequestHook`, `PreLLMHook`, and `PostLLMHook` - **[Writing a WASM plugin](./writing-wasm-plugin)** - Build a portable WASM plugin - **[Plugin architecture](../architecture/core/plugins)** - Deep dive into the plugin lifecycle and hook execution model diff --git a/docs/plugins/writing-go-plugin.mdx b/docs/plugins/writing-go-plugin.mdx index e9b2a514a7..d2462aa85f 100644 --- a/docs/plugins/writing-go-plugin.mdx +++ b/docs/plugins/writing-go-plugin.mdx @@ -152,6 +152,16 @@ func HTTPTransportStreamChunkHook(ctx *schemas.BifrostContext, req *schemas.HTTP } +// PreRequestHook is called once per top-level request (NOT per fallback attempt). +// This is the routing phase — use it for provider/model/fallback decisions. +// Mutations to req.Provider/req.Model/req.Fallbacks commit and propagate to every attempt. +// Errors are non-blocking (logged + skipped). +func PreRequestHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) error { + ctx.Log(schemas.LogLevelInfo, "PreRequestHook called") + // Plugins that don't participate in routing should just return nil + return nil +} + // PreLLMHook is called before the request is sent to the provider // This is where you can modify requests or short-circuit the flow func PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { @@ -268,6 +278,13 @@ func HTTPTransportStreamChunkHook(ctx *schemas.BifrostContext, req *schemas.HTTP } +// PreRequestHook is called once per top-level request (routing phase) +// Mutations to req.Provider/req.Model/req.Fallbacks commit across fallbacks +func PreRequestHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) error { + fmt.Println("PreRequestHook called") + return nil +} + // PreLLMHook is called before the request is sent to the provider // This is where you can modify requests or short-circuit the flow func PreLLMHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) (*schemas.BifrostRequest, *schemas.LLMPluginShortCircuit, error) { @@ -547,6 +564,58 @@ This function is **only called** when using `bifrost-http`. It's **not invoked** +#### `PreRequestHook(...)` v1.6.x+ + +Called **once per top-level request**, before any provider call and before `PreLLMHook`. This is the **routing phase**: it's where plugins decide which provider, model, and fallbacks the request should be sent to. + +Use this for: +- **Routing decisions**: governance rules, virtual-key load balancing, geo/tier routing +- **Provider resolution**: filling in `req.Provider` for unprefixed model names (the built-in `model-catalog-resolver` does this as the last routing layer) +- **Fallback chain construction**: populating `req.Fallbacks` based on policy + +**Why a separate hook from `PreLLMHook`:** + +| Aspect | `PreRequestHook` | `PreLLMHook` | +|---|---|---| +| Runs | Once per request | Once per provider attempt (re-runs on each fallback) | +| Provider/Model mutations | **Commit and propagate** to every attempt | No-op for `Provider`/`Model` (overwritten by core) | +| Use for | Routing decisions | Request transforms, caching, validation | +| Short-circuit | No | Yes (can return a synthetic response) | +| Error semantics | Non-blocking (logged, pipeline continues) | Non-blocking (logged, pipeline continues) | + +**Routing Example:** + +```go +func PreRequestHook(ctx *schemas.BifrostContext, req *schemas.BifrostRequest) error { + provider, model, _ := req.GetRequestFields() + + // Route premium-tier requests to a faster provider + if tier := ctx.Value(schemas.BifrostContextKey("x-tier")); tier == "premium" && provider == "openai" { + req.SetProvider(schemas.Anthropic) + req.SetModel("claude-3-5-sonnet") + // Emit a routing-engine log entry so users can see why this decision was made + ctx.AppendRoutingEngineLog(schemas.RoutingEngineRoutingRule, schemas.LogLevelInfo, + fmt.Sprintf("Routed %s to anthropic/claude-3-5-sonnet (tier=premium)", model)) + schemas.AppendToContextList(ctx, schemas.BifrostContextKeyRoutingEnginesUsed, + schemas.RoutingEngineRoutingRule) + } + return nil +} +``` + +**Two helpers worth knowing when writing routing logic:** + +- `ctx.AppendRoutingEngineLog(engine, level, message)` — emits a structured log entry visible in observability tools. Use it to explain *why* a routing decision was made. +- `schemas.AppendToContextList(ctx, schemas.BifrostContextKeyRoutingEnginesUsed, engineName)` — records which routing engine(s) participated. Surfaces in telemetry as `routing.engines_used`. + + +**Plugin order matters for routing.** Built-in routing plugins run in this order within `PreRequestHook`: governance routing rules → governance VK load balancing → enterprise load balancer → model-catalog-resolver (final fallback). Custom plugins can slot in via `placement` + `order` — see [Plugin Sequencing](./sequencing). + + + +**Errors are non-blocking.** Returning a non-nil error from `PreRequestHook` logs a warning but does NOT fail the request — the pipeline continues with the next plugin. The core validates `req.Provider` after all `PreRequestHook` plugins have run; an unresolved provider returns a 400 to the caller. + + #### `PreLLMHook(...)` Called before each provider request. Use this to: diff --git a/docs/providers/provider-routing.mdx b/docs/providers/provider-routing.mdx index 3f74aed602..1376cbd9c3 100644 --- a/docs/providers/provider-routing.mdx +++ b/docs/providers/provider-routing.mdx @@ -836,20 +836,21 @@ This is how Bifrost achieves **intelligent cross-provider routing** without manu v1.5.0-prerelease7 and above**. -When a request includes a bare model name without a `provider/` prefix (e.g., `"model": "gpt-4o"` instead of `"model": "openai/gpt-4o"`), Bifrost automatically resolves the provider using the Model Catalog. Note that this default behavior is applied **after all other routing engines** have run. +When a request includes a bare model name without a `provider/` prefix (e.g., `"model": "gpt-4o"` instead of `"model": "openai/gpt-4o"`), Bifrost automatically resolves the provider using the Model Catalog. This default behavior is applied **after all other routing engines** have run — the built-in `model-catalog-resolver` PreRequestHook plugin is registered as the last routing layer (order 9 within `builtin`), so governance routing rules, VK load balancing, and enterprise LB all get first crack. ### How It Works 1. **Request arrives** without a provider prefix (e.g., `"model": "gpt-4o"`) -2. **Catalog lookup**: Bifrost calls `GetProvidersForModel("gpt-4o")` to find all providers that support the model -3. **Provider selected**: A provider from the catalog's available list is used (e.g., `openai`) -4. **Request continues**: The resolved `provider/model` string is used for load balancing and fallback handling +2. Governance, VK LB, and enterprise LB all run first; if any of them sets `req.Provider`, the resolver no-ops +3. **Catalog lookup** (if `req.Provider` is still empty): Bifrost calls `GetProvidersForModel("gpt-4o")` to find all providers that support the model +4. **Provider selected**: If the request came in via an integration route (OpenAI / Anthropic / GenAI / Bedrock / Cohere) and the catalog includes that integration's canonical provider in the candidate list, it is preferred. Otherwise the first candidate is selected. +5. **Request continues**: The resolved `provider/model` is used for the provider call, fallback handling, and Level 2 key selection. This is logged as the **`model-catalog`** routing engine in telemetry and routing logs, with a message like: ``` No provider specified for model gpt-4o, found 3 options in model catalog: -[openai, azure, groq], selecting first: openai +[openai, azure, groq], selected: openai ``` ### Example @@ -1189,60 +1190,67 @@ This means key-level optimization works regardless of how the provider was chose flowchart TD Start["Request: gpt-4o"] - subgraph Governance["Governance Plugin (HTTPTransportIntercept)"] + subgraph PreReq["PreRequestHook Phase (once per request, registration order)"] HasVK{"Has VK with
provider_configs?"} - GovRoute["Provider Selection:
Weighted random"] - AddPrefix["Add prefix:
azure/gpt-4o"] + GovRoute["Governance:
Routing rules + VK weighted random"] + AddPrefix["Set req.Provider/Model:
azure/gpt-4o"] + PrefixCheck{"req.Provider
already set?"} + LBProvider["Enterprise LB:
Performance-based selection"] + AddLBPrefix["Set req.Provider/Model:
openai/gpt-4o"] + Resolver["model-catalog-resolver:
Fill from catalog (last fallback)"] end - subgraph LB1["Load Balancer Level 1 (Middleware)"] - PrefixCheck{"Has provider
prefix?"} - LBProvider["Provider Selection:
Performance-based"] - AddLBPrefix["Add prefix:
openai/gpt-4o"] - end - - subgraph LB2["Load Balancer Level 2 (Key Selector)"] + subgraph LB2["Load Balancer Level 2 (Key Selector, in core)"] GetKeys["Get available keys
for selected provider"] ScoreKeys["Score keys by
performance metrics"] SelectKey["Select best key"] end Start --> HasVK - HasVK -->|Yes| GovRoute --> AddPrefix + HasVK -->|Yes| GovRoute --> AddPrefix --> PrefixCheck HasVK -->|No| PrefixCheck - AddPrefix --> PrefixCheck - PrefixCheck -->|Yes, skip Level 1| GetKeys - PrefixCheck -->|No| LBProvider --> AddLBPrefix --> GetKeys + PrefixCheck -->|Yes, skip LB Level 1| Resolver + PrefixCheck -->|No| LBProvider --> AddLBPrefix --> Resolver + Resolver --> GetKeys GetKeys --> ScoreKeys --> SelectKey --> Execute["Execute request
with selected provider + key"] ``` ### Execution Order -1. **HTTPTransportIntercept** (Governance Plugin - Provider Level) - - Runs first in the request pipeline - - Checks if Virtual Key has `provider_configs` - - If yes: adds provider prefix (e.g., `azure/gpt-4o`) - - **Result**: Provider is selected by governance rules - -2. **Middleware** (Load Balancing Plugin - Provider Level / Direction) - - Runs after HTTPTransportIntercept - - Checks if model string contains "/" - - If yes: **skips provider selection** (already determined by governance or user) - - If no: performs performance-based provider selection - - **Result**: Provider prefix added if not already present - -3. **KeySelector** (Load Balancing - Key Level / Route) - - **Always runs** during request execution in Bifrost core - - Gets all keys for the selected provider - - Filters keys based on model restrictions +All three routing layers (governance, enterprise LB Level 1, model-catalog-resolver) now run inside a single **PreRequestHook** phase that fires **once per top-level request**, before any provider call and before per-attempt hooks. Within that phase, plugins execute in placement + order: + +1. **Governance Plugin** (PreRequestHook, builtin order 4) + - Evaluates routing rules (CEL expressions, scope hierarchy) + - If Virtual Key has `provider_configs`: performs weighted random provider selection + - **Result**: `req.Provider`/`req.Model` set; `req.Fallbacks` populated + +2. **Enterprise Load Balancer Level 1** (PreRequestHook, builtin) + - Runs after governance + - If `req.Provider` is already set (by governance or by an explicit `provider/model` prefix from the user): **skips provider selection** + - If not: performs performance-based provider selection across catalog providers + - **Result**: `req.Provider`/`req.Model` set if previously empty + +3. **model-catalog-resolver** (PreRequestHook, builtin order 9 — final fallback) + - Runs last + - If `req.Provider` is still empty: looks up the model in the catalog and picks a provider (preferring the integration's canonical provider when the request came in via an integration route) + - Emits a `model-catalog` routing-engine log entry + - **Result**: Always leaves `req.Provider` populated when the catalog knows about the model + +4. **Empty-provider validation** (core, after PreRequestHook) + - If `req.Provider` is still empty: returns 400 to the caller with a clear error + +5. **Load Balancer Level 2** (Key Selector — core, per provider attempt) + - **Always runs** during request execution + - Gets all keys for the selected provider, filters by model restrictions - Scores each key by performance metrics - Selects best key using weighted random + exploration - **Result**: Optimal key selected within the provider - **Important**: Even when governance specifies `azure/gpt-4o`, load balancing - **still optimizes which Azure key to use** based on performance metrics. This - is the power of the two-level architecture! + **Important**: Even when governance specifies `azure/gpt-4o` in PreRequestHook, + load balancing Level 2 **still optimizes which Azure key to use** based on + performance metrics. The two-level architecture is preserved — only the + *layer* where Level 1 runs has moved from a middleware to PreRequestHook. ### Example Scenarios @@ -1395,35 +1403,37 @@ Routing Rules provide sophisticated, expression-based control over request routi flowchart TD Start["Request: model + provider"] - subgraph Rules["1. Routing Rules Layer (Evaluated First)"] - RuleMatch{"CEL Expression
Matches?"} - RuleDecision["Override:
New provider/model/fallbacks"] - NoMatch["No match:
Continue to Governance"] + subgraph PreReq["PreRequestHook Phase (once per request)"] + direction TB + subgraph Gov["Governance Plugin"] + RuleMatch{"CEL Routing Rule
Matches?"} + RuleDecision["Override:
provider/model/fallbacks"] + VKValidation["Virtual Key Validation"] + GovRouting["VK Provider Selection
(weighted random)"] + end + LB1["Enterprise LB Level 1:
Provider Selection
(skipped if provider already set)"] + Resolver["model-catalog-resolver:
Fill provider from catalog
(final fallback)"] end - subgraph Gov["2. Governance Layer (if no routing rule matched)"] - VKValidation["Virtual Key Validation"] - GovRouting["Provider Governance Routing
(weighted random)"] - end - - subgraph LB["3. Load Balancing Layer"] - LB1["Level 1: Provider Selection"] - LB2["Level 2: Key Selection"] - end + LB2["LB Level 2: Key Selection
(core, per attempt)"] Start --> RuleMatch RuleMatch -->|Yes| RuleDecision --> LB1 - RuleMatch -->|No| NoMatch --> VKValidation --> GovRouting --> LB1 - LB1 --> LB2 --> Execute["Execute with
selected provider + key"] + RuleMatch -->|No| VKValidation --> GovRouting --> LB1 + LB1 --> Resolver --> LB2 --> Execute["Execute with
selected provider + key"] ``` ### How It Works +All routing layers below execute inside the **PreRequestHook** phase in registration order; routing rules run first within the governance plugin's hook body, before VK load balancing: + 1. **Routing rules evaluate first** in scope precedence order (VirtualKey → Team → Customer → Global) -2. **If a routing rule matches**: provider/model/fallbacks are overridden, governance provider_configs are skipped -3. **If no routing rule matches**: governance provider selection runs (weighted random) -4. **Load balancing Level 1**: skipped if provider already determined (has "/" prefix) -5. **Load balancing Level 2** (key selection): always runs to select the best key within the determined provider +2. **If a routing rule matches**: provider/model/fallbacks are overridden, the VK `provider_configs` weighted selection is skipped +3. **If no routing rule matches**: VK provider selection runs (weighted random) +4. **Enterprise LB Level 1**: skipped if `req.Provider` is already set; otherwise performs performance-based selection +5. **model-catalog-resolver**: last fallback — fills `req.Provider` from the catalog if no earlier plugin set it +6. **Empty-provider validation** (core): returns 400 if `req.Provider` is still empty after the phase +7. **Load balancing Level 2** (key selection, core, per attempt): always runs to select the best key within the determined provider ### Available CEL Variables From 1799a3ded532e3fbcdbdec9c789d3175888a3b05 Mon Sep 17 00:00:00 2001 From: Pratham Mishra <99235987+Pratham-Mishra04@users.noreply.github.com> Date: Tue, 9 Jun 2026 15:52:45 +0530 Subject: [PATCH 006/108] feat: add routing allowlist enforcement via `BifrostContextKeyRoutingAllowedProviders` (#4179) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Introduces a two-level routing allowlist enforcement mechanism that prevents any plugin or user-specified provider prefix from bypassing provider restrictions defined by a Virtual Key's `provider_configs`. Previously, governance could fail to pick a provider but a downstream routing layer could still select a non-permitted provider. This change closes that gap by publishing the VK's allowed-provider set to the request context and enforcing it both cooperatively (in the model catalog resolver) and as a hard guarantee (in core, after all `PreRequestHook` plugins have run). ## Changes - Adds `BifrostContextKeyRoutingAllowedProviders` context key (`[]ModelProvider`) that plugins can set to constrain which providers are valid for a request. An empty slice means "no provider is permitted" (fail-closed → HTTP 400). - Adds `enforceRoutingAllowlist` and `filterFallbacksByAllowlist` helpers in `core/bifrost.go`. After all pre-request hooks complete, `handleRequest` and `handleStreamRequest` validate the resolved provider against the allowlist and prune fallbacks to only allowed providers. A non-allowed primary provider returns HTTP 400. - The governance plugin now publishes the VK's `provider_configs` providers to `BifrostContextKeyRoutingAllowedProviders` during `PreRequestHook`, covering the case where governance cannot pick a provider itself but still needs to constrain downstream layers. - The model catalog resolver intersects its catalog candidates with the allowlist (when set) before selecting a provider, emitting routing-engine observability logs that explain which candidates were excluded and why. Returns `("", nil)` when the allowlist excludes all candidates. - Documents the two-level enforcement model, context key semantics, and custom plugin usage in `docs/providers/provider-routing.mdx`. ## Type of change - [ ] Bug fix - [x] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [ ] Transports (HTTP) - [ ] Providers/Integrations - [x] Plugins - [ ] UI (React) - [x] Docs ## How to test ```sh go test ./... ``` - Configure a Virtual Key with `provider_configs` restricted to a specific provider (e.g., `openai`). - Send a request with an explicit `provider/model` prefix targeting a non-allowed provider (e.g., `anthropic/claude-3`). Expect HTTP 400 with a message indicating the provider is not permitted. - Send a request with a model resolvable by the catalog to multiple providers where some are excluded by the VK allowlist. Verify routing-engine logs show the excluded candidates and the request routes only to an allowed provider. - Set `BifrostContextKeyRoutingAllowedProviders` to an empty slice from a custom plugin and confirm the request fails closed with HTTP 400. - Verify fallbacks targeting non-allowed providers are silently pruned and do not appear in the fallback chain. ## Breaking changes - [x] Yes - [ ] No Requests that previously succeeded by specifying an explicit `provider/model` prefix that bypassed a VK's `provider_configs` restrictions will now be rejected with HTTP 400. Any fallbacks targeting providers outside the VK's allowed set will be silently removed from the fallback chain. ## Security considerations This change strengthens provider-level access control enforced by Virtual Keys. Without this, a user or plugin could bypass governance-imposed provider restrictions by specifying an explicit provider prefix or by relying on a downstream routing layer to select a non-permitted provider. The hard enforcement in core ensures the allowlist is a guarantee rather than a best-effort constraint. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [ ] I added/updated tests where appropriate - [x] I updated documentation where needed - [ ] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable ## Summary by CodeRabbit ## New Features * Added provider routing allowlist enforcement: configured allowed providers now restrict request routing across primary and fallback options; requests using non-permitted primary providers return HTTP 400, and fallback providers are automatically filtered to only permitted options. ## Documentation * Added routing allowlist enforcement documentation covering enforcement mechanisms and fail-closed behavior. --- core/bifrost.go | 61 +++++++++++++++++++++++++ core/schemas/bifrost.go | 1 + docs/providers/provider-routing.mdx | 12 +++++ plugins/governance/main.go | 14 ++++++ plugins/modelcatalogresolver/main.go | 66 +++++++++++++++++++++++++++- 5 files changed, 152 insertions(+), 2 deletions(-) diff --git a/core/bifrost.go b/core/bifrost.go index f63bf3ebec..eaf296955b 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -4625,6 +4625,55 @@ func (bifrost *Bifrost) shouldContinueWithFallbacks(fallback schemas.Fallback, f return true } +// filterFallbacksByAllowlist returns a new slice containing only the fallbacks whose Provider +// is present in allowed. Used by handleRequest/handleStreamRequest to enforce the routing +// allowlist published via BifrostContextKeyRoutingAllowedProviders — a plugin earlier in +// PreRequestHook may have populated fallbacks that a later plugin's allowlist excludes. +func filterFallbacksByAllowlist(fallbacks []schemas.Fallback, allowed []schemas.ModelProvider) []schemas.Fallback { + if len(fallbacks) == 0 { + return fallbacks + } + filtered := make([]schemas.Fallback, 0, len(fallbacks)) + for _, fb := range fallbacks { + if slices.Contains(allowed, fb.Provider) { + filtered = append(filtered, fb) + } + } + return filtered +} + +// enforceRoutingAllowlist gates the resolved provider against the routing +// allowlist published via BifrostContextKeyRoutingAllowedProviders (e.g. by +// the governance plugin) and prunes the fallback list to allowed providers. +// Returns the (possibly filtered) fallback list, or a prepared BifrostError +// when the resolved provider isn't on the allowlist. When no allowlist is +// set on ctx, fallbacks pass through unchanged with a nil error. +// +// Side effect: when an allowlist is in effect, the filtered fallbacks are +// written back to req via SetFallbacks so downstream phases observe the +// pruned list. +func enforceRoutingAllowlist( + ctx *schemas.BifrostContext, + req *schemas.BifrostRequest, + provider schemas.ModelProvider, + model string, + fallbacks []schemas.Fallback, +) ([]schemas.Fallback, *schemas.BifrostError) { + allowed, ok := ctx.Value(schemas.BifrostContextKeyRoutingAllowedProviders).([]schemas.ModelProvider) + if !ok { + return fallbacks, nil + } + if !slices.Contains(allowed, provider) { + bifrostErr := newBifrostErrorFromMsg(fmt.Sprintf("provider %q is not permitted for this request (routing allowlist: %v)", provider, allowed)) + bifrostErr.PopulateExtraFields(req.RequestType, provider, model, model) + bifrostErr.StatusCode = schemas.Ptr(fasthttp.StatusBadRequest) + return nil, bifrostErr + } + filtered := filterFallbacksByAllowlist(fallbacks, allowed) + req.SetFallbacks(filtered) + return filtered, nil +} + // handleRequest handles the request to the provider based on the request type // It handles plugin hooks, request validation, response processing, and fallback providers. // If the primary provider fails, it will try each fallback provider in order until one succeeds. @@ -4663,6 +4712,13 @@ func (bifrost *Bifrost) handleRequest(ctx *schemas.BifrostContext, req *schemas. err.PopulateExtraFields(req.RequestType, provider, model, model) return nil, err } + // Enforce the routing-allowlist if any plugin published one. This guarantees no + // downstream layer (or user-specified provider prefix) can bypass governance VK + // restrictions or any other plugin-imposed allowlist. + var allowlistErr *schemas.BifrostError + if fallbacks, allowlistErr = enforceRoutingAllowlist(ctx, req, provider, model, fallbacks); allowlistErr != nil { + return nil, allowlistErr + } bifrost.logger.Debug(fmt.Sprintf("primary provider %s with model %s and %d fallbacks", provider, model, len(fallbacks))) @@ -4768,6 +4824,11 @@ func (bifrost *Bifrost) handleStreamRequest(ctx *schemas.BifrostContext, req *sc err.PopulateExtraFields(req.RequestType, provider, model, model) return nil, err } + // Enforce the routing-allowlist if any plugin published one. See handleRequest. + var allowlistErr *schemas.BifrostError + if fallbacks, allowlistErr = enforceRoutingAllowlist(ctx, req, provider, model, fallbacks); allowlistErr != nil { + return nil, allowlistErr + } bifrost.logger.Debug(fmt.Sprintf("primary provider %s with model %s and %d fallbacks", provider, model, len(fallbacks))) diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index ce42cb2990..201a600475 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -289,6 +289,7 @@ const ( BifrostIsAsyncRequest BifrostContextKey = "bifrost-is-async-request" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) - whether the request is an async request (only used in gateway) BifrostContextKeyRequestHeaders BifrostContextKey = "bifrost-request-headers" // map[string]string (all request headers with lowercased keys) BifrostContextKeyRequestQuery BifrostContextKey = "bifrost-request-query" // map[string]string (request query params with lowercased keys; consumed by governance routing CEL rules) + BifrostContextKeyRoutingAllowedProviders BifrostContextKey = "bifrost-routing-allowed-providers" // []ModelProvider; when set, downstream routing layers (enterprise LB, model-catalog-resolver) must intersect their candidate providers with this set. Plugins set this when they have an opinion about which providers are valid for the request — even if they couldn't pick one themselves. Empty slice means "no provider is permitted" (fail-closed). BifrostContextKeyAllowPerRequestStorageOverride BifrostContextKey = "bifrost-allow-per-request-storage-override" // bool (set by transport from config — gates whether x-bf-disable-content-logging and x-bf-store-raw-request-response per-request overrides are honored) BifrostContextKeyAllowPerRequestRawOverride BifrostContextKey = "bifrost-allow-per-request-raw-override" // bool (set by transport from config — gates whether x-bf-send-back-raw-request and x-bf-send-back-raw-response per-request overrides are honored) BifrostContextKeyDisableContentLogging BifrostContextKey = "x-bf-disable-content-logging" // bool (per-request override for content logging; only honored when BifrostContextKeyAllowPerRequestStorageOverride is true) diff --git a/docs/providers/provider-routing.mdx b/docs/providers/provider-routing.mdx index 1376cbd9c3..74108339db 100644 --- a/docs/providers/provider-routing.mdx +++ b/docs/providers/provider-routing.mdx @@ -874,6 +874,18 @@ curl -X POST http://localhost:8080/v1/chat/completions \ prefix. +### Routing allowlist enforcement + +When a Virtual Key has `provider_configs`, governance publishes the VK's allowed-provider set to the request context (`BifrostContextKeyRoutingAllowedProviders`). The constraint is then enforced at **two levels**: + +1. **Cooperative filtering (observability-first):** Enterprise LB and the model-catalog-resolver intersect their catalog candidates with the allowlist before picking a provider. This produces clean routing-engine logs explaining *why* a candidate was excluded ("filtered N catalog candidates by routing allowlist"). + +2. **Hard enforcement in core:** After all `PreRequestHook` plugins have run, the core validates the final `req.Provider` against the allowlist. If `req.Provider` isn't in the allowlist, the request fails with HTTP 400. Fallbacks that target non-allowed providers are silently filtered out. + +**Why two levels:** cooperative filtering surfaces routing decisions in observability; core enforcement makes the constraint a *guarantee* that no plugin (or user-specified `provider/model` prefix) can bypass. A user request for `model: "anthropic/claude-3"` against a VK that allows only `[openai, azure]` is rejected by core enforcement even though the user provided an explicit prefix. + +Custom routing plugins can set the same context key to constrain downstream routing for any reason — geo restrictions, A/B test cohorts, tier-based gating, etc. The semantics are **fail-closed**: setting `BifrostContextKeyRoutingAllowedProviders` to an empty slice means "no provider is permitted for this request" → HTTP 400. + --- ## Governance-based Routing diff --git a/plugins/governance/main.go b/plugins/governance/main.go index 53953ae32b..27ceb35703 100644 --- a/plugins/governance/main.go +++ b/plugins/governance/main.go @@ -1010,6 +1010,20 @@ func (p *GovernancePlugin) PreRequestHook(ctx *schemas.BifrostContext, req *sche stampGovernanceCtxFromVK(ctx, virtualKey) + // Publish the VK's allowed-provider set so downstream routing layers (enterprise LB, + // model-catalog-resolver) intersect their candidates with it. This guards against the case + // where governance fails to pick a provider (every VK entry rejected by allowed_models / + // budget / rate limit) and a downstream layer would otherwise pick a provider the VK does + // not permit. Empty slice means "no provider is permitted" → fail-closed via the empty- + // provider validation in handleRequest. + if virtualKey != nil { + allowed := make([]schemas.ModelProvider, 0, len(virtualKey.ProviderConfigs)) + for _, pc := range virtualKey.ProviderConfigs { + allowed = append(allowed, schemas.ModelProvider(pc.Provider)) + } + ctx.SetValue(schemas.BifrostContextKeyRoutingAllowedProviders, allowed) + } + // Large-payload mode: the body streams to the provider unparsed, so req.Model is // empty for routes where the model lives in the body (OpenAI/Anthropic chat, // responses, etc.). Route on LargePayloadMetadata.Model — the provider's diff --git a/plugins/modelcatalogresolver/main.go b/plugins/modelcatalogresolver/main.go index 104c6e9a77..df251edb30 100644 --- a/plugins/modelcatalogresolver/main.go +++ b/plugins/modelcatalogresolver/main.go @@ -120,7 +120,8 @@ func (p *Plugin) PreRequestHook(ctx *schemas.BifrostContext, req *schemas.Bifros // ResolveProviderFromCatalog performs the deterministic, integration-aware provider pick // that PreRequestHook does, exposed for transport paths that can't run through // PreRequestHook (realtime client_secrets, WebRTC). Returns the selected provider plus the -// full ordered candidate list. Returns ("", nil) when the catalog has no match for the model. +// candidate list (post-allowlist when an allowlist is in effect). Returns ("", nil) when +// the catalog has no match for the model, or when an allowlist excludes every candidate. // // The integration hint (BifrostContextKeyIntegrationType, when present and mapped) biases // the pick toward the integration's canonical provider if it is in the candidate set; @@ -130,6 +131,12 @@ func (p *Plugin) PreRequestHook(ctx *schemas.BifrostContext, req *schemas.Bifros // OpenAI SDK (BifrostContextKeyIsAzureUserAgent), schemas.Azure is preferred over // schemas.OpenAI when Azure is in the candidate list — the openai-format converters no // longer apply this default inline. +// +// When BifrostContextKeyRoutingAllowedProviders is set on ctx by an earlier plugin (e.g., +// governance VK config), the candidate list is intersected with the allowlist before +// selection — emitting routing-engine logs visible to callers when the allowlist prunes +// candidates. Side effect: routing-engine logs are written to ctx when allowlist filtering +// is applied (nil ctx skips logging). func ResolveProviderFromCatalog(ctx *schemas.BifrostContext, catalog *modelcatalog.ModelCatalog, model string) (schemas.ModelProvider, []schemas.ModelProvider) { if catalog == nil || model == "" { return "", nil @@ -146,13 +153,68 @@ func ResolveProviderFromCatalog(ctx *schemas.BifrostContext, catalog *modelcatal return strings.Compare(string(a), string(b)) }) - selected := providers[0] var integrationType string var isAzureUser bool + var allowed []schemas.ModelProvider + allowlistSet := false if ctx != nil { integrationType, _ = ctx.Value(schemas.BifrostContextKeyIntegrationType).(string) isAzureUser, _ = ctx.Value(schemas.BifrostContextKeyIsAzureUserAgent).(bool) + allowed, allowlistSet = ctx.Value(schemas.BifrostContextKeyRoutingAllowedProviders).([]schemas.ModelProvider) + } + + // Respect the routing-allowlist set by an earlier plugin (e.g., governance VK config): + // intersect catalog candidates with the allowlist so the VK's provider restrictions hold + // even when no earlier routing plugin set req.Provider. Emit observability logs for both + // the partial-prune and all-pruned cases — the two-level enforcement (cooperative here + + // hard core enforcement) is only useful if the cooperative pruning is visible in routing + // engine logs when it fires. + if allowlistSet { + preFilterCount := len(providers) + preFilterStrs := make([]string, preFilterCount) + for i, prov := range providers { + preFilterStrs[i] = string(prov) + } + allowedStrs := make([]string, len(allowed)) + for i, prov := range allowed { + allowedStrs[i] = string(prov) + } + filtered := make([]schemas.ModelProvider, 0, preFilterCount) + excluded := make([]schemas.ModelProvider, 0) + for _, prov := range providers { + if slices.Contains(allowed, prov) { + filtered = append(filtered, prov) + } else { + excluded = append(excluded, prov) + } + } + if len(excluded) > 0 && ctx != nil { + filteredStrs := make([]string, len(filtered)) + for i, prov := range filtered { + filteredStrs[i] = string(prov) + } + ctx.AppendRoutingEngineLog(schemas.RoutingEngineModelCatalog, schemas.LogLevelInfo, fmt.Sprintf( + "Catalog returned %d candidate provider(s) for model %s: [%s]; provider allowlist is [%s], so excluded %d; remaining providers are [%s]", + preFilterCount, model, strings.Join(preFilterStrs, ", "), + strings.Join(allowedStrs, ", "), + len(excluded), + strings.Join(filteredStrs, ", "), + )) + } + providers = filtered + if len(providers) == 0 { + if ctx != nil { + ctx.AppendRoutingEngineLog(schemas.RoutingEngineModelCatalog, schemas.LogLevelInfo, fmt.Sprintf( + "Catalog returned %d candidate provider(s) for model %s: [%s]; provider allowlist [%s] excluded all of them; leaving req.Provider empty", + preFilterCount, model, strings.Join(preFilterStrs, ", "), + strings.Join(allowedStrs, ", "), + )) + } + return "", nil + } } + + selected := providers[0] if integrationType != "" { if integrationDefault, mapped := integrationTypeToDefaultProvider[integrationType]; mapped && integrationDefault != "" { preferred := integrationDefault From 4bddf4096cae84f243c7b406da5b0cd91dbd350f Mon Sep 17 00:00:00 2001 From: Pratham Mishra <99235987+Pratham-Mishra04@users.noreply.github.com> Date: Tue, 9 Jun 2026 15:54:44 +0530 Subject: [PATCH 007/108] refactor: promote `KeyAliases` value type from `string` to `AliasConfig` with legacy wire-shape compatibility (#4180) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary `KeyAliases` previously mapped user-facing model names to plain strings (`map[string]string`). This PR promotes the value type to a rich `AliasConfig` struct, enabling per-alias metadata (`ModelName`, `ModelFamily`) and provider-specific overrides (`AzureAliasCfg`, `VertexAliasCfg`, `BedrockAliasCfg`, `ReplicateAliasCfg`) to be expressed directly on an alias entry rather than inferred from the wire model ID or duplicated at the key level. ## Changes - `KeyAliases` is now `map[string]AliasConfig` instead of `map[string]string`. `AliasConfig` carries `ModelID` (the wire identifier), optional `ModelName`, `ModelFamily` (a typed enum for routing decisions), `Description`, `Region`, and embedded provider sub-configs for Azure, Vertex, Bedrock, and Replicate. - `AliasConfig.MarshalJSON` emits the legacy `{"k":"v"}` string-valued wire shape when only `ModelID` is set, preserving byte-for-byte JSON compatibility with pre-refactor consumers and keeping `config_hash` stable for unenriched entries. - `KeyAliases.UnmarshalJSON` transparently accepts both the legacy string shape and the new object shape, promoting legacy string values to `AliasConfig{ModelID: }`. - `KeyAliases.Resolve` is preserved for backward compatibility. A new `ResolveConfig` method returns the full `AliasConfig` for callers that need more than the wire model string. - `KeyAliases.Validate` is extended to check `ModelName` whitespace and `ModelFamily` validity. - `ModelFamily` is introduced as a typed enum (`anthropic`, `openai`, `mistral`, `cohere`, `gemini`, `nova`, `titan`) with an `IsValid` method, enabling provider routing decisions without substring-sniffing the wire model ID. - All provider `ToBifrostListModelsResponse` and `ListModelsPipeline` call sites are updated to use `schemas.KeyAliases` and access `.ModelID` from `AliasConfig` values. - The JSON schema (`config.schema.json`) is updated so the `aliases` property accepts either the legacy string shape or the new object shape via `oneOf`. - New tests cover legacy/rich/mixed unmarshal, round-trip marshal stability, `Resolve`/`ResolveConfig` behavior, `Validate` error cases, `ModelFamily.IsValid`, DB persistence of both wire shapes, and hash stability guarantees. ## Type of change - [ ] Bug fix - [ ] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./core/schemas/... go test ./core/providers/... go test ./framework/configstore/... go test ./transports/bifrost-http/lib/... ``` Existing configs using the legacy `"alias": "model-id"` string shape require no changes — they deserialize and re-serialize identically. To opt into the rich shape, update an alias entry to the object form: ```json "aliases": { "my-model": { "model_id": "azure-deployment-xyz", "model_family": "anthropic", "model_name": "claude-3-5-sonnet", "api_version": "2024-08-01-preview" } } ``` ## Breaking changes - [x] Yes - [ ] No Any code that directly reads `KeyAliases` values as `string` (e.g. `aliases["key"]` expecting a `string`) must be updated to access `.ModelID` on the returned `AliasConfig`. The JSON wire format for unenriched aliases is unchanged. The `Resolve(model string) string` method signature is unchanged. ## Related issues ## Security considerations No new secrets or auth surfaces are introduced. Provider sub-config fields that accept `EnvVar` values follow the existing env-var resolution and encryption patterns already in place for key-level configs. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable --- core/internal/llmtests/account.go | 54 +-- core/providers/anthropic/models.go | 2 +- core/providers/azure/models.go | 2 +- core/providers/bedrock/models.go | 2 +- core/providers/bedrock/rerank_test.go | 4 +- core/providers/cohere/models.go | 2 +- core/providers/elevenlabs/models.go | 2 +- core/providers/gemini/models.go | 2 +- core/providers/huggingface/models.go | 2 +- core/providers/mistral/models.go | 2 +- core/providers/openai/models.go | 2 +- core/providers/openrouter/openrouter.go | 6 +- core/providers/replicate/models.go | 2 +- core/providers/utils/models.go | 20 +- core/providers/vertex/models.go | 4 +- core/providers/vertex/utils.go | 4 +- core/schemas/account.go | 196 ++++++++++- core/schemas/account_test.go | 308 ++++++++++++++++++ framework/configstore/encryption_test.go | 2 +- framework/configstore/keyhash_alias_test.go | 76 +++++ framework/configstore/migrations_test.go | 4 +- .../configstore/tables/encryption_test.go | 144 +++++++- transports/bifrost-http/lib/config_test.go | 130 ++++---- transports/config.schema.json | 67 +++- 24 files changed, 897 insertions(+), 142 deletions(-) create mode 100644 core/schemas/account_test.go create mode 100644 framework/configstore/keyhash_alias_test.go diff --git a/core/internal/llmtests/account.go b/core/internal/llmtests/account.go index c78ffbd508..6a2b0e5d8d 100644 --- a/core/internal/llmtests/account.go +++ b/core/internal/llmtests/account.go @@ -241,12 +241,12 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx context.Context, { Models: []string{"*"}, Weight: 1.0, - Aliases: map[string]string{ - "claude-3.7-sonnet": "us.anthropic.claude-3-7-sonnet-20250219-v1:0", - "claude-4-sonnet": "global.anthropic.claude-sonnet-4-20250514-v1:0", - "claude-4.5-sonnet": "global.anthropic.claude-sonnet-4-5-20250929-v1:0", - "claude-4.6-sonnet": "global.anthropic.claude-sonnet-4-6", - "claude-4.5-haiku": "global.anthropic.claude-haiku-4-5-20251001-v1:0", + Aliases: schemas.KeyAliases{ + "claude-3.7-sonnet": {ModelID: "us.anthropic.claude-3-7-sonnet-20250219-v1:0"}, + "claude-4-sonnet": {ModelID: "global.anthropic.claude-sonnet-4-20250514-v1:0"}, + "claude-4.5-sonnet": {ModelID: "global.anthropic.claude-sonnet-4-5-20250929-v1:0"}, + "claude-4.6-sonnet": {ModelID: "global.anthropic.claude-sonnet-4-6"}, + "claude-4.5-haiku": {ModelID: "global.anthropic.claude-haiku-4-5-20251001-v1:0"}, }, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("env.AWS_ACCESS_KEY_ID"), @@ -259,13 +259,13 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx context.Context, { Models: []string{"*"}, Weight: 1.0, - Aliases: map[string]string{ - "claude-3.5-sonnet": "anthropic.claude-3-5-sonnet-20240620-v1:0", - "claude-3.7-sonnet": "us.anthropic.claude-3-7-sonnet-20250219-v1:0", - "claude-4-sonnet": "global.anthropic.claude-sonnet-4-20250514-v1:0", - "claude-4.5-sonnet": "global.anthropic.claude-sonnet-4-5-20250929-v1:0", - "claude-4.6-sonnet": "global.anthropic.claude-sonnet-4-6", - "claude-4.5-haiku": "global.anthropic.claude-haiku-4-5-20251001-v1:0", + Aliases: schemas.KeyAliases{ + "claude-3.5-sonnet": {ModelID: "anthropic.claude-3-5-sonnet-20240620-v1:0"}, + "claude-3.7-sonnet": {ModelID: "us.anthropic.claude-3-7-sonnet-20250219-v1:0"}, + "claude-4-sonnet": {ModelID: "global.anthropic.claude-sonnet-4-20250514-v1:0"}, + "claude-4.5-sonnet": {ModelID: "global.anthropic.claude-sonnet-4-5-20250929-v1:0"}, + "claude-4.6-sonnet": {ModelID: "global.anthropic.claude-sonnet-4-6"}, + "claude-4.5-haiku": {ModelID: "global.anthropic.claude-haiku-4-5-20251001-v1:0"}, }, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("env.AWS_ACCESS_KEY_ID"), @@ -303,13 +303,13 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx context.Context, Models: []string{"*"}, Weight: 1.0, Aliases: schemas.KeyAliases{ - "gpt-4o": "gpt-4o", - "gpt-4o-backup": "gpt-4o-3", - "claude-opus-4-5": "claude-opus-4-5", - "o1": "o1", - "gpt-image-1": "gpt-image-1", - "text-embedding-ada-002": "text-embedding-ada-002", - "sora-2": "sora-2", + "gpt-4o": {ModelID: "gpt-4o"}, + "gpt-4o-backup": {ModelID: "gpt-4o-3"}, + "claude-opus-4-5": {ModelID: "claude-opus-4-5"}, + "o1": {ModelID: "o1"}, + "gpt-image-1": {ModelID: "gpt-image-1"}, + "text-embedding-ada-002": {ModelID: "text-embedding-ada-002"}, + "sora-2": {ModelID: "sora-2"}, }, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("env.AZURE_ENDPOINT"), @@ -324,10 +324,10 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx context.Context, Models: []string{"*"}, Weight: 1.0, Aliases: schemas.KeyAliases{ - "whisper": "whisper", - "whisper-1": "whisper", - "gpt-4o-mini-tts": "gpt-4o-mini-tts", - "gpt-4o-mini-audio-preview": "gpt-4o-mini-audio-preview", + "whisper": {ModelID: "whisper"}, + "whisper-1": {ModelID: "whisper"}, + "gpt-4o-mini-tts": {ModelID: "gpt-4o-mini-tts"}, + "gpt-4o-mini-audio-preview": {ModelID: "gpt-4o-mini-audio-preview"}, }, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("env.AZURE_ENDPOINT"), @@ -365,9 +365,9 @@ func (account *ComprehensiveTestAccount) GetKeysForProvider(ctx context.Context, Models: []string{"claude-sonnet-4-5", "claude-4.5-haiku", "claude-opus-4-5"}, Weight: 1.0, Aliases: schemas.KeyAliases{ - "claude-sonnet-4-5": "claude-sonnet-4-5", - "claude-4.5-haiku": "claude-haiku-4-5@20251001", - "claude-opus-4-5": "claude-opus-4-5", + "claude-sonnet-4-5": {ModelID: "claude-sonnet-4-5"}, + "claude-4.5-haiku": {ModelID: "claude-haiku-4-5@20251001"}, + "claude-opus-4-5": {ModelID: "claude-opus-4-5"}, }, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("env.VERTEX_PROJECT_ID"), diff --git a/core/providers/anthropic/models.go b/core/providers/anthropic/models.go index 3815a0244b..a09174c661 100644 --- a/core/providers/anthropic/models.go +++ b/core/providers/anthropic/models.go @@ -8,7 +8,7 @@ import ( "github.com/maximhq/bifrost/core/schemas" ) -func (response *AnthropicListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases map[string]string, unfiltered bool) *schemas.BifrostListModelsResponse { +func (response *AnthropicListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases schemas.KeyAliases, unfiltered bool) *schemas.BifrostListModelsResponse { if response == nil { return nil } diff --git a/core/providers/azure/models.go b/core/providers/azure/models.go index 5daca3836d..99f4c5ae86 100644 --- a/core/providers/azure/models.go +++ b/core/providers/azure/models.go @@ -7,7 +7,7 @@ import ( "github.com/maximhq/bifrost/core/schemas" ) -func (response *AzureListModelsResponse) ToBifrostListModelsResponse(allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases map[string]string, unfiltered bool) *schemas.BifrostListModelsResponse { +func (response *AzureListModelsResponse) ToBifrostListModelsResponse(allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases schemas.KeyAliases, unfiltered bool) *schemas.BifrostListModelsResponse { if response == nil { return nil } diff --git a/core/providers/bedrock/models.go b/core/providers/bedrock/models.go index 549db2e3bd..817a43f2b4 100644 --- a/core/providers/bedrock/models.go +++ b/core/providers/bedrock/models.go @@ -81,7 +81,7 @@ type BedrockRerankResponseDocument struct { TextDocument *BedrockRerankTextValue `json:"textDocument,omitempty"` } -func (response *BedrockListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases map[string]string, unfiltered bool) *schemas.BifrostListModelsResponse { +func (response *BedrockListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases schemas.KeyAliases, unfiltered bool) *schemas.BifrostListModelsResponse { if response == nil { return nil } diff --git a/core/providers/bedrock/rerank_test.go b/core/providers/bedrock/rerank_test.go index c1b7bb5480..5d55dccc69 100644 --- a/core/providers/bedrock/rerank_test.go +++ b/core/providers/bedrock/rerank_test.go @@ -196,7 +196,7 @@ func TestBedrockRerankRequestToBifrostRerankRequestNil(t *testing.T) { func TestResolveBedrockDeployment(t *testing.T) { key := schemas.Key{ Aliases: schemas.KeyAliases{ - "cohere-rerank": "arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0", + "cohere-rerank": {ModelID: "arn:aws:bedrock:us-east-1::foundation-model/cohere.rerank-v3-5:0"}, }, } @@ -211,7 +211,7 @@ func TestBedrockRerankRequiresARNModelIdentifier(t *testing.T) { ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) key := schemas.Key{ Aliases: schemas.KeyAliases{ - "cohere-rerank": "cohere.rerank-v3-5:0", + "cohere-rerank": {ModelID: "cohere.rerank-v3-5:0"}, }, } diff --git a/core/providers/cohere/models.go b/core/providers/cohere/models.go index 3b285f97b6..3312032888 100644 --- a/core/providers/cohere/models.go +++ b/core/providers/cohere/models.go @@ -45,7 +45,7 @@ type CohereRerankMeta struct { Tokens *CohereTokenUsage `json:"tokens,omitempty"` } -func (response *CohereListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases map[string]string, unfiltered bool) *schemas.BifrostListModelsResponse { +func (response *CohereListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases schemas.KeyAliases, unfiltered bool) *schemas.BifrostListModelsResponse { if response == nil { return nil } diff --git a/core/providers/elevenlabs/models.go b/core/providers/elevenlabs/models.go index f762d97ee8..7e3c1f8d34 100644 --- a/core/providers/elevenlabs/models.go +++ b/core/providers/elevenlabs/models.go @@ -7,7 +7,7 @@ import ( "github.com/maximhq/bifrost/core/schemas" ) -func (response *ElevenlabsListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases map[string]string, unfiltered bool) *schemas.BifrostListModelsResponse { +func (response *ElevenlabsListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases schemas.KeyAliases, unfiltered bool) *schemas.BifrostListModelsResponse { if response == nil { return nil } diff --git a/core/providers/gemini/models.go b/core/providers/gemini/models.go index 7b9f6410eb..e88387dc66 100644 --- a/core/providers/gemini/models.go +++ b/core/providers/gemini/models.go @@ -17,7 +17,7 @@ func toGeminiModelResourceName(modelID string) string { return "models/" + modelID } -func (response *GeminiListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases map[string]string, unfiltered bool) *schemas.BifrostListModelsResponse { +func (response *GeminiListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases schemas.KeyAliases, unfiltered bool) *schemas.BifrostListModelsResponse { if response == nil { return nil } diff --git a/core/providers/huggingface/models.go b/core/providers/huggingface/models.go index de615ccec2..3d04ce4936 100644 --- a/core/providers/huggingface/models.go +++ b/core/providers/huggingface/models.go @@ -14,7 +14,7 @@ const ( maxModelFetchLimit = 1000 ) -func (response *HuggingFaceListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, inferenceProvider inferenceProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases map[string]string, unfiltered bool) *schemas.BifrostListModelsResponse { +func (response *HuggingFaceListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, inferenceProvider inferenceProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases schemas.KeyAliases, unfiltered bool) *schemas.BifrostListModelsResponse { if response == nil { return nil } diff --git a/core/providers/mistral/models.go b/core/providers/mistral/models.go index 8d5fd7f3d6..7db5154aa6 100644 --- a/core/providers/mistral/models.go +++ b/core/providers/mistral/models.go @@ -7,7 +7,7 @@ import ( "github.com/maximhq/bifrost/core/schemas" ) -func (response *MistralListModelsResponse) ToBifrostListModelsResponse(allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases map[string]string, unfiltered bool) *schemas.BifrostListModelsResponse { +func (response *MistralListModelsResponse) ToBifrostListModelsResponse(allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases schemas.KeyAliases, unfiltered bool) *schemas.BifrostListModelsResponse { if response == nil { return nil } diff --git a/core/providers/openai/models.go b/core/providers/openai/models.go index a76d350d28..ed88255cb3 100644 --- a/core/providers/openai/models.go +++ b/core/providers/openai/models.go @@ -8,7 +8,7 @@ import ( ) // ToBifrostListModelsResponse converts an OpenAI list models response to a Bifrost list models response -func (response *OpenAIListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases map[string]string, unfiltered bool) *schemas.BifrostListModelsResponse { +func (response *OpenAIListModelsResponse) ToBifrostListModelsResponse(providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases schemas.KeyAliases, unfiltered bool) *schemas.BifrostListModelsResponse { if response == nil { return nil } diff --git a/core/providers/openrouter/openrouter.go b/core/providers/openrouter/openrouter.go index eae6de2bf9..516f52477e 100644 --- a/core/providers/openrouter/openrouter.go +++ b/core/providers/openrouter/openrouter.go @@ -206,9 +206,11 @@ func (provider *OpenRouterProvider) listModelsByKey(ctx *schemas.BifrostContext, for _, m := range key.BlacklistedModels { normalizedBlacklist = append(normalizedBlacklist, stripPrefix(m)) } - normalizedAliases := make(map[string]string, len(key.Aliases)) + normalizedAliases := make(schemas.KeyAliases, len(key.Aliases)) for k, v := range key.Aliases { - normalizedAliases[stripPrefix(k)] = stripPrefix(v) + cfg := v + cfg.ModelID = stripPrefix(v.ModelID) + normalizedAliases[stripPrefix(k)] = cfg } pipeline := &providerUtils.ListModelsPipeline{ diff --git a/core/providers/replicate/models.go b/core/providers/replicate/models.go index 6c0c14dbf7..3d2c4b6081 100644 --- a/core/providers/replicate/models.go +++ b/core/providers/replicate/models.go @@ -14,7 +14,7 @@ func ToBifrostListModelsResponse( providerKey schemas.ModelProvider, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, - aliases map[string]string, + aliases schemas.KeyAliases, unfiltered bool, ) *schemas.BifrostListModelsResponse { bifrostResponse := &schemas.BifrostListModelsResponse{ diff --git a/core/providers/utils/models.go b/core/providers/utils/models.go index dbbbd8071a..22555f11ae 100644 --- a/core/providers/utils/models.go +++ b/core/providers/utils/models.go @@ -168,9 +168,9 @@ type FilterResult struct { type ListModelsPipeline struct { AllowedModels schemas.WhiteList BlacklistedModels schemas.BlackList - // Aliases maps user-facing alias keys to provider-specific model IDs. - // e.g. {"my-gpt4": "gpt-4-turbo-2024-04-09"} - Aliases map[string]string + // Aliases maps user-facing alias keys to their AliasConfig. The pipeline + // reads AliasConfig.ModelID for matching and Alias surfacing. + Aliases schemas.KeyAliases Unfiltered bool ProviderKey schemas.ModelProvider // MatchFns is the ordered list of equivalence functions used for every @@ -224,9 +224,9 @@ type aliasMatch struct { // → [{key:"gpt-3.5-turbo", value:""}] func (p *ListModelsPipeline) resolveModelID(modelID string) []aliasMatch { var candidates []aliasMatch - for aliasKey, providerID := range p.Aliases { - if matches(modelID, providerID, p.MatchFns) { - candidates = append(candidates, aliasMatch{key: aliasKey, value: providerID}) + for aliasKey, alias := range p.Aliases { + if matches(modelID, alias.ModelID, p.MatchFns) { + candidates = append(candidates, aliasMatch{key: aliasKey, value: alias.ModelID}) } } if len(candidates) == 0 { @@ -369,9 +369,9 @@ func (p *ListModelsPipeline) BackfillModels(included map[string]bool) []schemas. Name: schemas.Ptr(ToDisplayName(entry)), } // If this allowlist entry has an alias, surface the provider-specific ID. - for aliasKey, providerID := range p.Aliases { + for aliasKey, alias := range p.Aliases { if matches(entry, aliasKey, p.MatchFns) { - m.Alias = schemas.Ptr(providerID) + m.Alias = schemas.Ptr(alias.ModelID) break } } @@ -382,7 +382,7 @@ func (p *ListModelsPipeline) BackfillModels(included map[string]bool) []schemas. // Case B: wildcard allowlist — backfill only explicitly configured aliases. if !p.Unfiltered && len(p.Aliases) > 0 { - for aliasKey, providerID := range p.Aliases { + for aliasKey, alias := range p.Aliases { if included[strings.ToLower(aliasKey)] { continue } @@ -400,7 +400,7 @@ func (p *ListModelsPipeline) BackfillModels(included map[string]bool) []schemas. result = append(result, schemas.Model{ ID: string(p.ProviderKey) + "/" + aliasKey, Name: schemas.Ptr(ToDisplayName(aliasKey)), - Alias: schemas.Ptr(providerID), + Alias: schemas.Ptr(alias.ModelID), }) } } diff --git a/core/providers/vertex/models.go b/core/providers/vertex/models.go index d373f58735..e27db45bd9 100644 --- a/core/providers/vertex/models.go +++ b/core/providers/vertex/models.go @@ -70,7 +70,7 @@ type vertexRerankOptions struct { // - If allowedModels is empty, all models are allowed // - If allowedModels is non-empty, only models/deployments with keys in allowedModels are included // - Deployments map is used to match model IDs to aliases and filter accordingly -func (response *VertexListModelsResponse) ToBifrostListModelsResponse(allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases map[string]string, unfiltered bool) *schemas.BifrostListModelsResponse { +func (response *VertexListModelsResponse) ToBifrostListModelsResponse(allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases schemas.KeyAliases, unfiltered bool) *schemas.BifrostListModelsResponse { if response == nil { return nil } @@ -140,7 +140,7 @@ func (response *VertexListModelsResponse) ToBifrostListModelsResponse(allowedMod // ToBifrostListModelsResponse converts a Vertex AI publisher models response to Bifrost's format. // This is for foundation models from the Model Garden (publishers.models.list endpoint). -func (response *VertexListPublisherModelsResponse) ToBifrostListModelsResponse(allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases map[string]string, unfiltered bool) *schemas.BifrostListModelsResponse { +func (response *VertexListPublisherModelsResponse) ToBifrostListModelsResponse(allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList, aliases schemas.KeyAliases, unfiltered bool) *schemas.BifrostListModelsResponse { if response == nil { return nil } diff --git a/core/providers/vertex/utils.go b/core/providers/vertex/utils.go index 7eb9161f6d..ef806ab750 100644 --- a/core/providers/vertex/utils.go +++ b/core/providers/vertex/utils.go @@ -242,7 +242,7 @@ func vertexServiceTierHeaderValue(region string, model string, tier schemas.Bifr // buildResponseFromConfig builds a list models response from configured deployments and allowedModels. // This is used when the user has explicitly configured which models they want to use. -func buildResponseFromConfig(deployments map[string]string, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList) *schemas.BifrostListModelsResponse { +func buildResponseFromConfig(deployments schemas.KeyAliases, allowedModels schemas.WhiteList, blacklistedModels schemas.BlackList) *schemas.BifrostListModelsResponse { response := &schemas.BifrostListModelsResponse{ Data: make([]schemas.Model, 0), } @@ -272,7 +272,7 @@ func buildResponseFromConfig(deployments map[string]string, allowedModels schema modelEntry := schemas.Model{ ID: modelID, Name: schemas.Ptr(modelName), - Alias: schemas.Ptr(deploymentValue), + Alias: schemas.Ptr(deploymentValue.ModelID), } response.Data = append(response.Data, modelEntry) diff --git a/core/schemas/account.go b/core/schemas/account.go index 449399f554..38d8083728 100644 --- a/core/schemas/account.go +++ b/core/schemas/account.go @@ -2,7 +2,9 @@ package schemas import ( + "bytes" "context" + "encoding/json" "fmt" "slices" "strings" @@ -141,22 +143,136 @@ type Key struct { Description string `json:"description,omitempty"` // Description of key } -type KeyAliases map[string]string +// ModelFamily is a typed enum identifying the underlying model family of an alias target. +// It enables provider routing decisions (request shape, response parsing, auth headers, +// URL construction) without substring-sniffing the wire model ID. +type ModelFamily string + +const ( + ModelFamilyAnthropic ModelFamily = "anthropic" + ModelFamilyOpenAI ModelFamily = "openai" + ModelFamilyMistral ModelFamily = "mistral" + ModelFamilyCohere ModelFamily = "cohere" + ModelFamilyGemini ModelFamily = "gemini" + ModelFamilyNova ModelFamily = "nova" + ModelFamilyTitan ModelFamily = "titan" +) + +// IsValid reports whether mf is a recognized model family. +func (mf *ModelFamily) IsValid() bool { + if mf == nil { + return false + } + switch *mf { + case ModelFamilyAnthropic, ModelFamilyOpenAI, ModelFamilyMistral, + ModelFamilyCohere, ModelFamilyGemini, ModelFamilyNova, ModelFamilyTitan: + return true + } + return false +} + +// AzureAliasCfg holds Azure-specific overrides that apply to a single alias. +// Each field, when non-nil, overrides the corresponding key-level default. +type AzureAliasCfg struct { + APIVersion *string `json:"api_version,omitempty"` // overrides the Azure OpenAI api-version query param for this alias + AnthropicVersion *string `json:"anthropic_version,omitempty"` // overrides the anthropic-version header for Claude-on-Azure deployments + Endpoint *EnvVar `json:"endpoint,omitempty"` // overrides AzureKeyConfig.Endpoint for this alias (allows one credential to span multiple Azure resources) +} + +// VertexAliasCfg holds Vertex-specific overrides that apply to a single alias. +type VertexAliasCfg struct { + ProjectID *EnvVar `json:"project_id,omitempty"` + ProjectNumber *EnvVar `json:"project_number,omitempty"` +} + +// BedrockAliasCfg holds Bedrock-specific overrides that apply to a single alias. +type BedrockAliasCfg struct { + InferenceProfileARN *EnvVar `json:"inference_profile_arn,omitempty"` +} + +// ReplicateAliasCfg holds Replicate-specific overrides that apply to a single alias. +type ReplicateAliasCfg struct { + UseDeploymentsEndpoint *bool `json:"use_deployments_endpoint,omitempty"` +} + +// AliasConfig is the rich value type held by KeyAliases. It carries everything +// needed to call a provider for an aliased model: the wire model identifier +// (ModelID), the canonical model name used for pricing/logging (ModelName), the +// family used for provider routing decisions (ModelFamily), and optional +// provider-specific overrides that override the key-level defaults. +type AliasConfig struct { + ModelID string `json:"model_id"` // wire model identifier sent to the provider + ModelName *string `json:"model_name,omitempty"` // canonical model name used for pricing, logging, and 2nd-tier family routing + ModelFamily *ModelFamily `json:"model_family,omitempty"` // 1st-tier family routing enum + Description string `json:"description,omitempty"` // description of the alias for users to understand its purpose (not used by bifrost) + Region *EnvVar `json:"region,omitempty"` + + *AzureAliasCfg + *VertexAliasCfg + *BedrockAliasCfg + *ReplicateAliasCfg +} + +// isLegacyShape reports whether this AliasConfig carries only ModelID and no +// other fields. Used by MarshalJSON to emit the legacy string-valued wire +// shape so older consumers that expect map[string]string keep working. +func (ac AliasConfig) isLegacyShape() bool { + return ac.ModelID != "" && + ac.ModelName == nil && + ac.ModelFamily == nil && + ac.Description == "" && + ac.Region == nil && + ac.AzureAliasCfg == nil && + ac.VertexAliasCfg == nil && + ac.BedrockAliasCfg == nil && + ac.ReplicateAliasCfg == nil +} + +// MarshalJSON emits the legacy string wire shape when only ModelID is set, so +// callers that haven't opted into the rich AliasConfig see no observable +// change on the wire. When any other field is populated, the full object is +// emitted. +func (ac AliasConfig) MarshalJSON() ([]byte, error) { + if ac.isLegacyShape() { + return Marshal(ac.ModelID) + } + type aliasConfigJSON AliasConfig + return Marshal(aliasConfigJSON(ac)) +} + +// KeyAliases maps a user-facing model name to its AliasConfig. +// +// Both the input (UnmarshalJSON) and the output (AliasConfig.MarshalJSON) +// transparently accept and emit two JSON wire shapes: +// - Legacy: {"my-model": "provider-model-id"} — value is a string +// - New: {"my-model": {"model_id": "provider-model-id", ... }} — value is an object +// +// Legacy entries deserialize to AliasConfig{ModelID: }; an AliasConfig +// that only has ModelID set serializes back to a plain string. This keeps the +// wire format byte-for-byte compatible with the pre-refactor flow until +// ModelName / ModelFamily / provider sub-configs are populated explicitly. +type KeyAliases map[string]AliasConfig func (ka KeyAliases) Validate() error { seen := make(map[string]struct{}, len(ka)) - for from, to := range ka { + for from, ac := range ka { if strings.TrimSpace(from) == "" { return fmt.Errorf("alias source cannot be empty") } - if strings.TrimSpace(to) == "" { - return fmt.Errorf("alias target for %q cannot be empty", from) + if strings.TrimSpace(ac.ModelID) == "" { + return fmt.Errorf("alias %q: model_id cannot be empty", from) } if strings.TrimSpace(from) != from { return fmt.Errorf("alias source %q cannot have leading or trailing whitespace", from) } - if strings.TrimSpace(to) != to { - return fmt.Errorf("alias target for %q cannot have leading or trailing whitespace", from) + if strings.TrimSpace(ac.ModelID) != ac.ModelID { + return fmt.Errorf("alias %q: model_id cannot have leading or trailing whitespace", from) + } + if ac.ModelName != nil && strings.TrimSpace(*ac.ModelName) != *ac.ModelName { + return fmt.Errorf("alias %q: model_name cannot have leading or trailing whitespace", from) + } + if ac.ModelFamily != nil && !ac.ModelFamily.IsValid() { + return fmt.Errorf("alias %q: invalid model_family %q", from, *ac.ModelFamily) } normalized := strings.ToLower(from) if _, ok := seen[normalized]; ok { @@ -167,20 +283,76 @@ func (ka KeyAliases) Validate() error { return nil } +// Resolve returns the wire model identifier for the given user-facing model name. +// If no alias matches, the input is returned unchanged. Case-insensitive fallback +// matches the prior behavior. +// +// This signature is preserved for backward compatibility with existing callers +// that only need the wire model string. For access to the full AliasConfig +// (ModelName, ModelFamily, provider overrides), use ResolveConfig. func (ka KeyAliases) Resolve(model string) string { + if ac := ka.ResolveConfig(model); ac != nil { + return ac.ModelID + } + return model +} + +// 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 { if ka == nil { - return model + return nil } - if alias, ok := ka[model]; ok { - return alias + if ac, ok := ka[model]; ok { + return &ac } - // Fall back to case-insensitive lookup for consistency with WhiteList.Contains for k, v := range ka { if strings.EqualFold(k, model) { - return v + return &v } } - return model + return nil +} + +// UnmarshalJSON accepts both the legacy {"k":"v"} and new {"k":{...}} wire +// shapes for KeyAliases. Legacy string values are promoted to +// AliasConfig{ModelID: }. +func (ka *KeyAliases) UnmarshalJSON(data []byte) error { + trimmed := bytes.TrimSpace(data) + if len(trimmed) == 0 || string(trimmed) == "null" { + *ka = nil + return nil + } + var raw map[string]json.RawMessage + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + result := make(KeyAliases, len(raw)) + for k, entry := range raw { + entryTrim := bytes.TrimSpace(entry) + if len(entryTrim) == 0 { + return fmt.Errorf("alias %q: empty value", k) + } + switch entryTrim[0] { + case '"': + // Legacy string value — promote to AliasConfig{ModelID: ...}. + var modelID string + if err := json.Unmarshal(entry, &modelID); err != nil { + return fmt.Errorf("alias %q: %w", k, err) + } + result[k] = AliasConfig{ModelID: modelID} + case '{': + var ac AliasConfig + if err := json.Unmarshal(entry, &ac); err != nil { + return fmt.Errorf("alias %q: %w", k, err) + } + result[k] = ac + default: + return fmt.Errorf("alias %q: value must be a string (legacy) or object", k) + } + } + *ka = result + return nil } type AzureAuthType string diff --git a/core/schemas/account_test.go b/core/schemas/account_test.go new file mode 100644 index 0000000000..18934bb41b --- /dev/null +++ b/core/schemas/account_test.go @@ -0,0 +1,308 @@ +package schemas + +import ( + "encoding/json" + "reflect" + "strings" + "testing" +) + +func TestKeyAliasesUnmarshalLegacyStringShape(t *testing.T) { + in := []byte(`{"best-model": "gpt-4o-deployment"}`) + var ka KeyAliases + if err := json.Unmarshal(in, &ka); err != nil { + t.Fatalf("unmarshal: %v", err) + } + want := KeyAliases{"best-model": AliasConfig{ModelID: "gpt-4o-deployment"}} + if !reflect.DeepEqual(ka, want) { + t.Fatalf("legacy shape mismatch: got %+v, want %+v", ka, want) + } +} + +func TestKeyAliasesUnmarshalRichShape(t *testing.T) { + // Provider sub-configs are embedded, so their fields appear at the top level of the JSON. + in := []byte(`{ + "best-model": { + "model_id": "azure-deployment-xyz", + "model_name": "claude-3-5-sonnet", + "model_family": "anthropic", + "description": "prod", + "api_version": "2024-08-01-preview" + } + }`) + var ka KeyAliases + if err := json.Unmarshal(in, &ka); err != nil { + t.Fatalf("unmarshal: %v", err) + } + got := ka["best-model"] + if got.ModelID != "azure-deployment-xyz" { + t.Fatalf("ModelID mismatch: %q", got.ModelID) + } + if got.ModelName == nil || *got.ModelName != "claude-3-5-sonnet" { + t.Fatalf("ModelName mismatch: %+v", got.ModelName) + } + if got.ModelFamily == nil || *got.ModelFamily != ModelFamilyAnthropic { + t.Fatalf("ModelFamily mismatch: %+v", got.ModelFamily) + } + if got.Description != "prod" { + t.Fatalf("Description mismatch: %q", got.Description) + } + if got.AzureAliasCfg == nil || got.APIVersion == nil || *got.APIVersion != "2024-08-01-preview" { + t.Fatalf("AzureAliasCfg.APIVersion mismatch: %+v", got.AzureAliasCfg) + } +} + +func TestKeyAliasesUnmarshalMixedShape(t *testing.T) { + in := []byte(`{ + "legacy": "gpt-4-deployment", + "rich": {"model_id": "azure-xyz", "model_family": "openai"} + }`) + var ka KeyAliases + if err := json.Unmarshal(in, &ka); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if got := ka["legacy"]; got.ModelID != "gpt-4-deployment" || got.ModelFamily != nil { + t.Fatalf("legacy entry wrong: %+v", got) + } + got := ka["rich"] + if got.ModelID != "azure-xyz" || got.ModelFamily == nil || *got.ModelFamily != ModelFamilyOpenAI { + t.Fatalf("rich entry wrong: %+v", got) + } +} + +func TestKeyAliasesUnmarshalEmptyAndNull(t *testing.T) { + cases := map[string]string{ + "empty-obj": `{}`, + "null": `null`, + } + for name, in := range cases { + var ka KeyAliases + if err := json.Unmarshal([]byte(in), &ka); err != nil { + t.Fatalf("%s: unmarshal: %v", name, err) + } + if len(ka) != 0 { + t.Fatalf("%s: want empty/nil, got %+v", name, ka) + } + } +} + +func TestKeyAliasesUnmarshalRoundTrip(t *testing.T) { + orig := KeyAliases{ + "best-model": AliasConfig{ + ModelID: "azure-xyz", + ModelName: Ptr("claude-3-5-sonnet"), + ModelFamily: Ptr(ModelFamilyAnthropic), + AzureAliasCfg: &AzureAliasCfg{ + APIVersion: Ptr("2024-08-01-preview"), + }, + }, + "simple": AliasConfig{ModelID: "gpt-4"}, + } + data, err := json.Marshal(orig) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var back KeyAliases + if err := json.Unmarshal(data, &back); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if !reflect.DeepEqual(orig, back) { + t.Fatalf("round-trip mismatch:\nwant: %+v\ngot: %+v", orig, back) + } +} + +func TestKeyAliasesMarshalLegacyShapeWhenOnlyModelIDSet(t *testing.T) { + // Only ModelID populated — should serialize to the legacy string-valued shape. + ka := KeyAliases{"best-model": AliasConfig{ModelID: "gpt-4o-deployment"}} + data, err := json.Marshal(ka) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if string(data) != `{"best-model":"gpt-4o-deployment"}` { + t.Fatalf("legacy shape mismatch: got %s", data) + } +} + +func TestKeyAliasesMarshalRichShapeWhenAnyExtraFieldSet(t *testing.T) { + cases := map[string]struct { + ac AliasConfig + wantKey string + wantValue any + }{ + "with_model_name": {AliasConfig{ModelID: "x", ModelName: Ptr("canonical")}, "model_name", "canonical"}, + "with_model_family": {AliasConfig{ModelID: "x", ModelFamily: Ptr(ModelFamilyAnthropic)}, "model_family", "anthropic"}, + "with_description": {AliasConfig{ModelID: "x", Description: "prod"}, "description", "prod"}, + "with_azure_subcfg": {AliasConfig{ModelID: "x", AzureAliasCfg: &AzureAliasCfg{APIVersion: Ptr("2024-08-01-preview")}}, "api_version", "2024-08-01-preview"}, + } + for name, c := range cases { + t.Run(name, func(t *testing.T) { + data, err := json.Marshal(c.ac) + if err != nil { + t.Fatalf("marshal: %v", err) + } + // Rich shape should be a JSON object, not a string. + if len(data) == 0 || data[0] != '{' { + t.Fatalf("want object shape, got %s", data) + } + var out map[string]any + if err := json.Unmarshal(data, &out); err != nil { + t.Fatalf("re-unmarshal: %v", err) + } + got, ok := out[c.wantKey] + if !ok { + t.Fatalf("expected key %q in serialized output, got %s", c.wantKey, data) + } + if got != c.wantValue { + t.Fatalf("field %q: want %v, got %v (raw: %s)", c.wantKey, c.wantValue, got, data) + } + }) + } +} + +func TestKeyAliasesMarshalUnmarshalLegacyRoundTrip(t *testing.T) { + // Legacy in → legacy out: byte-for-byte stable for the unenriched case. + in := []byte(`{"best-model":"gpt-4o-deployment"}`) + var ka KeyAliases + if err := json.Unmarshal(in, &ka); err != nil { + t.Fatalf("unmarshal: %v", err) + } + out, err := json.Marshal(ka) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if string(out) != string(in) { + t.Fatalf("round-trip drift:\n in: %s\nout: %s", in, out) + } +} + +func TestKeyAliasesUnmarshalInvalidValueType(t *testing.T) { + for _, in := range []string{ + `{"k": 123}`, + `{"k": [1,2]}`, + `{"k": true}`, + } { + var ka KeyAliases + if err := json.Unmarshal([]byte(in), &ka); err == nil { + t.Fatalf("expected error for %q, got nil", in) + } + } +} + +func TestKeyAliasesResolveBackwardCompat(t *testing.T) { + ka := KeyAliases{ + "best-model": AliasConfig{ModelID: "gpt-4o-deployment"}, + } + if got := ka.Resolve("best-model"); got != "gpt-4o-deployment" { + t.Fatalf("Resolve mismatch: %q", got) + } + if got := ka.Resolve("BEST-MODEL"); got != "gpt-4o-deployment" { + t.Fatalf("Resolve case-insensitive fallback failed: %q", got) + } + if got := ka.Resolve("unmapped"); got != "unmapped" { + t.Fatalf("Resolve unmatched mismatch: %q", got) + } + var nilKA KeyAliases + if got := nilKA.Resolve("x"); got != "x" { + t.Fatalf("nil Resolve mismatch: %q", got) + } +} + +func TestKeyAliasesResolveConfig(t *testing.T) { + ka := KeyAliases{ + "best-model": AliasConfig{ModelID: "azure-xyz", ModelFamily: Ptr(ModelFamilyAnthropic)}, + } + got := ka.ResolveConfig("best-model") + if got == nil || got.ModelID != "azure-xyz" || got.ModelFamily == nil || *got.ModelFamily != ModelFamilyAnthropic { + t.Fatalf("ResolveConfig mismatch: %+v", got) + } + if ka.ResolveConfig("unmapped") != nil { + t.Fatalf("ResolveConfig should return nil for unmapped") + } +} + +func TestKeyAliasesValidate(t *testing.T) { + madeUp := ModelFamily("made-up") + cases := []struct { + name string + ka KeyAliases + wantErr string + }{ + { + name: "ok", + ka: KeyAliases{"k": {ModelID: "v"}}, + }, + { + name: "empty source", + ka: KeyAliases{"": {ModelID: "v"}}, + wantErr: "alias source cannot be empty", + }, + { + name: "empty model id", + ka: KeyAliases{"k": {ModelID: ""}}, + wantErr: "model_id cannot be empty", + }, + { + name: "whitespace source", + ka: KeyAliases{" k ": {ModelID: "v"}}, + wantErr: "leading or trailing whitespace", + }, + { + name: "whitespace model_id", + ka: KeyAliases{"k": {ModelID: "v "}}, + wantErr: "model_id cannot have leading or trailing whitespace", + }, + { + name: "whitespace model_name", + ka: KeyAliases{"k": {ModelID: "v", ModelName: Ptr(" canonical ")}}, + wantErr: "model_name cannot have leading or trailing whitespace", + }, + { + name: "duplicate source case-insensitive", + ka: KeyAliases{"Key": {ModelID: "v"}, "key": {ModelID: "v"}}, + wantErr: "duplicate alias source", + }, + { + name: "invalid family", + ka: KeyAliases{"k": {ModelID: "v", ModelFamily: &madeUp}}, + wantErr: "invalid model_family", + }, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + err := c.ka.Validate() + if c.wantErr == "" { + if err != nil { + t.Fatalf("want ok, got %v", err) + } + return + } + if err == nil || !strings.Contains(err.Error(), c.wantErr) { + t.Fatalf("want error containing %q, got %v", c.wantErr, err) + } + }) + } +} + +func TestModelFamilyIsValid(t *testing.T) { + valid := []ModelFamily{ + ModelFamilyAnthropic, ModelFamilyOpenAI, ModelFamilyMistral, + ModelFamilyCohere, ModelFamilyGemini, ModelFamilyNova, ModelFamilyTitan, + } + for _, mf := range valid { + v := mf + if !v.IsValid() { + t.Fatalf("%q should be valid", mf) + } + } + for _, mf := range []ModelFamily{"", "unknown", "claude"} { + v := mf + if v.IsValid() { + t.Fatalf("%q should be invalid", mf) + } + } + // nil receiver is invalid. + var nilMF *ModelFamily + if nilMF.IsValid() { + t.Fatal("nil ModelFamily should be invalid") + } +} diff --git a/framework/configstore/encryption_test.go b/framework/configstore/encryption_test.go index 4b5c553904..6a795c6fca 100644 --- a/framework/configstore/encryption_test.go +++ b/framework/configstore/encryption_test.go @@ -775,7 +775,7 @@ func TestEncryptPlaintextKeys_BedrockFields_EncryptsAndDecryptsCorrectly(t *test assert.Equal(t, "us-west-2", found.BedrockKeyConfig.Region.GetValue()) require.NotNil(t, found.BedrockKeyConfig.ARN) assert.Equal(t, "arn:aws:iam::123456789:role/bedrock", found.BedrockKeyConfig.ARN.GetValue()) - assert.Equal(t, "profile-claude", found.Aliases["claude-3"]) + assert.Equal(t, "profile-claude", found.Aliases["claude-3"].ModelID) require.NotNil(t, found.BedrockKeyConfig.BatchS3Config) require.Len(t, found.BedrockKeyConfig.BatchS3Config.Buckets, 1) assert.Equal(t, "my-bucket", found.BedrockKeyConfig.BatchS3Config.Buckets[0].BucketName) diff --git a/framework/configstore/keyhash_alias_test.go b/framework/configstore/keyhash_alias_test.go new file mode 100644 index 0000000000..604ab03fbe --- /dev/null +++ b/framework/configstore/keyhash_alias_test.go @@ -0,0 +1,76 @@ +package configstore + +import ( + "testing" + + "github.com/bytedance/sonic" + "github.com/maximhq/bifrost/core/schemas" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestGenerateKeyHash_LegacyAliasesPreserveByteShape proves that an +// unenriched alias (only ModelID set) marshals into the hasher as the legacy +// {"k":"v"} string-valued shape, which is what keeps config_hash byte-stable +// across the refactor. If MarshalJSON ever stops emitting the legacy form for +// ModelID-only entries, this test fires. +// +// Strategy: hash the same Key with two equivalent alias representations — the +// rich KeyAliases{"k": {ModelID: "v"}} and (a hand-rolled JSON for) the +// legacy "k": "v" shape — and confirm both feed identical bytes into the +// hasher by checking that the marshaled outputs match. We don't recompute the +// full SHA256 since GenerateKeyHash composes many field bytes; the marshaling +// stability of the alias map alone is the regression-prone surface. +func TestGenerateKeyHash_LegacyAliasesPreserveByteShape(t *testing.T) { + key := schemas.Key{ + Name: "openai-key", + Value: *schemas.NewEnvVar("sk-test"), + Weight: 1.0, + Aliases: schemas.KeyAliases{"best-model": {ModelID: "gpt-4o-deployment"}}, + } + + gotMarshal, err := sonic.Marshal(key.Aliases) + require.NoError(t, err) + assert.Equal(t, + `{"best-model":"gpt-4o-deployment"}`, + string(gotMarshal), + "unenriched alias should marshal to the legacy string-valued wire shape; otherwise GenerateKeyHash drifts from pre-refactor rows", + ) + + // And GenerateKeyHash itself runs cleanly with the new types. + hash, err := GenerateKeyHash(key) + require.NoError(t, err) + assert.NotEmpty(t, hash) +} + +// TestGenerateKeyHash_RichAliasesProduceDifferentHash sanity-checks the other +// side: enriching an alias with ModelName/Family/etc. *does* change the hash, +// so genuine config changes are still detected. +func TestGenerateKeyHash_RichAliasesProduceDifferentHash(t *testing.T) { + canonical := "gpt-4o" + family := schemas.ModelFamilyOpenAI + + legacy := schemas.Key{ + Name: "k", + Value: *schemas.NewEnvVar("sk"), + Weight: 1.0, + Aliases: schemas.KeyAliases{"x": {ModelID: "y"}}, + } + rich := schemas.Key{ + Name: "k", + Value: *schemas.NewEnvVar("sk"), + Weight: 1.0, + Aliases: schemas.KeyAliases{"x": { + ModelID: "y", + ModelName: &canonical, + ModelFamily: &family, + }}, + } + + legacyHash, err := GenerateKeyHash(legacy) + require.NoError(t, err) + richHash, err := GenerateKeyHash(rich) + require.NoError(t, err) + + assert.NotEqual(t, legacyHash, richHash, "enriching an alias must change the key hash so config diffs are detected") +} diff --git a/framework/configstore/migrations_test.go b/framework/configstore/migrations_test.go index f522c7cb95..5a274d533b 100644 --- a/framework/configstore/migrations_test.go +++ b/framework/configstore/migrations_test.go @@ -1090,8 +1090,8 @@ func TestMigrationDropDeploymentColumnsAndAddAliases_BedrockEncrypted(t *testing // Verify the aliases contain the original deployment data (not double-encrypted) aliases := keys[0].Aliases assert.Contains(t, aliases, "claude") - assert.Equal(t, "dep-claude", aliases["claude"]) - assert.Equal(t, "dep-instant", aliases["claude-instant"]) + assert.Equal(t, "dep-claude", aliases["claude"].ModelID) + assert.Equal(t, "dep-instant", aliases["claude-instant"].ModelID) } // ============================================================================ diff --git a/framework/configstore/tables/encryption_test.go b/framework/configstore/tables/encryption_test.go index 9b329fbe51..2454f5cfd1 100644 --- a/framework/configstore/tables/encryption_test.go +++ b/framework/configstore/tables/encryption_test.go @@ -178,7 +178,7 @@ func TestTableKey_BedrockFieldsEncryptDecrypt(t *testing.T) { Provider: "bedrock", KeyID: "bedrock-uuid-1", Value: *schemas.NewEnvVar("bedrock-val"), - Aliases: schemas.KeyAliases{"model-a": "profile-a"}, + Aliases: schemas.KeyAliases{"model-a": {ModelID: "profile-a"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -224,7 +224,7 @@ func TestTableKey_BedrockFieldsEncryptDecrypt(t *testing.T) { assert.Equal(t, "us-west-2", found.BedrockKeyConfig.Region.GetValue()) require.NotNil(t, found.BedrockKeyConfig.ARN) assert.Equal(t, "arn:aws:iam::123456789:role/test", found.BedrockKeyConfig.ARN.GetValue()) - assert.Equal(t, "profile-a", found.Aliases["model-a"]) + assert.Equal(t, "profile-a", found.Aliases["model-a"].ModelID) require.NotNil(t, found.BedrockKeyConfig.BatchS3Config) require.Len(t, found.BedrockKeyConfig.BatchS3Config.Buckets, 1) assert.Equal(t, "my-batch-bucket", found.BedrockKeyConfig.BatchS3Config.Buckets[0].BucketName) @@ -1156,7 +1156,7 @@ func TestTableKey_AllProviderConfigs_EncryptDecrypt(t *testing.T) { Provider: "custom", KeyID: "multi-uuid", Value: *schemas.NewEnvVar("multi-api-key"), - Aliases: schemas.KeyAliases{"claude-3": "profile-claude"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "profile-claude"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://azure.endpoint.com"), ClientID: schemas.NewEnvVar("multi-azure-cid"), @@ -1230,7 +1230,7 @@ func TestTableKey_AllProviderConfigs_EncryptDecrypt(t *testing.T) { assert.Equal(t, "eu-west-1", found.BedrockKeyConfig.Region.GetValue()) require.NotNil(t, found.BedrockKeyConfig.ARN) assert.Equal(t, "arn:aws:bedrock:eu-west-1:123:role", found.BedrockKeyConfig.ARN.GetValue()) - assert.Equal(t, "profile-claude", found.Aliases["claude-3"]) + assert.Equal(t, "profile-claude", found.Aliases["claude-3"].ModelID) } // ============================================================================ @@ -1919,3 +1919,139 @@ func TestTableKey_VertexPlainValue_RoundTrip(t *testing.T) { assert.False(t, found.VertexKeyConfig.ProjectID.FromEnv) assert.Equal(t, "us-central1", found.VertexKeyConfig.Region.GetValue()) } + +// TestTableKey_AliasesJSON_LegacyWireShape verifies that a KeyAliases value +// containing only ModelID (the unenriched shape) is persisted to the DB as the +// legacy {"k":"v"} string-valued JSON, preserving byte-for-byte wire compat +// with pre-refactor consumers and keeping config_hash stable. +func TestTableKey_AliasesJSON_LegacyWireShape(t *testing.T) { + db := setupTestDB(t) + + key := &TableKey{ + Name: "openai-key", + ProviderID: 1, + Provider: "openai", + KeyID: "openai-uuid-aliases-legacy", + Value: *schemas.NewEnvVar("sk-test"), + Aliases: schemas.KeyAliases{ + "best-model": {ModelID: "gpt-4o-deployment"}, + "backup": {ModelID: "gpt-3.5-turbo"}, + }, + } + require.NoError(t, db.Create(key).Error) + + raw := rawRow(t, db, "config_keys", key.ID) + rawAliasesVal := raw["aliases_json"] + var rawAliasesStr string + switch v := rawAliasesVal.(type) { + case string: + rawAliasesStr = v + case []byte: + rawAliasesStr = string(v) + } + require.NotEmpty(t, rawAliasesStr) + + plaintext, err := encrypt.Decrypt(rawAliasesStr) + require.NoError(t, err, "aliases_json should be decryptable") + + // Both expected shapes are valid JSON encodings (map iteration order is not stable). + candidates := []string{ + `{"best-model":"gpt-4o-deployment","backup":"gpt-3.5-turbo"}`, + `{"backup":"gpt-3.5-turbo","best-model":"gpt-4o-deployment"}`, + } + assert.Contains(t, candidates, plaintext, "legacy ModelID-only aliases should marshal to the string-valued legacy wire shape") +} + +// TestTableKey_AliasesJSON_RichRoundTrip verifies that an enriched AliasConfig +// (with ModelName/ModelFamily/sub-config populated) survives the full DB +// encrypt → save → load → decrypt round-trip with no loss of information. +func TestTableKey_AliasesJSON_RichRoundTrip(t *testing.T) { + db := setupTestDB(t) + + apiVersion := "2024-08-01-preview" + canonical := "claude-3-5-sonnet" + family := schemas.ModelFamilyAnthropic + + key := &TableKey{ + Name: "azure-rich", + ProviderID: 1, + Provider: "azure", + KeyID: "azure-uuid-aliases-rich", + Value: *schemas.NewEnvVar("sk-test"), + Aliases: schemas.KeyAliases{ + "best-model": { + ModelID: "azure-deployment-xyz", + ModelName: &canonical, + ModelFamily: &family, + Description: "prod summarizer", + AzureAliasCfg: &schemas.AzureAliasCfg{ + APIVersion: &apiVersion, + }, + }, + "plain": {ModelID: "gpt-4o-fallback"}, + }, + AzureKeyConfig: &schemas.AzureKeyConfig{ + Endpoint: *schemas.NewEnvVar("https://example.openai.azure.com"), + }, + } + require.NoError(t, db.Create(key).Error) + + var found TableKey + require.NoError(t, db.First(&found, key.ID).Error) + require.NotNil(t, found.Aliases) + require.Len(t, found.Aliases, 2) + + rich := found.Aliases["best-model"] + assert.Equal(t, "azure-deployment-xyz", rich.ModelID) + require.NotNil(t, rich.ModelName) + assert.Equal(t, canonical, *rich.ModelName) + require.NotNil(t, rich.ModelFamily) + assert.Equal(t, schemas.ModelFamilyAnthropic, *rich.ModelFamily) + assert.Equal(t, "prod summarizer", rich.Description) + require.NotNil(t, rich.AzureAliasCfg) + require.NotNil(t, rich.AzureAliasCfg.APIVersion) + assert.Equal(t, apiVersion, *rich.AzureAliasCfg.APIVersion) + + // The unenriched sibling stays a legacy-shape entry — proves marshaling + // only escalates to the rich object form for entries that need it. + plain := found.Aliases["plain"] + assert.Equal(t, "gpt-4o-fallback", plain.ModelID) + assert.Nil(t, plain.ModelName) + assert.Nil(t, plain.ModelFamily) + assert.Nil(t, plain.AzureAliasCfg) +} + +// TestTableKey_AliasesJSON_LegacyInputRoundTrip simulates a row written before +// the refactor — raw legacy {"k":"v"} JSON in the aliases_json column — and +// verifies AfterFind promotes it to AliasConfig{ModelID: v} transparently. +func TestTableKey_AliasesJSON_LegacyInputRoundTrip(t *testing.T) { + db := setupTestDB(t) + + // First create a key without aliases so the row exists. + key := &TableKey{ + Name: "openai-key", + ProviderID: 1, + Provider: "openai", + KeyID: "openai-uuid-aliases-legacy-input", + Value: *schemas.NewEnvVar("sk-test"), + } + require.NoError(t, db.Create(key).Error) + + // Then write the legacy-shaped JSON directly into the aliases_json column, + // bypassing BeforeSave — this is what a pre-refactor row looks like. + legacy := `{"best-model":"gpt-4o-deployment"}` + encrypted, err := encrypt.Encrypt(legacy) + require.NoError(t, err) + require.NoError(t, db.Exec("UPDATE config_keys SET aliases_json = ? WHERE id = ?", encrypted, key.ID).Error) + + // Read back through GORM — AfterFind should decrypt + UnmarshalJSON should + // promote the legacy string value into AliasConfig{ModelID: ...}. + var found TableKey + require.NoError(t, db.First(&found, key.ID).Error) + require.NotNil(t, found.Aliases) + require.Len(t, found.Aliases, 1) + got := found.Aliases["best-model"] + assert.Equal(t, "gpt-4o-deployment", got.ModelID) + assert.Nil(t, got.ModelName) + assert.Nil(t, got.ModelFamily) +} diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go index e9429c823a..77d0692392 100644 --- a/transports/bifrost-http/lib/config_test.go +++ b/transports/bifrost-http/lib/config_test.go @@ -2695,7 +2695,7 @@ func TestGenerateKeyHash(t *testing.T) { Value: *schemas.NewEnvVar("sk-123"), Models: []string{"gpt-4", "gpt-3.5-turbo"}, Weight: 1.5, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, } hashWithAliases, err := configstore.GenerateKeyHash(keyWithAliases) @@ -5139,7 +5139,7 @@ func TestKeyHashComparison_AzureConfigSyncScenarios(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -5150,7 +5150,7 @@ func TestKeyHashComparison_AzureConfigSyncScenarios(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -5172,7 +5172,7 @@ func TestKeyHashComparison_AzureConfigSyncScenarios(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -5183,7 +5183,7 @@ func TestKeyHashComparison_AzureConfigSyncScenarios(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://different-azure.openai.azure.com"), // Changed! }, @@ -5205,7 +5205,7 @@ func TestKeyHashComparison_AzureConfigSyncScenarios(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -5216,7 +5216,7 @@ func TestKeyHashComparison_AzureConfigSyncScenarios(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment", "gpt-3.5-turbo": "gpt-35-turbo-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}, "gpt-3.5-turbo": {ModelID: "gpt-35-turbo-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -5304,7 +5304,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5317,7 +5317,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5341,7 +5341,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5354,7 +5354,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAI44QH8DHBEXAMPLE"), // Changed! SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5378,7 +5378,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5391,7 +5391,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("differentSecretKey/NEWKEY/bPxRfiCYEXAMPLEKEY"), // Changed! @@ -5415,7 +5415,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5428,7 +5428,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5452,7 +5452,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5466,7 +5466,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5491,7 +5491,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5504,7 +5504,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile", "claude-3.5": "claude-35-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}, "claude-3.5": {ModelID: "claude-35-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5594,7 +5594,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5608,7 +5608,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5634,7 +5634,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar(""), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar(""), // Empty for IAM role auth SecretKey: *schemas.NewEnvVar(""), // Empty for IAM role auth @@ -5648,7 +5648,7 @@ func TestKeyHashComparison_BedrockConfigSyncScenarios(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar(""), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "claude-3-inference-profile"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "claude-3-inference-profile"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5674,7 +5674,7 @@ func TestProviderHashComparison_AzureProviderFullLifecycle(t *testing.T) { Name: "azure-openai-key", Value: *schemas.NewEnvVar("azure-api-key-initial"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -5708,7 +5708,7 @@ func TestProviderHashComparison_AzureProviderFullLifecycle(t *testing.T) { Name: "azure-openai-key", Value: *schemas.NewEnvVar("azure-api-key-dashboard-edited"), // Changed via dashboard! Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -5740,7 +5740,7 @@ func TestProviderHashComparison_AzureProviderFullLifecycle(t *testing.T) { Name: "azure-openai-key", Value: *schemas.NewEnvVar("azure-api-key-initial"), // Original value from file Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -5776,7 +5776,7 @@ func TestProviderHashComparison_AzureProviderFullLifecycle(t *testing.T) { Name: "azure-openai-key", Value: *schemas.NewEnvVar("azure-api-key-initial"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment", "gpt-4o": "gpt-4o-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}, "gpt-4o": {ModelID: "gpt-4o-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://new-azure.openai.azure.com"), // Changed! }, @@ -5882,7 +5882,7 @@ func TestProviderHashComparison_BedrockProviderFullLifecycle(t *testing.T) { Name: "aws-bedrock-key", Value: *schemas.NewEnvVar(""), // Empty for Bedrock with IAM or AccessKey auth Weight: 1, - Aliases: schemas.KeyAliases{"claude-3-sonnet": "anthropic.claude-3-sonnet-20240229-v1:0"}, + Aliases: schemas.KeyAliases{"claude-3-sonnet": {ModelID: "anthropic.claude-3-sonnet-20240229-v1:0"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5917,7 +5917,7 @@ func TestProviderHashComparison_BedrockProviderFullLifecycle(t *testing.T) { Name: "aws-bedrock-key-eu", Value: *schemas.NewEnvVar(""), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3-sonnet": "anthropic.claude-3-sonnet-20240229-v1:0"}, + Aliases: schemas.KeyAliases{"claude-3-sonnet": {ModelID: "anthropic.claude-3-sonnet-20240229-v1:0"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAI44QH8DHBEXAMPLE"), SecretKey: *schemas.NewEnvVar("je7MtGbClwBF/2Zp9Utk/h3yCo8nvbEXAMPLEKEY"), @@ -5944,7 +5944,7 @@ func TestProviderHashComparison_BedrockProviderFullLifecycle(t *testing.T) { Name: "aws-bedrock-key", Value: *schemas.NewEnvVar(""), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3-sonnet": "anthropic.claude-3-sonnet-20240229-v1:0"}, + Aliases: schemas.KeyAliases{"claude-3-sonnet": {ModelID: "anthropic.claude-3-sonnet-20240229-v1:0"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -5983,7 +5983,7 @@ func TestProviderHashComparison_BedrockProviderFullLifecycle(t *testing.T) { Name: "aws-bedrock-key", Value: *schemas.NewEnvVar(""), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3-sonnet": "anthropic.claude-3-sonnet-20240229-v1:0", "claude-3-opus": "anthropic.claude-3-opus-20240229-v1:0"}, + Aliases: schemas.KeyAliases{"claude-3-sonnet": {ModelID: "anthropic.claude-3-sonnet-20240229-v1:0"}, "claude-3-opus": {ModelID: "anthropic.claude-3-opus-20240229-v1:0"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -6111,7 +6111,7 @@ func TestProviderHashComparison_BedrockProviderFullLifecycle(t *testing.T) { Name: "aws-bedrock-key", Value: *schemas.NewEnvVar(""), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3-sonnet": "anthropic.claude-3-sonnet-20240229-v1:0", "claude-3-opus": "anthropic.claude-3-opus-20240229-v1:0"}, + Aliases: schemas.KeyAliases{"claude-3-sonnet": {ModelID: "anthropic.claude-3-sonnet-20240229-v1:0"}, "claude-3-opus": {ModelID: "anthropic.claude-3-opus-20240229-v1:0"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -6154,7 +6154,7 @@ func TestProviderHashComparison_AzureNewProviderFromConfig(t *testing.T) { Name: "azure-openai-key", Value: *schemas.NewEnvVar("azure-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -6220,7 +6220,7 @@ func TestProviderHashComparison_BedrockNewProviderFromConfig(t *testing.T) { Name: "aws-bedrock-key", Value: *schemas.NewEnvVar(""), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "anthropic.claude-3-sonnet-20240229-v1:0"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "anthropic.claude-3-sonnet-20240229-v1:0"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -6289,7 +6289,7 @@ func TestProviderHashComparison_AzureDBValuePreservedWhenHashMatches(t *testing. Name: "azure-openai-key", Value: *schemas.NewEnvVar("DASHBOARD-EDITED-SECRET-KEY"), // Dashboard edited this! Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -6317,7 +6317,7 @@ func TestProviderHashComparison_AzureDBValuePreservedWhenHashMatches(t *testing. Name: "azure-openai-key", Value: *schemas.NewEnvVar("original-key-from-file"), // Different value than DB! Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), // Same }, @@ -6373,7 +6373,7 @@ func TestProviderHashComparison_BedrockDBValuePreservedWhenHashMatches(t *testin Name: "aws-bedrock-key", Value: *schemas.NewEnvVar(""), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "anthropic.claude-3-sonnet-20240229-v1:0"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "anthropic.claude-3-sonnet-20240229-v1:0"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("DASHBOARD-EDITED-ACCESS-KEY"), // Dashboard edited! SecretKey: *schemas.NewEnvVar("DASHBOARD-EDITED-SECRET-KEY"), // Dashboard edited! @@ -6403,7 +6403,7 @@ func TestProviderHashComparison_BedrockDBValuePreservedWhenHashMatches(t *testin Name: "aws-bedrock-key", Value: *schemas.NewEnvVar(""), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "anthropic.claude-3-sonnet-20240229-v1:0"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "anthropic.claude-3-sonnet-20240229-v1:0"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), // Different! SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), // Different! @@ -6491,7 +6491,7 @@ func TestProviderHashComparison_AzureConfigChangedInFile(t *testing.T) { Name: "azure-openai-key", Value: *schemas.NewEnvVar("azure-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4o": "gpt-4o-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4o": {ModelID: "gpt-4o-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://NEW-azure.openai.azure.com"), // Changed! }, @@ -6575,7 +6575,7 @@ func TestProviderHashComparison_BedrockConfigChangedInFile(t *testing.T) { Name: "aws-bedrock-key", Value: *schemas.NewEnvVar(""), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3-opus": "anthropic.claude-3-opus-20240229-v1:0"}, + Aliases: schemas.KeyAliases{"claude-3-opus": {ModelID: "anthropic.claude-3-opus-20240229-v1:0"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), @@ -14037,7 +14037,7 @@ func TestGenerateKeyHash_RuntimeVsMigrationParity(t *testing.T) { Value: *schemas.NewEnvVar("azure-key-value"), Weight: ptrFloat64(1.0), AzureKeyConfig: azureConfig, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, } schemaKey := schemas.Key{ @@ -14989,7 +14989,7 @@ func TestKeyHashComparison_VertexConfigSyncScenarios(t *testing.T) { Name: "vertex-key", Value: *schemas.NewEnvVar("vertex-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gemini-pro": "gemini-pro-endpoint"}, + Aliases: schemas.KeyAliases{"gemini-pro": {ModelID: "gemini-pro-endpoint"}}, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("my-project-123"), ProjectNumber: *schemas.NewEnvVar("123456789"), @@ -15003,7 +15003,7 @@ func TestKeyHashComparison_VertexConfigSyncScenarios(t *testing.T) { Name: "vertex-key", Value: *schemas.NewEnvVar("vertex-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gemini-pro": "gemini-pro-endpoint"}, + Aliases: schemas.KeyAliases{"gemini-pro": {ModelID: "gemini-pro-endpoint"}}, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("my-project-123"), ProjectNumber: *schemas.NewEnvVar("123456789"), @@ -15133,7 +15133,7 @@ func TestKeyHashComparison_VertexConfigSyncScenarios(t *testing.T) { Name: "vertex-key", Value: *schemas.NewEnvVar("vertex-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gemini-pro": "gemini-pro-endpoint"}, + Aliases: schemas.KeyAliases{"gemini-pro": {ModelID: "gemini-pro-endpoint"}}, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("my-project-123"), Region: *schemas.NewEnvVar("us-central1"), @@ -15145,7 +15145,7 @@ func TestKeyHashComparison_VertexConfigSyncScenarios(t *testing.T) { Name: "vertex-key", Value: *schemas.NewEnvVar("vertex-api-key-123"), Weight: 1, - Aliases: schemas.KeyAliases{"gemini-pro": "gemini-pro-endpoint", "gemini-1.5-pro": "gemini-15-pro-endpoint"}, + Aliases: schemas.KeyAliases{"gemini-pro": {ModelID: "gemini-pro-endpoint"}, "gemini-1.5-pro": {ModelID: "gemini-15-pro-endpoint"}}, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("my-project-123"), Region: *schemas.NewEnvVar("us-central1"), @@ -15532,7 +15532,7 @@ func TestKeyHashComparison_AzureDeploymentsChange(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -15543,7 +15543,7 @@ func TestKeyHashComparison_AzureDeploymentsChange(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment", "gpt-4o": "gpt-4o-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}, "gpt-4o": {ModelID: "gpt-4o-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -15563,7 +15563,7 @@ func TestKeyHashComparison_AzureDeploymentsChange(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment", "gpt-4o": "gpt-4o-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}, "gpt-4o": {ModelID: "gpt-4o-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -15574,7 +15574,7 @@ func TestKeyHashComparison_AzureDeploymentsChange(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -15594,7 +15594,7 @@ func TestKeyHashComparison_AzureDeploymentsChange(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment-v1"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment-v1"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -15605,7 +15605,7 @@ func TestKeyHashComparison_AzureDeploymentsChange(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment-v2"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment-v2"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -15635,7 +15635,7 @@ func TestKeyHashComparison_AzureDeploymentsChange(t *testing.T) { Name: "azure-key", Value: *schemas.NewEnvVar("azure-api-key"), Weight: 1, - Aliases: schemas.KeyAliases{"gpt-4": "gpt-4-deployment"}, + Aliases: schemas.KeyAliases{"gpt-4": {ModelID: "gpt-4-deployment"}}, AzureKeyConfig: &schemas.AzureKeyConfig{ Endpoint: *schemas.NewEnvVar("https://myazure.openai.azure.com"), }, @@ -15658,7 +15658,7 @@ func TestKeyHashComparison_BedrockDeploymentsChange(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-key"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "arn:aws:bedrock:us-east-1::inference-profile/claude-3"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "arn:aws:bedrock:us-east-1::inference-profile/claude-3"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI"), @@ -15671,7 +15671,7 @@ func TestKeyHashComparison_BedrockDeploymentsChange(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-key"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "arn:aws:bedrock:us-east-1::inference-profile/claude-3", "claude-3.5": "arn:aws:bedrock:us-east-1::inference-profile/claude-3.5"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "arn:aws:bedrock:us-east-1::inference-profile/claude-3"}, "claude-3.5": {ModelID: "arn:aws:bedrock:us-east-1::inference-profile/claude-3.5"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI"), @@ -15693,7 +15693,7 @@ func TestKeyHashComparison_BedrockDeploymentsChange(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-key"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "arn:aws:bedrock:us-east-1::inference-profile/claude-3", "claude-3.5": "arn:aws:bedrock:us-east-1::inference-profile/claude-3.5"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "arn:aws:bedrock:us-east-1::inference-profile/claude-3"}, "claude-3.5": {ModelID: "arn:aws:bedrock:us-east-1::inference-profile/claude-3.5"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI"), @@ -15706,7 +15706,7 @@ func TestKeyHashComparison_BedrockDeploymentsChange(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-key"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "arn:aws:bedrock:us-east-1::inference-profile/claude-3"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "arn:aws:bedrock:us-east-1::inference-profile/claude-3"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI"), @@ -15728,7 +15728,7 @@ func TestKeyHashComparison_BedrockDeploymentsChange(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-key"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "arn:aws:bedrock:us-east-1::inference-profile/claude-3-old"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "arn:aws:bedrock:us-east-1::inference-profile/claude-3-old"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI"), @@ -15741,7 +15741,7 @@ func TestKeyHashComparison_BedrockDeploymentsChange(t *testing.T) { Name: "bedrock-key", Value: *schemas.NewEnvVar("bedrock-key"), Weight: 1, - Aliases: schemas.KeyAliases{"claude-3": "arn:aws:bedrock:us-east-1::inference-profile/claude-3-new"}, + Aliases: schemas.KeyAliases{"claude-3": {ModelID: "arn:aws:bedrock:us-east-1::inference-profile/claude-3-new"}}, BedrockKeyConfig: &schemas.BedrockKeyConfig{ AccessKey: *schemas.NewEnvVar("AKIAIOSFODNN7EXAMPLE"), SecretKey: *schemas.NewEnvVar("wJalrXUtnFEMI"), @@ -15766,7 +15766,7 @@ func TestKeyHashComparison_VertexDeploymentsChange(t *testing.T) { Name: "vertex-key", Value: *schemas.NewEnvVar("vertex-creds"), Weight: 1, - Aliases: schemas.KeyAliases{"gemini-pro": "gemini-pro-endpoint"}, + Aliases: schemas.KeyAliases{"gemini-pro": {ModelID: "gemini-pro-endpoint"}}, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("my-project"), Region: *schemas.NewEnvVar("us-central1"), @@ -15778,7 +15778,7 @@ func TestKeyHashComparison_VertexDeploymentsChange(t *testing.T) { Name: "vertex-key", Value: *schemas.NewEnvVar("vertex-creds"), Weight: 1, - Aliases: schemas.KeyAliases{"gemini-pro": "gemini-pro-endpoint", "gemini-1.5-pro": "gemini-15-pro-endpoint"}, + Aliases: schemas.KeyAliases{"gemini-pro": {ModelID: "gemini-pro-endpoint"}, "gemini-1.5-pro": {ModelID: "gemini-15-pro-endpoint"}}, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("my-project"), Region: *schemas.NewEnvVar("us-central1"), @@ -15799,7 +15799,7 @@ func TestKeyHashComparison_VertexDeploymentsChange(t *testing.T) { Name: "vertex-key", Value: *schemas.NewEnvVar("vertex-creds"), Weight: 1, - Aliases: schemas.KeyAliases{"gemini-pro": "gemini-pro-endpoint", "gemini-1.5-pro": "gemini-15-pro-endpoint"}, + Aliases: schemas.KeyAliases{"gemini-pro": {ModelID: "gemini-pro-endpoint"}, "gemini-1.5-pro": {ModelID: "gemini-15-pro-endpoint"}}, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("my-project"), Region: *schemas.NewEnvVar("us-central1"), @@ -15811,7 +15811,7 @@ func TestKeyHashComparison_VertexDeploymentsChange(t *testing.T) { Name: "vertex-key", Value: *schemas.NewEnvVar("vertex-creds"), Weight: 1, - Aliases: schemas.KeyAliases{"gemini-pro": "gemini-pro-endpoint"}, + Aliases: schemas.KeyAliases{"gemini-pro": {ModelID: "gemini-pro-endpoint"}}, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("my-project"), Region: *schemas.NewEnvVar("us-central1"), @@ -15832,7 +15832,7 @@ func TestKeyHashComparison_VertexDeploymentsChange(t *testing.T) { Name: "vertex-key", Value: *schemas.NewEnvVar("vertex-creds"), Weight: 1, - Aliases: schemas.KeyAliases{"gemini-pro": "gemini-pro-endpoint-v1"}, + Aliases: schemas.KeyAliases{"gemini-pro": {ModelID: "gemini-pro-endpoint-v1"}}, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("my-project"), Region: *schemas.NewEnvVar("us-central1"), @@ -15844,7 +15844,7 @@ func TestKeyHashComparison_VertexDeploymentsChange(t *testing.T) { Name: "vertex-key", Value: *schemas.NewEnvVar("vertex-creds"), Weight: 1, - Aliases: schemas.KeyAliases{"gemini-pro": "gemini-pro-endpoint-v2"}, + Aliases: schemas.KeyAliases{"gemini-pro": {ModelID: "gemini-pro-endpoint-v2"}}, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("my-project"), Region: *schemas.NewEnvVar("us-central1"), @@ -15876,7 +15876,7 @@ func TestKeyHashComparison_VertexDeploymentsChange(t *testing.T) { Name: "vertex-key", Value: *schemas.NewEnvVar("vertex-creds"), Weight: 1, - Aliases: schemas.KeyAliases{"gemini-pro": "gemini-pro-endpoint"}, + Aliases: schemas.KeyAliases{"gemini-pro": {ModelID: "gemini-pro-endpoint"}}, VertexKeyConfig: &schemas.VertexKeyConfig{ ProjectID: *schemas.NewEnvVar("my-project"), Region: *schemas.NewEnvVar("us-central1"), diff --git a/transports/config.schema.json b/transports/config.schema.json index a28081efb9..f4d9a95a34 100644 --- a/transports/config.schema.json +++ b/transports/config.schema.json @@ -2398,13 +2398,74 @@ "aliases": { "type": "object", "additionalProperties": { - "type": "string", - "minLength": 1 + "oneOf": [ + { + "type": "string", + "minLength": 1, + "description": "Legacy shape: a bare provider-specific identifier. Equivalent to {\"model_id\": \"\"}." + }, + { + "type": "object", + "properties": { + "model_id": { + "type": "string", + "minLength": 1, + "description": "Provider-specific identifier sent on the wire (deployment name, inference profile ID, fine-tuned model ID, etc.)." + }, + "model_name": { + "type": "string", + "description": "Canonical model name used for pricing, logging, and family inference." + }, + "model_family": { + "type": "string", + "enum": ["anthropic", "openai", "mistral", "cohere", "gemini", "nova", "titan"], + "description": "Underlying model family. Used by provider routing without substring-sniffing the wire model ID." + }, + "description": { + "type": "string" + }, + "region": { + "type": "string", + "description": "Per-alias region override (can use env. prefix)." + }, + "api_version": { + "type": "string", + "description": "Azure OpenAI api-version override for this alias." + }, + "anthropic_version": { + "type": "string", + "description": "Azure anthropic-version header override for Claude-on-Azure deployments." + }, + "endpoint": { + "type": "string", + "description": "Per-alias Azure endpoint override (can use env. prefix)." + }, + "project_id": { + "type": "string", + "description": "Per-alias Vertex project ID override (can use env. prefix)." + }, + "project_number": { + "type": "string", + "description": "Per-alias Vertex project number override (can use env. prefix)." + }, + "inference_profile_arn": { + "type": "string", + "description": "Per-alias Bedrock inference profile ARN (can use env. prefix)." + }, + "use_deployments_endpoint": { + "type": "boolean", + "description": "Replicate: use the deployments endpoint instead of the predictions endpoint for this alias." + } + }, + "required": ["model_id"], + "additionalProperties": false + } + ] }, "propertyNames": { "minLength": 1 }, - "description": "Model alias mappings: maps a model name to a provider-specific identifier (deployment name, inference profile ID, fine-tuned model ID, etc.)" + "description": "Model alias mappings: each entry maps a user-facing model name to either a bare provider identifier (legacy string shape) or an AliasConfig object carrying the wire identifier plus optional canonical name, family, and provider-specific overrides." } }, "required": ["name", "weight"] From 84634a91abd00b66151cf0b585d9d6c448964d99 Mon Sep 17 00:00:00 2001 From: Pratham Mishra <99235987+Pratham-Mishra04@users.noreply.github.com> Date: Tue, 9 Jun 2026 15:56:48 +0530 Subject: [PATCH 008/108] feat: add per-alias Azure endpoint, API version, and Anthropic version overrides with context-aware `ResolveFamily` (#4181) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary This PR introduces per-alias Azure overrides for endpoint, API version, and Anthropic version, and adds a `ResolvedAlias` context value that providers can read to determine model family routing. Previously, model family detection relied solely on substring matching against the wire model ID, which broke when Azure deployment IDs were opaque strings (e.g., `12345-azure-deployment`) even if the alias key or config clearly indicated the model family (e.g., `best-claude`). This change fixes that by walking a precedence chain — explicit `ModelFamily` field → `ModelName` → `ModelID` → alias key — before falling back to substring matching. ## Changes - Introduced `ResolvedAlias` struct and `BifrostContextKeyResolvedAlias` context key; `bifrost.go` now stashes the full `AliasConfig` (and the user-facing alias key) into context after each key-level alias resolution, for both streaming and non-streaming paths. - Added `ResolveFamily`, `IsAnthropicModelFamily`, and `GetResolvedAlias` helpers in `schemas/account.go` that walk the alias precedence chain for model family detection. All Azure provider call sites that previously called `schemas.IsAnthropicModel(model)` now call `schemas.IsAnthropicModelFamily(ctx, model)` or `schemas.ResolveFamily(ctx, model)`. - Added `AzureAliasCfg` fields (`APIVersion`, `AnthropicVersion`, `Endpoint`) and three resolver helpers (`resolveAzureEndpoint`, `resolveAPIVersion`, `resolveAnthropicVersion`) in `core/providers/azure/utils.go`. All Azure provider methods now use these helpers instead of reading `key.AzureKeyConfig.Endpoint.GetValue()` directly, enabling per-alias endpoint and version overrides. - `buildPassthroughURL` now returns an error when the endpoint is empty and accepts a `*BifrostContext` to apply alias-level `api-version` overrides on passthrough routes. - `buildContainerURL` now accepts a `*BifrostContext` for the same reason. - `KeyAliases.UnmarshalJSON` switched from `encoding/json` to `sonic` for consistency with the rest of the codebase. - Added `KeyAliases.ResolveConfig` (returns `*AliasConfig`) used by `bifrost.go` to populate `ResolvedAlias`; the existing `Resolve` (returns string) is preserved for backward compatibility. ## Type of change - [ ] Bug fix - [x] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [ ] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./core/... ./core/providers/azure/... ./core/schemas/... ``` Key scenarios to validate: 1. **Opaque Azure deployment alias** — configure an alias like `best-claude → { model_id: "12345-deployment", azure_alias_cfg: {} }` with alias key containing "claude"; verify the request is routed through the Anthropic path (correct URL, `anthropic-version` header set). 2. **Alias-level endpoint override** — configure two aliases pointing to different Azure cognitive-services resources under the same key; verify each request hits the correct endpoint. 3. **Alias-level `api_version` override** — configure an alias with `azure_alias_cfg.api_version: "2024-10-21"`; verify the `api-version` query parameter on `/openai/deployments/` and `/openai/v1/responses` routes uses the override rather than the route default. 4. **Caller-supplied `api-version` wins** — pass `api-version` explicitly in a passthrough `rawQuery`; verify the alias override does not overwrite it. 5. **No alias matched** — verify existing substring-based family detection is unchanged. ## Breaking changes - [ ] Yes - [x] No `buildPassthroughURL` now returns `(string, error)` instead of `string`. This is an internal method on `AzureProvider` and is not part of any exported interface. ## Related issues ## Security considerations The `ResolvedAlias` value stored in `BifrostContext` is set exclusively by the core request worker and is documented as read-only for plugins. Alias-level endpoint overrides are resolved from `EnvVar` (supporting environment variable indirection), consistent with how key-level endpoints are handled, so secrets are not inlined in config. ## Checklist - [x] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable --- core/bifrost.go | 16 +- core/providers/azure/azure.go | 224 ++++++++++++------ .../providers/azure/azure_passthrough_test.go | 145 +++++++++++- core/providers/azure/utils.go | 42 ++++ core/schemas/account.go | 87 ++++++- core/schemas/account_test.go | 120 ++++++++++ core/schemas/bifrost.go | 1 + 7 files changed, 556 insertions(+), 79 deletions(-) diff --git a/core/bifrost.go b/core/bifrost.go index eaf296955b..d1fd698c2f 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -6177,7 +6177,13 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas // returned to the pool via its deferred finalizer. if IsStreamRequestType(req.RequestType) { stream, bifrostError = executeRequestWithRetries(req.Context, config, func(k schemas.Key) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { - resolvedModel = k.Aliases.Resolve(originalModelRequested) + if aliasConfig := k.Aliases.ResolveConfig(originalModelRequested); aliasConfig != nil { + resolvedModel = aliasConfig.ModelID + req.Context.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{Key: originalModelRequested, Config: aliasConfig}) + } else { + resolvedModel = originalModelRequested + req.Context.SetValue(schemas.BifrostContextKeyResolvedAlias, nil) + } req.SetModel(resolvedModel) // Snapshot per-attempt so postHookRunner doesn't observe a later retry's // alias while this attempt's provider goroutine is still emitting chunks. @@ -6236,7 +6242,13 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas }, keyProvider, req.RequestType, provider.GetProviderKey(), model, &req.BifrostRequest, bifrost.logger) } else { result, bifrostError = executeRequestWithRetries(req.Context, config, func(k schemas.Key) (*schemas.BifrostResponse, *schemas.BifrostError) { - resolvedModel = k.Aliases.Resolve(originalModelRequested) + if aliasConfig := k.Aliases.ResolveConfig(originalModelRequested); aliasConfig != nil { + resolvedModel = aliasConfig.ModelID + req.Context.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{Key: originalModelRequested, Config: aliasConfig}) + } else { + resolvedModel = originalModelRequested + req.Context.SetValue(schemas.BifrostContextKeyResolvedAlias, nil) + } req.SetModel(resolvedModel) return bifrost.handleProviderRequest(provider, config, req, k, keys) }, keyProvider, req.RequestType, provider.GetProviderKey(), model, &req.BifrostRequest, bifrost.logger) diff --git a/core/providers/azure/azure.go b/core/providers/azure/azure.go index c1e7781325..9bd1e2c80f 100644 --- a/core/providers/azure/azure.go +++ b/core/providers/azure/azure.go @@ -226,11 +226,10 @@ func (provider *AzureProvider) completeRequest( }() var url string - isAnthropicModel := schemas.IsAnthropicModel(model) // Set any extra headers from network config. // For Anthropic models, exclude anthropic-beta — it is merged and filtered explicitly below. - if isAnthropicModel { + if schemas.IsAnthropicModelFamily(ctx, model) { providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, []string{anthropic.AnthropicBetaHeader}) } else { providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) @@ -239,7 +238,7 @@ func (provider *AzureProvider) completeRequest( req.Header.SetContentType("application/json") // Get authentication headers - authHeaders, bifrostErr := provider.getAzureAuthHeaders(ctx, key, isAnthropicModel) + authHeaders, bifrostErr := provider.getAzureAuthHeaders(ctx, key, schemas.IsAnthropicModelFamily(ctx, model)) if bifrostErr != nil { return nil, 0, nil, bifrostErr } @@ -249,13 +248,13 @@ func (provider *AzureProvider) completeRequest( req.Header.Set(k, v) } - endpoint := key.AzureKeyConfig.Endpoint.GetValue() + endpoint := resolveAzureEndpoint(ctx, key) if endpoint == "" { return nil, 0, nil, providerUtils.NewConfigurationError("endpoint not set") } - if isAnthropicModel { - req.Header.Set("anthropic-version", AzureAnthropicAPIVersionDefault) + if schemas.IsAnthropicModelFamily(ctx, model) { + req.Header.Set("anthropic-version", resolveAnthropicVersion(ctx)) url = fmt.Sprintf("%s/%s", endpoint, path) // Merge ExtraHeaders + context anthropic-beta, filter for Azure, then set as HTTP header @@ -307,6 +306,11 @@ func (provider *AzureProvider) completeRequest( // listModelsByKey performs a list models request for a single key. // Returns the response and latency, or an error if the request fails. func (provider *AzureProvider) listModelsByKey(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostListModelsRequest) (*schemas.BifrostListModelsResponse, *schemas.BifrostError) { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + return nil, providerUtils.NewConfigurationError("endpoint not set") + } + // Create the request req := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() @@ -316,7 +320,7 @@ func (provider *AzureProvider) listModelsByKey(ctx *schemas.BifrostContext, key // Set any extra headers from network config providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) - req.SetRequestURI(key.AzureKeyConfig.Endpoint.GetValue() + providerUtils.GetPathFromContext(ctx, "/openai/v1/models")) + req.SetRequestURI(endpoint + providerUtils.GetPathFromContext(ctx, "/openai/v1/models")) req.Header.SetMethod(http.MethodGet) req.Header.SetContentType("application/json") @@ -460,7 +464,11 @@ func (provider *AzureProvider) TextCompletion(ctx *schemas.BifrostContext, key s // It formats the request, sends it to Azure, and processes the response. // Returns a channel of BifrostStreamChunk objects or an error if the request fails. func (provider *AzureProvider) TextCompletionStream(ctx *schemas.BifrostContext, postHookRunner schemas.PostHookRunner, postHookSpanFinalizer func(context.Context), key schemas.Key, request *schemas.BifrostTextCompletionRequest) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { - url := fmt.Sprintf("%s/openai/v1/completions", key.AzureKeyConfig.Endpoint.GetValue()) + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + return nil, providerUtils.NewConfigurationError("endpoint not set") + } + url := fmt.Sprintf("%s/openai/v1/completions", endpoint) // Get Azure authentication headers authHeader, err := provider.getAzureAuthHeaders(ctx, key, false) @@ -496,7 +504,7 @@ func (provider *AzureProvider) ChatCompletion(ctx *schemas.BifrostContext, key s ctx, request, func() (providerUtils.RequestBodyWithExtraParams, error) { - if schemas.IsAnthropicModel(request.Model) { + if schemas.ResolveFamily(ctx, request.Model) == schemas.ModelFamilyAnthropic { reqBody, err := anthropic.ToAnthropicChatRequest(ctx, request) if err != nil { return nil, err @@ -515,7 +523,7 @@ func (provider *AzureProvider) ChatCompletion(ctx *schemas.BifrostContext, key s } var path string - if schemas.IsAnthropicModel(request.Model) { + if schemas.ResolveFamily(ctx, request.Model) == schemas.ModelFamilyAnthropic { path = "anthropic/v1/messages" } else { path = "openai/v1/chat/completions" @@ -550,7 +558,7 @@ func (provider *AzureProvider) ChatCompletion(ctx *schemas.BifrostContext, key s var rawRequest interface{} var rawResponse interface{} - if schemas.IsAnthropicModel(request.Model) { + if schemas.ResolveFamily(ctx, request.Model) == schemas.ModelFamilyAnthropic { anthropicResponse := anthropic.AcquireAnthropicMessageResponse() defer anthropic.ReleaseAnthropicMessageResponse(anthropicResponse) rawRequest, rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, anthropicResponse, jsonData, providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse)) @@ -586,14 +594,18 @@ func (provider *AzureProvider) ChatCompletion(ctx *schemas.BifrostContext, key s // Uses Azure-specific URL construction with deployments and supports both api-key and Bearer token authentication. // Returns a channel containing BifrostResponse objects representing the stream or an error if the request fails. func (provider *AzureProvider) ChatCompletionStream(ctx *schemas.BifrostContext, postHookRunner schemas.PostHookRunner, postHookSpanFinalizer func(context.Context), key schemas.Key, request *schemas.BifrostChatRequest) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + return nil, providerUtils.NewConfigurationError("endpoint not set") + } var url string - if schemas.IsAnthropicModel(request.Model) { + if schemas.ResolveFamily(ctx, request.Model) == schemas.ModelFamilyAnthropic { authHeader, err := provider.getAzureAuthHeaders(ctx, key, true) if err != nil { return nil, err } - authHeader["anthropic-version"] = AzureAnthropicAPIVersionDefault - url = fmt.Sprintf("%s/anthropic/v1/messages", key.AzureKeyConfig.Endpoint.GetValue()) + authHeader["anthropic-version"] = resolveAnthropicVersion(ctx) + url = fmt.Sprintf("%s/anthropic/v1/messages", endpoint) jsonData, err := providerUtils.CheckContextAndGetRequestBody( ctx, @@ -637,7 +649,7 @@ func (provider *AzureProvider) ChatCompletionStream(ctx *schemas.BifrostContext, if err != nil { return nil, err } - url = fmt.Sprintf("%s/openai/v1/chat/completions", key.AzureKeyConfig.Endpoint.GetValue()) + url = fmt.Sprintf("%s/openai/v1/chat/completions", endpoint) // Use shared streaming logic from OpenAI return openai.HandleOpenAIChatCompletionStreaming( @@ -669,7 +681,7 @@ func (provider *AzureProvider) ChatCompletionStream(ctx *schemas.BifrostContext, func (provider *AzureProvider) Responses(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) { var jsonData []byte var bifrostErr *schemas.BifrostError - if schemas.IsAnthropicModel(request.Model) { + if schemas.IsAnthropicModelFamily(ctx, request.Model) { jsonData, bifrostErr = getRequestBodyForAnthropicResponses(ctx, request, request.Model, false, provider.sendBackRawRequest, provider.sendBackRawResponse) } else { jsonData, bifrostErr = providerUtils.CheckContextAndGetRequestBody( @@ -685,10 +697,10 @@ func (provider *AzureProvider) Responses(ctx *schemas.BifrostContext, key schema } var path string - if schemas.IsAnthropicModel(request.Model) { + if schemas.IsAnthropicModelFamily(ctx, request.Model) { path = "anthropic/v1/messages" } else { - path = fmt.Sprintf("openai/v1/responses?api-version=%s", AzureAPIVersionPreview) + path = fmt.Sprintf("openai/v1/responses?api-version=%s", resolveAPIVersion(ctx, AzureAPIVersionPreview)) } responseBody, latency, providerResponseHeaders, err := provider.completeRequest( @@ -720,7 +732,7 @@ func (provider *AzureProvider) Responses(ctx *schemas.BifrostContext, key schema var rawRequest interface{} var rawResponse interface{} - if schemas.IsAnthropicModel(request.Model) { + if schemas.IsAnthropicModelFamily(ctx, request.Model) { anthropicResponse := anthropic.AcquireAnthropicMessageResponse() defer anthropic.ReleaseAnthropicMessageResponse(anthropicResponse) rawRequest, rawResponse, bifrostErr = providerUtils.HandleProviderResponse(responseBody, anthropicResponse, jsonData, providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse)) @@ -753,14 +765,18 @@ func (provider *AzureProvider) Responses(ctx *schemas.BifrostContext, key schema // ResponsesStream performs a streaming responses request to Azure's API. func (provider *AzureProvider) ResponsesStream(ctx *schemas.BifrostContext, postHookRunner schemas.PostHookRunner, postHookSpanFinalizer func(context.Context), key schemas.Key, request *schemas.BifrostResponsesRequest) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + return nil, providerUtils.NewConfigurationError("endpoint not set") + } var url string - if schemas.IsAnthropicModel(request.Model) { + if schemas.IsAnthropicModelFamily(ctx, request.Model) { authHeader, err := provider.getAzureAuthHeaders(ctx, key, true) if err != nil { return nil, err } - authHeader["anthropic-version"] = AzureAnthropicAPIVersionDefault - url = fmt.Sprintf("%s/anthropic/v1/messages", key.AzureKeyConfig.Endpoint.GetValue()) + authHeader["anthropic-version"] = resolveAnthropicVersion(ctx) + url = fmt.Sprintf("%s/anthropic/v1/messages", endpoint) jsonData, bifrostErr := getRequestBodyForAnthropicResponses(ctx, request, request.Model, true, provider.sendBackRawRequest, provider.sendBackRawResponse) if bifrostErr != nil { @@ -790,7 +806,7 @@ func (provider *AzureProvider) ResponsesStream(ctx *schemas.BifrostContext, post if err != nil { return nil, err } - url = fmt.Sprintf("%s/openai/v1/responses?api-version=%s", key.AzureKeyConfig.Endpoint.GetValue(), AzureAPIVersionPreview) + url = fmt.Sprintf("%s/openai/v1/responses?api-version=%s", endpoint, resolveAPIVersion(ctx, AzureAPIVersionPreview)) // Use shared streaming logic from OpenAI return openai.HandleOpenAIResponsesStreaming( @@ -881,7 +897,7 @@ func (provider *AzureProvider) Embedding(ctx *schemas.BifrostContext, key schema // Speech is not supported by the Azure provider. func (provider *AzureProvider) Speech(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostSpeechRequest) (*schemas.BifrostSpeechResponse, *schemas.BifrostError) { - endpoint := key.AzureKeyConfig.Endpoint.GetValue() + endpoint := resolveAzureEndpoint(ctx, key) if endpoint == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -922,13 +938,18 @@ func (provider *AzureProvider) OCR(ctx *schemas.BifrostContext, key schemas.Key, // SpeechStream handles streaming for speech synthesis with Azure. // Azure sends raw binary audio bytes in SSE format, unlike OpenAI which sends JSON. func (provider *AzureProvider) SpeechStream(ctx *schemas.BifrostContext, postHookRunner schemas.PostHookRunner, postHookSpanFinalizer func(context.Context), key schemas.Key, request *schemas.BifrostSpeechRequest) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + return nil, providerUtils.NewConfigurationError("endpoint not set") + } + // Get Azure authentication headers authHeader, err := provider.getAzureAuthHeaders(ctx, key, false) if err != nil { return nil, err } - url := fmt.Sprintf("%s/openai/v1/audio/speech", key.AzureKeyConfig.Endpoint.GetValue()) + url := fmt.Sprintf("%s/openai/v1/audio/speech", endpoint) // Create HTTP request for streaming req := fasthttp.AcquireRequest() @@ -1209,7 +1230,11 @@ func (provider *AzureProvider) SpeechStream(ctx *schemas.BifrostContext, postHoo // Transcription is not supported by the Azure provider. func (provider *AzureProvider) Transcription(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostTranscriptionRequest) (*schemas.BifrostTranscriptionResponse, *schemas.BifrostError) { - url := fmt.Sprintf("%s/openai/deployments/%s/audio/transcriptions?api-version=%s", key.AzureKeyConfig.Endpoint.GetValue(), request.Model, DefaultAzureAPIVersion) + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + return nil, providerUtils.NewConfigurationError("endpoint not set") + } + url := fmt.Sprintf("%s/openai/deployments/%s/audio/transcriptions?api-version=%s", endpoint, request.Model, resolveAPIVersion(ctx, DefaultAzureAPIVersion)) response, err := openai.HandleOpenAITranscriptionRequest( ctx, @@ -1241,7 +1266,7 @@ func (provider *AzureProvider) TranscriptionStream(ctx *schemas.BifrostContext, // Returns a BifrostResponse containing the bifrost response or an error if the request fails. func (provider *AzureProvider) ImageGeneration(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostImageGenerationRequest) (*schemas.BifrostImageGenerationResponse, *schemas.BifrostError) { - endpoint := key.AzureKeyConfig.Endpoint.GetValue() + endpoint := resolveAzureEndpoint(ctx, key) if endpoint == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -1275,7 +1300,7 @@ func (provider *AzureProvider) ImageGenerationStream( key schemas.Key, request *schemas.BifrostImageGenerationRequest, ) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { - endpoint := key.AzureKeyConfig.Endpoint.GetValue() + endpoint := resolveAzureEndpoint(ctx, key) if endpoint == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -1311,7 +1336,7 @@ func (provider *AzureProvider) ImageGenerationStream( // ImageEdit performs an image edit request to Azure's API. func (provider *AzureProvider) ImageEdit(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostImageEditRequest) (*schemas.BifrostImageGenerationResponse, *schemas.BifrostError) { - endpoint := key.AzureKeyConfig.Endpoint.GetValue() + endpoint := resolveAzureEndpoint(ctx, key) if endpoint == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -1338,7 +1363,7 @@ func (provider *AzureProvider) ImageEdit(ctx *schemas.BifrostContext, key schema // ImageEditStream performs a streaming image edit request to Azure's API. func (provider *AzureProvider) ImageEditStream(ctx *schemas.BifrostContext, postHookRunner schemas.PostHookRunner, postHookSpanFinalizer func(context.Context), key schemas.Key, request *schemas.BifrostImageEditRequest) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { - endpoint := key.AzureKeyConfig.Endpoint.GetValue() + endpoint := resolveAzureEndpoint(ctx, key) if endpoint == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -1380,7 +1405,7 @@ func (provider *AzureProvider) ImageVariation(ctx *schemas.BifrostContext, key s // VideoGeneration creates a video using Azure's OpenAI-compatible Sora API. // This delegates to the OpenAI handler with Azure-specific URL and authentication. func (provider *AzureProvider) VideoGeneration(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostVideoGenerationRequest) (*schemas.BifrostVideoGenerationResponse, *schemas.BifrostError) { - endpoint := key.AzureKeyConfig.Endpoint.GetValue() + endpoint := resolveAzureEndpoint(ctx, key) if endpoint == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -1415,7 +1440,7 @@ func (provider *AzureProvider) VideoRetrieve(ctx *schemas.BifrostContext, key sc } videoID := providerUtils.StripVideoIDProviderSuffix(request.ID, providerName) - endpoint := key.AzureKeyConfig.Endpoint.GetValue() + endpoint := resolveAzureEndpoint(ctx, key) if endpoint == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -1450,7 +1475,7 @@ func (provider *AzureProvider) VideoDownload(ctx *schemas.BifrostContext, key sc } videoID := providerUtils.StripVideoIDProviderSuffix(request.ID, providerName) - endpoint := key.AzureKeyConfig.Endpoint.GetValue() + endpoint := resolveAzureEndpoint(ctx, key) if endpoint == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -1525,7 +1550,7 @@ func (provider *AzureProvider) VideoDelete(ctx *schemas.BifrostContext, key sche } videoID := providerUtils.StripVideoIDProviderSuffix(request.ID, providerName) - endpoint := key.AzureKeyConfig.Endpoint.GetValue() + endpoint := resolveAzureEndpoint(ctx, key) if endpoint == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -1554,7 +1579,7 @@ func (provider *AzureProvider) VideoDelete(ctx *schemas.BifrostContext, key sche // VideoList lists videos from Azure's OpenAI-compatible API. func (provider *AzureProvider) VideoList(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostVideoListRequest) (*schemas.BifrostVideoListResponse, *schemas.BifrostError) { - endpoint := key.AzureKeyConfig.Endpoint.GetValue() + endpoint := resolveAzureEndpoint(ctx, key) if endpoint == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -1588,6 +1613,10 @@ func (provider *AzureProvider) VideoRemix(_ *schemas.BifrostContext, _ schemas.K // FileUpload uploads a file to Azure OpenAI. func (provider *AzureProvider) FileUpload(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostFileUploadRequest) (*schemas.BifrostFileUploadResponse, *schemas.BifrostError) { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + return nil, providerUtils.NewConfigurationError("endpoint not set") + } if len(request.File) == 0 { return nil, providerUtils.NewBifrostOperationError("file content is required", nil) } @@ -1629,7 +1658,7 @@ func (provider *AzureProvider) FileUpload(ctx *schemas.BifrostContext, key schem defer fasthttp.ReleaseResponse(resp) // Build URL - requestURL := fmt.Sprintf("%s/openai/v1/files", key.AzureKeyConfig.Endpoint.GetValue()) + requestURL := fmt.Sprintf("%s/openai/v1/files", endpoint) // Set headers providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) @@ -1700,6 +1729,11 @@ func (provider *AzureProvider) FileList(ctx *schemas.BifrostContext, keys []sche }, nil } + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + return nil, providerUtils.NewConfigurationError("endpoint not set") + } + // Create request req := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() @@ -1707,7 +1741,7 @@ func (provider *AzureProvider) FileList(ctx *schemas.BifrostContext, keys []sche defer fasthttp.ReleaseResponse(resp) // Build URL with query params - requestURL := fmt.Sprintf("%s/openai/v1/files", key.AzureKeyConfig.Endpoint.GetValue()) + requestURL := fmt.Sprintf("%s/openai/v1/files", endpoint) values := url.Values{} if request.Purpose != "" { values.Set("purpose", string(request.Purpose)) @@ -1803,12 +1837,17 @@ func (provider *AzureProvider) FileRetrieve(ctx *schemas.BifrostContext, keys [] var lastErr *schemas.BifrostError for _, key := range keys { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + lastErr = providerUtils.NewConfigurationError("endpoint not set") + continue + } // Create request req := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() // Build URL - requestURL := fmt.Sprintf("%s/openai/v1/files/%s", key.AzureKeyConfig.Endpoint.GetValue(), url.PathEscape(request.FileID)) + requestURL := fmt.Sprintf("%s/openai/v1/files/%s", endpoint, url.PathEscape(request.FileID)) // Set headers providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) @@ -1887,12 +1926,17 @@ func (provider *AzureProvider) FileDelete(ctx *schemas.BifrostContext, keys []sc var lastErr *schemas.BifrostError for _, key := range keys { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + lastErr = providerUtils.NewConfigurationError("endpoint not set") + continue + } // Create request req := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() // Build URL - requestURL := fmt.Sprintf("%s/openai/v1/files/%s", key.AzureKeyConfig.Endpoint.GetValue(), url.PathEscape(request.FileID)) + requestURL := fmt.Sprintf("%s/openai/v1/files/%s", endpoint, url.PathEscape(request.FileID)) // Set headers providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) @@ -2000,12 +2044,17 @@ func (provider *AzureProvider) FileContent(ctx *schemas.BifrostContext, keys []s var lastErr *schemas.BifrostError for _, key := range keys { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + lastErr = providerUtils.NewConfigurationError("endpoint not set") + continue + } // Create request req := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() // Build URL - requestURL := fmt.Sprintf("%s/openai/v1/files/%s/content", key.AzureKeyConfig.Endpoint.GetValue(), url.PathEscape(request.FileID)) + requestURL := fmt.Sprintf("%s/openai/v1/files/%s/content", endpoint, url.PathEscape(request.FileID)) // Set headers providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) @@ -2075,6 +2124,10 @@ func (provider *AzureProvider) FileContent(ctx *schemas.BifrostContext, keys []s // BatchCreate creates a new batch job on Azure OpenAI. // Azure Batch API uses the same format as OpenAI but with Azure-specific URL patterns. func (provider *AzureProvider) BatchCreate(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostBatchCreateRequest) (*schemas.BifrostBatchCreateResponse, *schemas.BifrostError) { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + return nil, providerUtils.NewConfigurationError("endpoint not set") + } inputFileID := request.InputFileID // If no file_id provided but inline requests are available, upload them first @@ -2110,7 +2163,7 @@ func (provider *AzureProvider) BatchCreate(ctx *schemas.BifrostContext, key sche defer fasthttp.ReleaseResponse(resp) // Build URL - requestURL := fmt.Sprintf("%s/openai/v1/batches", key.AzureKeyConfig.Endpoint.GetValue()) + requestURL := fmt.Sprintf("%s/openai/v1/batches", endpoint) // Set headers providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) @@ -2209,6 +2262,11 @@ func (provider *AzureProvider) BatchList(ctx *schemas.BifrostContext, keys []sch }, nil } + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + return nil, providerUtils.NewConfigurationError("endpoint not set") + } + // Create request req := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() @@ -2216,7 +2274,7 @@ func (provider *AzureProvider) BatchList(ctx *schemas.BifrostContext, keys []sch defer fasthttp.ReleaseResponse(resp) // Build URL with query params - baseURL := fmt.Sprintf("%s/openai/v1/batches", key.AzureKeyConfig.Endpoint.GetValue()) + baseURL := fmt.Sprintf("%s/openai/v1/batches", endpoint) values := url.Values{} if request.Limit > 0 { values.Set("limit", fmt.Sprintf("%d", request.Limit)) @@ -2303,12 +2361,17 @@ func (provider *AzureProvider) BatchRetrieve(ctx *schemas.BifrostContext, keys [ var lastErr *schemas.BifrostError for _, key := range keys { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + lastErr = providerUtils.NewConfigurationError("endpoint not set") + continue + } // Create request req := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() // Build URL - requestURL := fmt.Sprintf("%s/openai/v1/batches/%s", key.AzureKeyConfig.Endpoint.GetValue(), url.PathEscape(request.BatchID)) + requestURL := fmt.Sprintf("%s/openai/v1/batches/%s", endpoint, url.PathEscape(request.BatchID)) // Set headers providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) @@ -2388,12 +2451,17 @@ func (provider *AzureProvider) BatchCancel(ctx *schemas.BifrostContext, keys []s var lastErr *schemas.BifrostError for _, key := range keys { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + lastErr = providerUtils.NewConfigurationError("endpoint not set") + continue + } // Create request req := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() // Build URL - requestURL := fmt.Sprintf("%s/openai/v1/batches/%s/cancel", key.AzureKeyConfig.Endpoint.GetValue(), url.PathEscape(request.BatchID)) + requestURL := fmt.Sprintf("%s/openai/v1/batches/%s/cancel", endpoint, url.PathEscape(request.BatchID)) // Set headers providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) @@ -2729,8 +2797,9 @@ func (provider *AzureProvider) Compaction(ctx *schemas.BifrostContext, key schem // buildContainerURL constructs the Azure container API URL. // Container endpoints are not per-deployment, so they use the openai/v1 prefix directly. -func (provider *AzureProvider) buildContainerURL(key schemas.Key, path string) string { - endpoint := strings.TrimRight(key.AzureKeyConfig.Endpoint.GetValue(), "/") +// ctx carries the resolved alias so per-alias Endpoint overrides are honored. +func (provider *AzureProvider) buildContainerURL(ctx *schemas.BifrostContext, key schemas.Key, path string) string { + endpoint := strings.TrimRight(resolveAzureEndpoint(ctx, key), "/") return fmt.Sprintf("%s/openai/v1%s", endpoint, path) } @@ -2742,7 +2811,7 @@ func (provider *AzureProvider) ContainerCreate(ctx *schemas.BifrostContext, key if request.Name == "" { return nil, providerUtils.NewBifrostOperationError("invalid request: name is required", nil) } - if key.AzureKeyConfig == nil || key.AzureKeyConfig.Endpoint.GetValue() == "" { + if resolveAzureEndpoint(ctx, key) == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -2779,7 +2848,7 @@ func (provider *AzureProvider) ContainerCreate(ctx *schemas.BifrostContext, key defer fasthttp.ReleaseResponse(resp) providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) - req.SetRequestURI(provider.buildContainerURL(key, "/containers")) + req.SetRequestURI(provider.buildContainerURL(ctx, key, "/containers")) req.Header.SetMethod(http.MethodPost) req.Header.SetContentType("application/json") req.SetBody(jsonBody) @@ -2863,7 +2932,7 @@ func (provider *AzureProvider) ContainerRetrieve(ctx *schemas.BifrostContext, ke var lastErr *schemas.BifrostError for _, key := range keys { - if key.AzureKeyConfig == nil || key.AzureKeyConfig.Endpoint.GetValue() == "" { + if resolveAzureEndpoint(ctx, key) == "" { lastErr = providerUtils.NewConfigurationError("endpoint not set") continue } @@ -2872,7 +2941,7 @@ func (provider *AzureProvider) ContainerRetrieve(ctx *schemas.BifrostContext, ke resp := fasthttp.AcquireResponse() providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) - req.SetRequestURI(provider.buildContainerURL(key, "/containers/"+url.PathEscape(request.ContainerID))) + req.SetRequestURI(provider.buildContainerURL(ctx, key, "/containers/"+url.PathEscape(request.ContainerID))) req.Header.SetMethod(http.MethodGet) authHeaders, bifrostErr := provider.getAzureAuthHeaders(ctx, key, false) @@ -2969,7 +3038,7 @@ func (provider *AzureProvider) ContainerDelete(ctx *schemas.BifrostContext, keys var lastErr *schemas.BifrostError for _, key := range keys { - if key.AzureKeyConfig == nil || key.AzureKeyConfig.Endpoint.GetValue() == "" { + if resolveAzureEndpoint(ctx, key) == "" { lastErr = providerUtils.NewConfigurationError("endpoint not set") continue } @@ -2978,7 +3047,7 @@ func (provider *AzureProvider) ContainerDelete(ctx *schemas.BifrostContext, keys resp := fasthttp.AcquireResponse() providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) - req.SetRequestURI(provider.buildContainerURL(key, "/containers/"+url.PathEscape(request.ContainerID))) + req.SetRequestURI(provider.buildContainerURL(ctx, key, "/containers/"+url.PathEscape(request.ContainerID))) req.Header.SetMethod(http.MethodDelete) req.Header.SetContentType("application/json") @@ -3061,7 +3130,7 @@ func (provider *AzureProvider) ContainerFileCreate(ctx *schemas.BifrostContext, if len(request.File) == 0 { return nil, providerUtils.NewBifrostOperationError("invalid request: file is required", nil) } - if key.AzureKeyConfig == nil || key.AzureKeyConfig.Endpoint.GetValue() == "" { + if resolveAzureEndpoint(ctx, key) == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } @@ -3084,7 +3153,7 @@ func (provider *AzureProvider) ContainerFileCreate(ctx *schemas.BifrostContext, defer fasthttp.ReleaseResponse(resp) providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) - req.SetRequestURI(provider.buildContainerURL(key, fmt.Sprintf("/containers/%s/files", url.PathEscape(request.ContainerID)))) + req.SetRequestURI(provider.buildContainerURL(ctx, key, fmt.Sprintf("/containers/%s/files", url.PathEscape(request.ContainerID)))) req.Header.SetMethod(http.MethodPost) req.Header.Set("Content-Type", writer.FormDataContentType()) req.SetBody(body.Bytes()) @@ -3169,11 +3238,11 @@ func (provider *AzureProvider) ContainerFileList(ctx *schemas.BifrostContext, ke if !ok { return &schemas.BifrostContainerFileListResponse{Object: "list", Data: []schemas.ContainerFileObject{}, HasMore: false}, nil } - if key.AzureKeyConfig == nil || key.AzureKeyConfig.Endpoint.GetValue() == "" { + if resolveAzureEndpoint(ctx, key) == "" { return nil, providerUtils.NewConfigurationError("endpoint not set") } - requestURL := provider.buildContainerURL(key, fmt.Sprintf("/containers/%s/files", url.PathEscape(request.ContainerID))) + requestURL := provider.buildContainerURL(ctx, key, fmt.Sprintf("/containers/%s/files", url.PathEscape(request.ContainerID))) queryParams := url.Values{} if request.Limit > 0 { queryParams.Set("limit", fmt.Sprintf("%d", request.Limit)) @@ -3275,7 +3344,7 @@ func (provider *AzureProvider) ContainerFileRetrieve(ctx *schemas.BifrostContext var lastErr *schemas.BifrostError for _, key := range keys { - if key.AzureKeyConfig == nil || key.AzureKeyConfig.Endpoint.GetValue() == "" { + if resolveAzureEndpoint(ctx, key) == "" { lastErr = providerUtils.NewConfigurationError("endpoint not set") continue } @@ -3284,7 +3353,7 @@ func (provider *AzureProvider) ContainerFileRetrieve(ctx *schemas.BifrostContext resp := fasthttp.AcquireResponse() providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) - req.SetRequestURI(provider.buildContainerURL(key, fmt.Sprintf("/containers/%s/files/%s", url.PathEscape(request.ContainerID), url.PathEscape(request.FileID)))) + req.SetRequestURI(provider.buildContainerURL(ctx, key, fmt.Sprintf("/containers/%s/files/%s", url.PathEscape(request.ContainerID), url.PathEscape(request.FileID)))) req.Header.SetMethod(http.MethodGet) authHeaders, bifrostErr := provider.getAzureAuthHeaders(ctx, key, false) @@ -3380,7 +3449,7 @@ func (provider *AzureProvider) ContainerFileContent(ctx *schemas.BifrostContext, var lastErr *schemas.BifrostError for _, key := range keys { - if key.AzureKeyConfig == nil || key.AzureKeyConfig.Endpoint.GetValue() == "" { + if resolveAzureEndpoint(ctx, key) == "" { lastErr = providerUtils.NewConfigurationError("endpoint not set") continue } @@ -3389,7 +3458,7 @@ func (provider *AzureProvider) ContainerFileContent(ctx *schemas.BifrostContext, resp := fasthttp.AcquireResponse() providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) - req.SetRequestURI(provider.buildContainerURL(key, fmt.Sprintf("/containers/%s/files/%s/content", url.PathEscape(request.ContainerID), url.PathEscape(request.FileID)))) + req.SetRequestURI(provider.buildContainerURL(ctx, key, fmt.Sprintf("/containers/%s/files/%s/content", url.PathEscape(request.ContainerID), url.PathEscape(request.FileID)))) req.Header.SetMethod(http.MethodGet) authHeaders, bifrostErr := provider.getAzureAuthHeaders(ctx, key, false) @@ -3470,7 +3539,7 @@ func (provider *AzureProvider) ContainerFileDelete(ctx *schemas.BifrostContext, var lastErr *schemas.BifrostError for _, key := range keys { - if key.AzureKeyConfig == nil || key.AzureKeyConfig.Endpoint.GetValue() == "" { + if resolveAzureEndpoint(ctx, key) == "" { lastErr = providerUtils.NewConfigurationError("endpoint not set") continue } @@ -3479,7 +3548,7 @@ func (provider *AzureProvider) ContainerFileDelete(ctx *schemas.BifrostContext, resp := fasthttp.AcquireResponse() providerUtils.SetExtraHeaders(ctx, req, provider.networkConfig.ExtraHeaders, nil) - req.SetRequestURI(provider.buildContainerURL(key, fmt.Sprintf("/containers/%s/files/%s", url.PathEscape(request.ContainerID), url.PathEscape(request.FileID)))) + req.SetRequestURI(provider.buildContainerURL(ctx, key, fmt.Sprintf("/containers/%s/files/%s", url.PathEscape(request.ContainerID), url.PathEscape(request.FileID)))) req.Header.SetMethod(http.MethodDelete) authHeaders, bifrostErr := provider.getAzureAuthHeaders(ctx, key, false) @@ -3556,7 +3625,10 @@ func (provider *AzureProvider) Passthrough( key schemas.Key, req *schemas.BifrostPassthroughRequest, ) (*schemas.BifrostPassthroughResponse, *schemas.BifrostError) { - url := provider.buildPassthroughURL(key, req.Path, req.RawQuery) + url, err := provider.buildPassthroughURL(ctx, key, req.Path, req.RawQuery) + if err != nil { + return nil, providerUtils.NewConfigurationError(fmt.Sprintf("failed to build passthrough URL: %s", err.Error())) + } fasthttpReq := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() @@ -3572,7 +3644,7 @@ func (provider *AzureProvider) Passthrough( fasthttpReq.Header.Set(k, v) } - authHeaders, bifrostErr := provider.getAzureAuthHeaders(ctx, key, schemas.IsAnthropicModel(req.Model)) + authHeaders, bifrostErr := provider.getAzureAuthHeaders(ctx, key, schemas.IsAnthropicModelFamily(ctx, req.Model)) if bifrostErr != nil { return nil, bifrostErr } @@ -3625,7 +3697,10 @@ func (provider *AzureProvider) PassthroughStream( key schemas.Key, req *schemas.BifrostPassthroughRequest, ) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { - url := provider.buildPassthroughURL(key, req.Path, req.RawQuery) + url, err := provider.buildPassthroughURL(ctx, key, req.Path, req.RawQuery) + if err != nil { + return nil, providerUtils.NewConfigurationError(fmt.Sprintf("failed to build passthrough URL: %s", err.Error())) + } fasthttpReq := fasthttp.AcquireRequest() resp := fasthttp.AcquireResponse() @@ -3643,7 +3718,7 @@ func (provider *AzureProvider) PassthroughStream( fasthttpReq.Header.Set("Connection", "close") - authHeaders, bifrostErr := provider.getAzureAuthHeaders(ctx, key, schemas.IsAnthropicModel(req.Model)) + authHeaders, bifrostErr := provider.getAzureAuthHeaders(ctx, key, schemas.IsAnthropicModelFamily(ctx, req.Model)) if bifrostErr != nil { return nil, bifrostErr } @@ -3686,7 +3761,7 @@ func (provider *AzureProvider) PassthroughStream( } var anthropicUsage *anthropic.AnthropicPassthroughStreamUsage - if schemas.IsAnthropicModel(req.Model) { + if schemas.IsAnthropicModelFamily(ctx, req.Model) { anthropicUsage = &anthropic.AnthropicPassthroughStreamUsage{} } return providerUtils.StreamPassthrough( @@ -3716,8 +3791,13 @@ func (provider *AzureProvider) PassthroughStream( } // buildPassthroughURL constructs the full Azure URL for a passthrough request. -func (provider *AzureProvider) buildPassthroughURL(key schemas.Key, path, rawQuery string) string { - endpoint := key.AzureKeyConfig.Endpoint.GetValue() +// ctx carries the resolved alias used to pick a per-alias api-version override +// when the caller did not supply one in rawQuery. +func (provider *AzureProvider) buildPassthroughURL(ctx *schemas.BifrostContext, key schemas.Key, path, rawQuery string) (string, error) { + endpoint := resolveAzureEndpoint(ctx, key) + if endpoint == "" { + return "", fmt.Errorf("endpoint not set") + } // Normalise paths emitted by the Azure SDK. path = strings.Replace(path, "/openai/responses", "/openai/v1/responses", 1) @@ -3734,14 +3814,14 @@ func (provider *AzureProvider) buildPassthroughURL(key schemas.Key, path, rawQue // Responses API requires api-version=preview. values, _ := url.ParseQuery(rawQuery) if values.Get("api-version") == "" { - values.Set("api-version", AzureAPIVersionPreview) + values.Set("api-version", resolveAPIVersion(ctx, AzureAPIVersionPreview)) rawQuery = values.Encode() } case strings.Contains(path, "/openai/deployments/"): // Classic /deployments/ routes require api-version. Inject a default if absent. values, _ := url.ParseQuery(rawQuery) if values.Get("api-version") == "" { - values.Set("api-version", DefaultAzureAPIVersion) + values.Set("api-version", resolveAPIVersion(ctx, DefaultAzureAPIVersion)) rawQuery = values.Encode() } } @@ -3750,7 +3830,7 @@ func (provider *AzureProvider) buildPassthroughURL(key schemas.Key, path, rawQue if rawQuery != "" { fullURL += "?" + rawQuery } - return fullURL + return fullURL, nil } // extractAzurePassthroughUsage dispatches usage extraction by the upstream API the diff --git a/core/providers/azure/azure_passthrough_test.go b/core/providers/azure/azure_passthrough_test.go index b070c49f9f..684c02dfd1 100644 --- a/core/providers/azure/azure_passthrough_test.go +++ b/core/providers/azure/azure_passthrough_test.go @@ -119,10 +119,153 @@ func TestBuildPassthroughURL(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { t.Parallel() - got := provider.buildPassthroughURL(makeKey(endpoint), tt.path, tt.rawQuery) + got, _ := provider.buildPassthroughURL(nil, makeKey(endpoint), tt.path, tt.rawQuery) if got != tt.want { t.Errorf("\ngot: %s\nwant: %s", got, tt.want) } }) } } + +// TestBuildPassthroughURL_AliasAPIVersionOverride verifies that when the +// resolved alias carries an AzureAliasCfg.APIVersion override, it takes +// precedence over the route default (DefaultAzureAPIVersion for /deployments/, +// AzureAPIVersionPreview for /openai/v1/responses) — only in the path where +// the caller did NOT supply api-version themselves. Caller-supplied wins over +// alias override; alias override wins over route default. +func TestBuildPassthroughURL_AliasAPIVersionOverride(t *testing.T) { + t.Parallel() + + provider := &AzureProvider{} + endpoint := "https://my-resource.openai.azure.com" + makeKey := schemas.Key{ + AzureKeyConfig: &schemas.AzureKeyConfig{ + Endpoint: *schemas.NewEnvVar(endpoint), + }, + } + + // Build a ctx carrying an alias with APIVersion override. + overrideVer := "2024-10-21" + ctx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "best-model", + Config: &schemas.AliasConfig{ + ModelID: "gpt-4o-deployment", + AzureAliasCfg: &schemas.AzureAliasCfg{ + APIVersion: &overrideVer, + }, + }, + }) + + t.Run("deployments route: alias APIVersion overrides default", func(t *testing.T) { + got, _ := provider.buildPassthroughURL(ctx, makeKey, "/openai/deployments/gpt-4o/chat/completions", "") + want := endpoint + "/openai/deployments/gpt-4o/chat/completions?api-version=" + overrideVer + if got != want { + t.Errorf("\ngot: %s\nwant: %s", got, want) + } + }) + + t.Run("responses route: alias APIVersion overrides preview default", func(t *testing.T) { + got, _ := provider.buildPassthroughURL(ctx, makeKey, "/openai/v1/responses", "") + want := endpoint + "/openai/v1/responses?api-version=" + overrideVer + if got != want { + t.Errorf("\ngot: %s\nwant: %s", got, want) + } + }) + + t.Run("caller-supplied api-version wins over alias override", func(t *testing.T) { + got, _ := provider.buildPassthroughURL(ctx, makeKey, "/openai/deployments/gpt-4o/chat/completions", "api-version=2023-01-01") + want := endpoint + "/openai/deployments/gpt-4o/chat/completions?api-version=2023-01-01" + if got != want { + t.Errorf("\ngot: %s\nwant: %s", got, want) + } + }) +} + +// TestResolveAPIVersion_NoAlias verifies the helper returns the route default +// when no resolved alias is in ctx (covers the legacy code path). +func TestResolveAPIVersion_NoAlias(t *testing.T) { + if got := resolveAPIVersion(nil, DefaultAzureAPIVersion); got != DefaultAzureAPIVersion { + t.Errorf("got %q, want %q", got, DefaultAzureAPIVersion) + } + ctx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + if got := resolveAPIVersion(ctx, AzureAPIVersionPreview); got != AzureAPIVersionPreview { + t.Errorf("got %q, want %q", got, AzureAPIVersionPreview) + } +} + +// TestResolveAzureEndpoint_AliasOverride verifies the Endpoint override path. +// Lets one Azure credential cover deployments hosted on multiple cognitive- +// services resources. +func TestResolveAzureEndpoint_AliasOverride(t *testing.T) { + keyEndpoint := "https://primary.openai.azure.com" + aliasEndpoint := "https://anthropic-resource.openai.azure.com" + key := schemas.Key{ + AzureKeyConfig: &schemas.AzureKeyConfig{ + Endpoint: *schemas.NewEnvVar(keyEndpoint), + }, + } + + // No alias: falls back to key-level endpoint. + if got := resolveAzureEndpoint(nil, key); got != keyEndpoint { + t.Errorf("nil ctx: got %q, want key-level %q", got, keyEndpoint) + } + ctx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + if got := resolveAzureEndpoint(ctx, key); got != keyEndpoint { + t.Errorf("empty ctx: got %q, want key-level %q", got, keyEndpoint) + } + + // With alias-level Endpoint override, alias wins. + override := schemas.NewEnvVar(aliasEndpoint) + ctx.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "best-claude", + Config: &schemas.AliasConfig{ + ModelID: "claude-deployment", + AzureAliasCfg: &schemas.AzureAliasCfg{ + Endpoint: override, + }, + }, + }) + if got := resolveAzureEndpoint(ctx, key); got != aliasEndpoint { + t.Errorf("alias override: got %q, want %q", got, aliasEndpoint) + } + + // Alias with empty Endpoint value falls through to key-level — guards against + // a misconfigured alias accidentally erasing the endpoint. + emptyOverride := schemas.NewEnvVar("") + ctx2 := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx2.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "x", + Config: &schemas.AliasConfig{ + ModelID: "x", + AzureAliasCfg: &schemas.AzureAliasCfg{ + Endpoint: emptyOverride, + }, + }, + }) + if got := resolveAzureEndpoint(ctx2, key); got != keyEndpoint { + t.Errorf("empty alias endpoint should fall through: got %q, want %q", got, keyEndpoint) + } +} + +// TestResolveAnthropicVersion_AliasOverride verifies the AnthropicVersion +// override path mirrors the APIVersion behavior. +func TestResolveAnthropicVersion_AliasOverride(t *testing.T) { + if got := resolveAnthropicVersion(nil); got != AzureAnthropicAPIVersionDefault { + t.Errorf("nil ctx: got %q, want default", got) + } + override := "2024-10-22" + ctx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "best-claude", + Config: &schemas.AliasConfig{ + ModelID: "claude-deployment", + AzureAliasCfg: &schemas.AzureAliasCfg{ + AnthropicVersion: &override, + }, + }, + }) + if got := resolveAnthropicVersion(ctx); got != override { + t.Errorf("got %q, want %q", got, override) + } +} diff --git a/core/providers/azure/utils.go b/core/providers/azure/utils.go index e1ef5d02f1..2aecf5e04f 100644 --- a/core/providers/azure/utils.go +++ b/core/providers/azure/utils.go @@ -37,3 +37,45 @@ func getAzureScopes(configuredScopes []string) []string { } return scopes } + +// resolveAnthropicVersion returns the anthropic-version header value for the +// current attempt. Uses the AzureAliasCfg.AnthropicVersion override from the +// resolved alias when present, otherwise the Azure default. +func resolveAnthropicVersion(ctx *schemas.BifrostContext) string { + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.AzureAliasCfg != nil && ra.Config.AzureAliasCfg.AnthropicVersion != nil && *ra.Config.AzureAliasCfg.AnthropicVersion != "" { + return *ra.Config.AzureAliasCfg.AnthropicVersion + } + return AzureAnthropicAPIVersionDefault +} + +// resolveAPIVersion returns the Azure api-version query parameter value for +// the current attempt. Uses the AzureAliasCfg.APIVersion override from the +// resolved alias when present, otherwise the provided default. Different +// Azure routes have different defaults (DefaultAzureAPIVersion for classic +// /openai/deployments/, AzureAPIVersionPreview for /openai/v1/responses); +// callers pass the route's default so the override can take precedence +// without losing the route-specific fallback. +func resolveAPIVersion(ctx *schemas.BifrostContext, defaultVersion string) string { + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.AzureAliasCfg != nil && ra.Config.AzureAliasCfg.APIVersion != nil && *ra.Config.AzureAliasCfg.APIVersion != "" { + return *ra.Config.AzureAliasCfg.APIVersion + } + return defaultVersion +} + +// resolveAzureEndpoint returns the Azure cognitive-services endpoint URL for +// the current attempt. Uses the AzureAliasCfg.Endpoint override from the +// resolved alias when present, otherwise the key-level endpoint. Lets one +// Azure credential (ClientID/Secret/TenantID or API key) span deployments +// hosted on different Azure resources (e.g. OpenAI on east-us, Anthropic on +// west-us2). +func resolveAzureEndpoint(ctx *schemas.BifrostContext, key schemas.Key) string { + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.AzureAliasCfg != nil && ra.Config.AzureAliasCfg.Endpoint != nil { + if v := ra.Config.AzureAliasCfg.Endpoint.GetValue(); v != "" { + return v + } + } + if key.AzureKeyConfig != nil { + return key.AzureKeyConfig.Endpoint.GetValue() + } + return "" +} diff --git a/core/schemas/account.go b/core/schemas/account.go index 38d8083728..6a6e1bb519 100644 --- a/core/schemas/account.go +++ b/core/schemas/account.go @@ -8,6 +8,8 @@ import ( "fmt" "slices" "strings" + + "github.com/bytedance/sonic" ) type KeyStatusType string @@ -176,7 +178,7 @@ func (mf *ModelFamily) IsValid() bool { type AzureAliasCfg struct { APIVersion *string `json:"api_version,omitempty"` // overrides the Azure OpenAI api-version query param for this alias AnthropicVersion *string `json:"anthropic_version,omitempty"` // overrides the anthropic-version header for Claude-on-Azure deployments - Endpoint *EnvVar `json:"endpoint,omitempty"` // overrides AzureKeyConfig.Endpoint for this alias (allows one credential to span multiple Azure resources) + Endpoint *EnvVar `json:"endpoint,omitempty"` // overrides AzureKeyConfig.Endpoint for this alias — lets one credential span deployments on multiple Azure resources } // VertexAliasCfg holds Vertex-specific overrides that apply to a single alias. @@ -297,6 +299,83 @@ func (ka KeyAliases) Resolve(model string) string { return model } +// ResolvedAlias is what core stashes in BifrostContext after key-level alias +// resolution. Key is the user-facing model name the client sent (LHS of the +// alias map). Config is the matched AliasConfig. +// +// Carrying the alias key alongside the config lets providers consult it as +// the lowest-precedence tier for family detection — common case: an admin +// names their alias "best-claude" but the wire ModelID is an opaque Azure +// deployment ID, so neither the config fields nor request.Model carry the +// "claude" substring; the alias key does. +type ResolvedAlias struct { + Key string + Config *AliasConfig +} + +// GetResolvedAlias returns the ResolvedAlias that core stashed in ctx after +// key-level alias resolution, or nil if no alias matched or ctx is nil. +// +// This is set by bifrost.go alongside req.SetModel(resolved). Plugins must +// not write to this key directly. +func GetResolvedAlias(ctx *BifrostContext) *ResolvedAlias { + if ctx == nil { + return nil + } + v := ctx.Value(BifrostContextKeyResolvedAlias) + if v == nil { + return nil + } + ra, _ := v.(*ResolvedAlias) + return ra +} + +// ResolveFamily returns the model family for the current attempt, walking +// the precedence: explicit alias ModelFamily → alias ModelName → alias +// ModelID → alias Key. When no alias matched, falls back to substring +// matching against fallbackModel (typically request.Model), preserving +// pre-refactor behavior. +// +// Returns an empty ModelFamily if nothing matches. +func ResolveFamily(ctx *BifrostContext, fallbackModel string) ModelFamily { + ra := GetResolvedAlias(ctx) + var candidates []string + if ra != nil && ra.Config != nil { + if ra.Config.ModelFamily != nil && *ra.Config.ModelFamily != "" { + return *ra.Config.ModelFamily + } + if ra.Config.ModelName != nil { + candidates = append(candidates, *ra.Config.ModelName) + } + candidates = append(candidates, ra.Config.ModelID, ra.Key) + } else { + candidates = append(candidates, fallbackModel) + } + for _, s := range candidates { + switch { + case IsAnthropicModel(s): + return ModelFamilyAnthropic + case IsMistralModel(s): + return ModelFamilyMistral + case IsGeminiModel(s): + return ModelFamilyGemini + case IsNovaModel(s): + return ModelFamilyNova + } + } + return "" +} + +// IsAnthropicModelFamily reports whether the current attempt resolves to the +// Anthropic model family. Thin wrapper over ResolveFamily so provider code +// reads uniformly at the many call sites that branch on Anthropic vs +// non-Anthropic (request shape, response parsing, anthropic-version header, +// URL path construction). model is passed as the substring-match fallback +// used when no alias is resolved in ctx — typically request.Model. +func IsAnthropicModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyAnthropic +} + // 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 { @@ -324,7 +403,7 @@ func (ka *KeyAliases) UnmarshalJSON(data []byte) error { return nil } var raw map[string]json.RawMessage - if err := json.Unmarshal(data, &raw); err != nil { + if err := sonic.Unmarshal(data, &raw); err != nil { return err } result := make(KeyAliases, len(raw)) @@ -337,13 +416,13 @@ func (ka *KeyAliases) UnmarshalJSON(data []byte) error { case '"': // Legacy string value — promote to AliasConfig{ModelID: ...}. var modelID string - if err := json.Unmarshal(entry, &modelID); err != nil { + if err := sonic.Unmarshal(entry, &modelID); err != nil { return fmt.Errorf("alias %q: %w", k, err) } result[k] = AliasConfig{ModelID: modelID} case '{': var ac AliasConfig - if err := json.Unmarshal(entry, &ac); err != nil { + if err := sonic.Unmarshal(entry, &ac); err != nil { return fmt.Errorf("alias %q: %w", k, err) } result[k] = ac diff --git a/core/schemas/account_test.go b/core/schemas/account_test.go index 18934bb41b..39634863a4 100644 --- a/core/schemas/account_test.go +++ b/core/schemas/account_test.go @@ -283,6 +283,126 @@ func TestKeyAliasesValidate(t *testing.T) { } } +func TestResolveFamilyPrecedence(t *testing.T) { + familyOpenAI := ModelFamilyOpenAI + + // Helper to build a BifrostContext carrying a ResolvedAlias. + withAlias := func(ra *ResolvedAlias) *BifrostContext { + bc := NewBifrostContext(nil, NoDeadline) + if ra != nil { + bc.SetValue(BifrostContextKeyResolvedAlias, ra) + } + return bc + } + + cases := []struct { + name string + ra *ResolvedAlias + fallback string + want ModelFamily + }{ + { + name: "tier 1: explicit ModelFamily wins over everything", + ra: &ResolvedAlias{ + Key: "some-claude-name", + Config: &AliasConfig{ + ModelID: "opaque-id", + ModelFamily: &familyOpenAI, // wins despite name/key smelling like Claude + }, + }, + fallback: "claude-3-5-sonnet", + want: ModelFamilyOpenAI, + }, + { + name: "tier 2: ModelName substring when no explicit family", + ra: &ResolvedAlias{ + Key: "best-model", + Config: &AliasConfig{ + ModelID: "opaque-id", + ModelName: Ptr("claude-3-5-sonnet"), + }, + }, + fallback: "opaque-id", + want: ModelFamilyAnthropic, + }, + { + name: "tier 3: ModelID substring when name absent", + ra: &ResolvedAlias{ + Key: "best-model", + Config: &AliasConfig{ + ModelID: "us.anthropic.claude-3-5-sonnet-20241022-v2:0", + }, + }, + fallback: "best-model", + want: ModelFamilyAnthropic, + }, + { + name: "tier 4: alias key substring when nothing else hits — the legacy 'best-claude→opaque-deployment-id' case the refactor is specifically meant to fix", + ra: &ResolvedAlias{ + Key: "best-claude", + Config: &AliasConfig{ + ModelID: "12345-azure-deployment", + }, + }, + fallback: "12345-azure-deployment", + want: ModelFamilyAnthropic, + }, + { + name: "no alias matched: fall back to substring on fallbackModel — preserves pre-refactor behavior", + ra: nil, + fallback: "claude-3-5-sonnet", + want: ModelFamilyAnthropic, + }, + { + name: "no alias and no substring hit anywhere", + ra: nil, + fallback: "totally-unknown-model", + want: "", + }, + { + name: "explicit empty ModelFamily pointer is treated as absent (falls through to name)", + ra: &ResolvedAlias{ + Key: "x", + Config: &AliasConfig{ + ModelID: "opaque-id", + ModelName: Ptr("claude-3-5-sonnet"), + ModelFamily: Ptr(ModelFamily("")), + }, + }, + fallback: "x", + want: ModelFamilyAnthropic, + }, + { + name: "first matching candidate wins (ModelName matches Anthropic before ModelID could match anything else)", + ra: &ResolvedAlias{ + Key: "x", + Config: &AliasConfig{ + ModelID: "mistral-large-2407", // would match Mistral but ModelName is checked first + ModelName: Ptr("claude-3-5-sonnet"), + }, + }, + fallback: "x", + want: ModelFamilyAnthropic, + }, + { + name: "uses fallback when ResolvedAlias.Config is nil (defensive)", + ra: &ResolvedAlias{Key: "x", Config: nil}, + // With Config==nil the candidates list is empty for the alias branch, + // so we drop to fallback substring matching. + fallback: "claude-3-haiku", + want: ModelFamilyAnthropic, + }, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + got := ResolveFamily(withAlias(c.ra), c.fallback) + if got != c.want { + t.Fatalf("ResolveFamily: got %q, want %q", got, c.want) + } + }) + } +} + func TestModelFamilyIsValid(t *testing.T) { valid := []ModelFamily{ ModelFamilyAnthropic, ModelFamilyOpenAI, ModelFamilyMistral, diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 201a600475..15f3793b2d 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -228,6 +228,7 @@ const ( BifrostContextKeyGovernanceIncludeOnlyKeys BifrostContextKey = "bf-governance-include-only-keys" // []string (to store the include-only key IDs for provider config routing (set by bifrost governance plugin - DO NOT SET THIS MANUALLY)) BifrostContextKeyNumberOfRetries BifrostContextKey = "bifrost-number-of-retries" // int (to store the number of retries (set by bifrost - DO NOT SET THIS MANUALLY)) BifrostContextKeyFallbackIndex BifrostContextKey = "bifrost-fallback-index" // int (to store the fallback index (set by bifrost - DO NOT SET THIS MANUALLY)) 0 for primary, 1 for first fallback, etc. + BifrostContextKeyResolvedAlias BifrostContextKey = "bifrost-resolved-alias" // *ResolvedAlias (set by bifrost after key-level alias resolution — providers read this for model_family routing and provider-specific overrides; nil/absent when no alias matched) BifrostContextKeyStreamEndIndicator BifrostContextKey = "bifrost-stream-end-indicator" // bool (set by bifrost - DO NOT SET THIS MANUALLY)) BifrostContextKeyStreamIdleTimeout BifrostContextKey = "bifrost-stream-idle-timeout" // time.Duration (per-chunk idle timeout for streaming) BifrostContextKeySkipKeySelection BifrostContextKey = "bifrost-skip-key-selection" // bool (will pass an empty key to the provider) From de4b699ac1a04237742e530ada9dbe097ffb8cb7 Mon Sep 17 00:00:00 2001 From: Pratham Mishra <99235987+Pratham-Mishra04@users.noreply.github.com> Date: Tue, 9 Jun 2026 15:58:44 +0530 Subject: [PATCH 009/108] feat: add alias-level region/ARN overrides and context-aware model family resolution for Bedrock (#4182) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Model-family detection on the Bedrock provider previously relied on substring matching against the raw model string. This meant that aliases pointing to opaque Bedrock deployments (e.g., inference profiles or cross-region ARNs) could not be correctly routed to the right request/response shape. This PR threads `*BifrostContext` through all family-detection call sites so that the resolved alias — including its explicit `ModelFamily`, `Region`, and `BedrockAliasCfg.InferenceProfileARN` — takes precedence over substring heuristics. ## Changes - Replaced all calls to `schemas.IsAnthropicModel`, `IsMistralModel`, `IsNovaModel`, `IsLlamaModel`, `IsCohereModel` (bare string matchers) with new context-aware variants: `IsAnthropicModelFamily`, `IsMistralModelFamily`, `IsNovaModelFamily`, `IsLlamaModelFamily`, `IsCohereModelFamily`, `IsTitanModelFamily`. These consult `ResolveFamily(ctx, model)` first, which reads the alias family tag before falling back to substring detection. - Added `ModelFamilyLlama`, `ModelFamilyTitan`, and `ModelFamilyCohere` to the `ModelFamily` enum and wired them into `ResolveFamily` so alias-tagged models route correctly. - Added `IsCohereModel` and `IsTitanModel` substring helpers used as the final fallback inside `ResolveFamily`. - Introduced `resolveBedrockARN(ctx, key)` to resolve the inference-profile ARN with priority: alias-level `BedrockAliasCfg.InferenceProfileARN` > key-level `BedrockKeyConfig.ARN`. Removed the inline ARN lookup from `getModelPathAndRegion`. - Updated `resolveBedrockRegion` and `getModelPathAndRegion` to accept `*BifrostContext` and honor the alias-level `Region` override between the model-string prefix (highest) and the key-level region (lower). - Propagated `*BifrostContext` into `DetermineEmbeddingModelType` and `ToBedrockEmbeddingInvokeResponse` so embedding model routing uses the same family resolution path. - Changed `convertToolConfigFromFiltered` to accept `*schemas.BifrostContext` instead of `context.Context` so family gates inside it can read the resolved alias. - Added nil-guard checks on `BifrostContext` in `bedrockAliasToolName` and `bedrockRestoreToolName`. - Added tests covering alias-level `Region` and `InferenceProfileARN` override priority in `region_test.go`. ## Type of change - [ ] Bug fix - [x] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./core/providers/bedrock/... ./core/schemas/... ``` To validate alias-level overrides end-to-end, configure an alias with an explicit `ModelFamily`, `Region`, or `BedrockAliasCfg.InferenceProfileARN` and confirm that: - Requests are signed and routed to the alias-specified region rather than the key-level region. - The ARN is prepended to the model path when `InferenceProfileARN` is set on the alias. - Embedding requests to an alias tagged `cohere` or `titan` use the correct request/response envelope regardless of the wire model string. ## Breaking changes - [x] Yes - [ ] No `ToBedrockEmbeddingInvokeResponse` now requires a `*schemas.BifrostContext` as its first argument. Any external callers of this function must be updated to pass the context. The `convertToolConfigFromFiltered` signature changed from `context.Context` to `*schemas.BifrostContext`. ## Related issues ## Security considerations No new secrets or auth surfaces introduced. ARN and region values continue to flow through the existing `EnvVar`/`GetValue()` resolution path. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable --- core/providers/bedrock/bedrock.go | 65 +++++++------ core/providers/bedrock/chat.go | 4 +- core/providers/bedrock/embedding.go | 13 ++- core/providers/bedrock/invoke.go | 6 +- core/providers/bedrock/mantle.go | 8 +- core/providers/bedrock/region_test.go | 92 ++++++++++++++++++- core/providers/bedrock/responses.go | 28 +++--- core/providers/bedrock/utils.go | 66 ++++++++++--- core/schemas/account.go | 44 ++++++++- core/schemas/utils.go | 13 +++ .../bifrost-http/integrations/bedrock.go | 2 +- 11 files changed, 262 insertions(+), 79 deletions(-) diff --git a/core/providers/bedrock/bedrock.go b/core/providers/bedrock/bedrock.go index 8b7c169b04..8457a5bd13 100644 --- a/core/providers/bedrock/bedrock.go +++ b/core/providers/bedrock/bedrock.go @@ -205,7 +205,7 @@ var retryableBedrockExceptions = map[string]int{ // Returns the response body, request latency, or an error if the request fails. func (provider *BedrockProvider) completeRequest(ctx *schemas.BifrostContext, jsonData []byte, path string, key schemas.Key, model string) ([]byte, time.Duration, map[string]string, *schemas.BifrostError) { config := key.BedrockKeyConfig - region := resolveBedrockRegion(key, model) + region := resolveBedrockRegion(ctx, key, model) // Create the request with the JSON body requestURL := fmt.Sprintf("https://bedrock-runtime.%s.amazonaws.com/model/%s", region, path) @@ -434,7 +434,7 @@ func (provider *BedrockProvider) completeAgentRuntimeRequest(ctx *schemas.Bifros // Returns the response body and an error if the request fails. func (provider *BedrockProvider) makeStreamingRequest(ctx *schemas.BifrostContext, jsonData []byte, key schemas.Key, model string, action string) (*http.Response, *schemas.BifrostError) { // Parse region and path in one pass to avoid running the regex twice. - path, region := provider.getModelPathAndRegion(action, model, key) + path, region := provider.getModelPathAndRegion(ctx, action, model, key) // Create HTTP request for streaming requestURL := fmt.Sprintf("https://bedrock-runtime.%s.amazonaws.com/model/%s", region, path) @@ -860,7 +860,7 @@ func (provider *BedrockProvider) TextCompletion(ctx *schemas.BifrostContext, key return nil, bifrostErr } - path, _ := provider.getModelPathAndRegion("invoke", request.Model, key) + path, _ := provider.getModelPathAndRegion(ctx, "invoke", request.Model, key) body, latency, providerResponseHeaders, err := provider.completeRequest(ctx, jsonData, path, key, request.Model) if providerResponseHeaders != nil { ctx.SetValue(schemas.BifrostContextKeyProviderResponseHeaders, providerResponseHeaders) @@ -872,14 +872,14 @@ func (provider *BedrockProvider) TextCompletion(ctx *schemas.BifrostContext, key // Handle model-specific response conversion var bifrostResponse *schemas.BifrostTextCompletionResponse switch { - case schemas.IsAnthropicModel(request.Model): + case schemas.IsAnthropicModelFamily(ctx, request.Model): var response BedrockAnthropicTextResponse if err := sonic.Unmarshal(body, &response); err != nil { return nil, providerUtils.NewBifrostOperationError("error parsing anthropic response", err) } bifrostResponse = response.ToBifrostTextCompletionResponse() - case schemas.IsMistralModel(request.Model): + case schemas.IsMistralModelFamily(ctx, request.Model): var response BedrockMistralTextResponse if err := sonic.Unmarshal(body, &response); err != nil { return nil, providerUtils.NewBifrostOperationError("error parsing mistral response", err) @@ -1089,7 +1089,7 @@ func (provider *BedrockProvider) ChatCompletion(ctx *schemas.BifrostContext, key } // Format the path with proper model identifier - path, _ := provider.getModelPathAndRegion("converse", request.Model, key) + path, _ := provider.getModelPathAndRegion(ctx, "converse", request.Model, key) // Create the signed request responseBody, latency, providerResponseHeaders, bifrostErr := provider.completeRequest(ctx, jsonData, path, key, request.Model) @@ -1477,7 +1477,7 @@ func (provider *BedrockProvider) Responses(ctx *schemas.BifrostContext, key sche } // Format the path with proper model identifier - path, _ := provider.getModelPathAndRegion("converse", request.Model, key) + path, _ := provider.getModelPathAndRegion(ctx, "converse", request.Model, key) // Create the signed request responseBody, latency, providerResponseHeaders, bifrostErr := provider.completeRequest(ctx, jsonData, path, key, request.Model) @@ -1837,7 +1837,7 @@ func (provider *BedrockProvider) Embedding(ctx *schemas.BifrostContext, key sche } // Determine model type - modelType, err := DetermineEmbeddingModelType(request.Model) + modelType, err := DetermineEmbeddingModelType(ctx, request.Model) if err != nil { return nil, providerUtils.NewConfigurationError(err.Error()) } @@ -1861,7 +1861,7 @@ func (provider *BedrockProvider) Embedding(ctx *schemas.BifrostContext, key sche if bifrostError != nil { return nil, bifrostError } - path, _ = provider.getModelPathAndRegion("invoke", request.Model, key) + path, _ = provider.getModelPathAndRegion(ctx, "invoke", request.Model, key) rawResponse, latency, providerResponseHeaders, bifrostError = provider.completeRequest(ctx, jsonData, path, key, request.Model) case "cohere": @@ -1874,7 +1874,7 @@ func (provider *BedrockProvider) Embedding(ctx *schemas.BifrostContext, key sche if bifrostError != nil { return nil, bifrostError } - path, _ = provider.getModelPathAndRegion("invoke", request.Model, key) + path, _ = provider.getModelPathAndRegion(ctx, "invoke", request.Model, key) rawResponse, latency, providerResponseHeaders, bifrostError = provider.completeRequest(ctx, jsonData, path, key, request.Model) default: @@ -2027,7 +2027,7 @@ func (provider *BedrockProvider) ImageGeneration(ctx *schemas.BifrostContext, ke var providerResponseHeaders map[string]string var path string - path, _ = provider.getModelPathAndRegion("invoke", request.Model, key) + path, _ = provider.getModelPathAndRegion(ctx, "invoke", request.Model, key) jsonData, bifrostError = providerUtils.CheckContextAndGetRequestBody( ctx, @@ -2100,7 +2100,7 @@ func (provider *BedrockProvider) ImageEdit(ctx *schemas.BifrostContext, key sche var bifrostError *schemas.BifrostError // Stability AI routing and task-type inference use the actual model ID. - path, _ := provider.getModelPathAndRegion("invoke", request.Model, key) + path, _ := provider.getModelPathAndRegion(ctx, "invoke", request.Model, key) jsonData, bifrostError = providerUtils.CheckContextAndGetRequestBody( ctx, @@ -2182,7 +2182,7 @@ func (provider *BedrockProvider) ImageVariation(ctx *schemas.BifrostContext, key } // Make API request (same URL as image generation) - path, _ := provider.getModelPathAndRegion("invoke", request.Model, key) + path, _ := provider.getModelPathAndRegion(ctx, "invoke", request.Model, key) rawResponse, latency, providerResponseHeaders, bifrostError := provider.completeRequest(ctx, jsonData, path, key, request.Model) if providerResponseHeaders != nil { ctx.SetValue(schemas.BifrostContextKeyProviderResponseHeaders, providerResponseHeaders) @@ -3575,32 +3575,29 @@ func (provider *BedrockProvider) BatchResults(ctx *schemas.BifrostContext, keys return batchResultsResp, nil } -// resolveBedrockRegion returns the AWS region to use for a request. -// the priority is: model string region > key configured region > default region -func resolveBedrockRegion(key schemas.Key, model string) string { - if region, _ := parseBedrockRegionAndModel(model); region != "" { - return region - } - if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.Region != nil && key.BedrockKeyConfig.Region.GetValue() != "" { - return key.BedrockKeyConfig.Region.GetValue() - } - return DefaultBedrockRegion -} - // getModelPathAndRegion is a helper that calls parseBedrockRegionAndModel -// once and returns both the request path and the AWS signing region -func (provider *BedrockProvider) getModelPathAndRegion(basePath, model string, key schemas.Key) (path, region string) { +// once and returns both the request path and the AWS signing region. +// Honors per-alias Region and BedrockAliasCfg.InferenceProfileARN overrides +// via the resolved alias in ctx. +func (provider *BedrockProvider) getModelPathAndRegion(ctx *schemas.BifrostContext, basePath, model string, key schemas.Key) (path, region string) { r, bareModel := parseBedrockRegionAndModel(model) if r == "" { - if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.Region != nil && key.BedrockKeyConfig.Region.GetValue() != "" { - r = key.BedrockKeyConfig.Region.GetValue() - } else { - r = DefaultBedrockRegion + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.Region != nil { + if v := ra.Config.Region.GetValue(); v != "" { + r = v + } + } + if r == "" { + if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.Region != nil && key.BedrockKeyConfig.Region.GetValue() != "" { + r = key.BedrockKeyConfig.Region.GetValue() + } else { + r = DefaultBedrockRegion + } } } p := fmt.Sprintf("%s/%s", bareModel, basePath) - if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.ARN != nil && key.BedrockKeyConfig.ARN.GetValue() != "" { - encodedModelIdentifier := url.PathEscape(fmt.Sprintf("%s/%s", key.BedrockKeyConfig.ARN.GetValue(), bareModel)) + if arn := resolveBedrockARN(ctx, key); arn != "" { + encodedModelIdentifier := url.PathEscape(fmt.Sprintf("%s/%s", arn, bareModel)) p = fmt.Sprintf("%s/%s", encodedModelIdentifier, basePath) } return p, r @@ -3627,7 +3624,7 @@ func (provider *BedrockProvider) CountTokens(ctx *schemas.BifrostContext, key sc } // Format the path with proper model identifier - path, _ := provider.getModelPathAndRegion("count-tokens", request.Model, key) + path, _ := provider.getModelPathAndRegion(ctx, "count-tokens", request.Model, key) // Send the request responseBody, latency, providerResponseHeaders, bifrostErr := provider.completeRequest(ctx, jsonData, path, key, request.Model) diff --git a/core/providers/bedrock/chat.go b/core/providers/bedrock/chat.go index ed089e7098..175b5973c2 100644 --- a/core/providers/bedrock/chat.go +++ b/core/providers/bedrock/chat.go @@ -25,7 +25,7 @@ func ToBedrockChatCompletionRequest(ctx *schemas.BifrostContext, bifrostReq *sch } input := bifrostReq.Input - if schemas.IsAnthropicModel(bifrostReq.Model) && ctx.Value(schemas.BifrostContextKeySupportsAssistantPrefill) == false { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) && ctx.Value(schemas.BifrostContextKeySupportsAssistantPrefill) == false { trimmed := len(input) for trimmed > 0 && input[trimmed-1].Role == schemas.ChatMessageRoleAssistant { trimmed-- @@ -46,7 +46,7 @@ func ToBedrockChatCompletionRequest(ctx *schemas.BifrostContext, bifrostReq *sch // Trim trailing whitespace from the last assistant message text blocks // (only for Anthropic models which use text-based prefill) lastMsgIndex := len(bedrockReq.Messages) - 1 - if schemas.IsAnthropicModel(bifrostReq.Model) && lastMsgIndex >= 0 && bedrockReq.Messages[lastMsgIndex].Role == BedrockMessageRoleAssistant { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) && lastMsgIndex >= 0 && bedrockReq.Messages[lastMsgIndex].Role == BedrockMessageRoleAssistant { blocks := bedrockReq.Messages[lastMsgIndex].Content for j := len(blocks) - 1; j >= 0; j-- { if blocks[j].Text != nil { diff --git a/core/providers/bedrock/embedding.go b/core/providers/bedrock/embedding.go index cb4ef19e88..b049bcfe61 100644 --- a/core/providers/bedrock/embedding.go +++ b/core/providers/bedrock/embedding.go @@ -3,7 +3,6 @@ package bedrock import ( "encoding/json" "fmt" - "strings" "github.com/maximhq/bifrost/core/schemas" ) @@ -159,12 +158,16 @@ func ToBedrockCohereEmbeddingRequest(bifrostReq *schemas.BifrostEmbeddingRequest return req, nil } -// DetermineEmbeddingModelType determines the embedding model type from the model name -func DetermineEmbeddingModelType(model string) (string, error) { +// DetermineEmbeddingModelType determines the embedding model type for the +// current attempt. It consults the resolved alias family first +// (model_family / model_name / model_id / alias key) and falls back to the +// substring detectors against the wire model — so an alias to an opaque +// Bedrock deployment that's tagged with the right family routes correctly. +func DetermineEmbeddingModelType(ctx *schemas.BifrostContext, model string) (string, error) { switch { - case strings.Contains(model, "amazon.titan-embed-text"): + case schemas.IsTitanModelFamily(ctx, model): return "titan", nil - case strings.Contains(model, "cohere.embed"): + case schemas.IsCohereModelFamily(ctx, model): return "cohere", nil default: return "", fmt.Errorf("unsupported embedding model: %s", model) diff --git a/core/providers/bedrock/invoke.go b/core/providers/bedrock/invoke.go index 65a8665af9..3edcd4e12c 100644 --- a/core/providers/bedrock/invoke.go +++ b/core/providers/bedrock/invoke.go @@ -956,7 +956,7 @@ func ToBedrockInvokeImagesResponse(ctx *schemas.BifrostContext, resp *schemas.Bi // Bedrock invoke API response format. // Single-embedding (Titan) responses use: {"embedding": [...], "inputTextTokenCount": N} // Multi-embedding (Cohere) responses use: {"embeddings": [[...],[...]], "response_type": "embeddings_floats"} -func ToBedrockEmbeddingInvokeResponse(resp *schemas.BifrostEmbeddingResponse) (interface{}, error) { +func ToBedrockEmbeddingInvokeResponse(ctx *schemas.BifrostContext, resp *schemas.BifrostEmbeddingResponse) (interface{}, error) { if resp == nil { return nil, fmt.Errorf("bifrost embedding response is nil") } @@ -975,7 +975,7 @@ func ToBedrockEmbeddingInvokeResponse(resp *schemas.BifrostEmbeddingResponse) (i return &BedrockInvokeEmbeddingResp{InputTextTokenCount: tokenCount}, nil } - // Use model name to distinguish Cohere from Titan — not batch size. + // Use the resolved family to distinguish Cohere from Titan — not batch size. // A single-input Cohere request must still return the Cohere envelope format. model := resp.Model if model == "" { @@ -986,7 +986,7 @@ func ToBedrockEmbeddingInvokeResponse(resp *schemas.BifrostEmbeddingResponse) (i } } - if strings.Contains(strings.ToLower(model), "cohere") { + if schemas.IsCohereModelFamily(ctx, model) { floats := make([][]float32, 0, len(resp.Data)) for _, d := range resp.Data { float32Emb := make([]float32, len(d.Embedding.EmbeddingArray)) diff --git a/core/providers/bedrock/mantle.go b/core/providers/bedrock/mantle.go index 5c85a1cbdc..932a9df544 100644 --- a/core/providers/bedrock/mantle.go +++ b/core/providers/bedrock/mantle.go @@ -60,7 +60,7 @@ func (provider *BedrockProvider) chatCompletionViaMantle( key schemas.Key, request *schemas.BifrostChatRequest, ) (*schemas.BifrostChatResponse, *schemas.BifrostError) { - region := resolveBedrockRegion(key, request.Model) + region := resolveBedrockRegion(ctx, key, request.Model) url := mantleURL(region, "chat/completions") // Build extraHeaders: always start with network-config headers, then overlay SigV4 if needed. @@ -106,7 +106,7 @@ func (provider *BedrockProvider) chatCompletionStreamViaMantle( key schemas.Key, request *schemas.BifrostChatRequest, ) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { - region := resolveBedrockRegion(key, request.Model) + region := resolveBedrockRegion(ctx, key, request.Model) url := mantleURL(region, "chat/completions") // Bearer: identical to Groq / any OpenAI-compatible provider. @@ -162,7 +162,7 @@ func (provider *BedrockProvider) responsesViaMantle( key schemas.Key, request *schemas.BifrostResponsesRequest, ) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) { - region := resolveBedrockRegion(key, request.Model) + region := resolveBedrockRegion(ctx, key, request.Model) url := mantleURL(region, "responses") extraHeaders := make(map[string]string, len(provider.networkConfig.ExtraHeaders)) @@ -204,7 +204,7 @@ func (provider *BedrockProvider) responsesStreamViaMantle( key schemas.Key, request *schemas.BifrostResponsesRequest, ) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { - region := resolveBedrockRegion(key, request.Model) + region := resolveBedrockRegion(ctx, key, request.Model) url := mantleURL(region, "responses") // Bearer: identical to Groq / any OpenAI-compatible provider. diff --git a/core/providers/bedrock/region_test.go b/core/providers/bedrock/region_test.go index cbd8b3d1cc..55bdf538e9 100644 --- a/core/providers/bedrock/region_test.go +++ b/core/providers/bedrock/region_test.go @@ -59,7 +59,7 @@ func TestGetModelPathStripsRegion(t *testing.T) { } for _, tc := range cases { t.Run(tc.model, func(t *testing.T) { - got, _ := provider.getModelPathAndRegion(tc.basePath, tc.model, key) + got, _ := provider.getModelPathAndRegion(nil, tc.basePath, tc.model, key) assert.Equal(t, tc.wantPath, got) }) } @@ -91,12 +91,98 @@ func TestGetModelPathStripsRegionWithARN(t *testing.T) { } for _, tc := range cases { t.Run(tc.model, func(t *testing.T) { - got, _ := provider.getModelPathAndRegion("converse", tc.model, key) + got, _ := provider.getModelPathAndRegion(nil, "converse", tc.model, key) assert.Equal(t, tc.wantPath, got) }) } } +// TestResolveBedrockRegion_AliasOverride verifies the per-alias Region +// override slots between the model-string prefix (highest priority) and the +// key-level Region (lower priority). +func TestResolveBedrockRegion_AliasOverride(t *testing.T) { + keyRegion := "us-east-1" + aliasRegion := "us-west-2" + key := schemas.Key{ + BedrockKeyConfig: &schemas.BedrockKeyConfig{ + Region: schemas.NewEnvVar(keyRegion), + }, + } + + // Build ctx carrying an alias with Region override. + ctx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "best-claude", + Config: &schemas.AliasConfig{ + ModelID: "anthropic.claude-3-5-sonnet-20241022-v2:0", + Region: schemas.NewEnvVar(aliasRegion), + }, + }) + + // Bare model — alias.Region wins over key.Region. + if got := resolveBedrockRegion(ctx, key, "anthropic.claude-3-5-sonnet-20241022-v2:0"); got != aliasRegion { + t.Errorf("alias override should win over key region: got %q, want %q", got, aliasRegion) + } + + // Model string with explicit region prefix — wins over alias override. + if got := resolveBedrockRegion(ctx, key, "eu-west-1/anthropic.claude-v2"); got != "eu-west-1" { + t.Errorf("model-string region should win over alias override: got %q", got) + } + + // No alias in ctx — falls through to key.Region. + emptyCtx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + if got := resolveBedrockRegion(emptyCtx, key, "anthropic.claude-v2"); got != keyRegion { + t.Errorf("no alias: should use key.Region: got %q, want %q", got, keyRegion) + } +} + +// TestResolveBedrockARN_AliasOverride verifies the BedrockAliasCfg +// InferenceProfileARN override takes precedence over key.BedrockKeyConfig.ARN. +func TestResolveBedrockARN_AliasOverride(t *testing.T) { + keyARN := "arn:aws:bedrock:us-east-1:1234567890:resource-config/default" + aliasARN := "arn:aws:bedrock:us-east-1:1234567890:inference-profile/us.anthropic.claude-3-7-sonnet" + key := schemas.Key{ + BedrockKeyConfig: &schemas.BedrockKeyConfig{ + ARN: schemas.NewEnvVar(keyARN), + }, + } + + // No alias — falls back to key.ARN. + if got := resolveBedrockARN(nil, key); got != keyARN { + t.Errorf("nil ctx: got %q, want key ARN %q", got, keyARN) + } + + // Alias with InferenceProfileARN override wins. + ctx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "best-claude", + Config: &schemas.AliasConfig{ + ModelID: "anthropic.claude-3-7-sonnet-20250219-v1:0", + BedrockAliasCfg: &schemas.BedrockAliasCfg{ + InferenceProfileARN: schemas.NewEnvVar(aliasARN), + }, + }, + }) + if got := resolveBedrockARN(ctx, key); got != aliasARN { + t.Errorf("alias override should win: got %q, want %q", got, aliasARN) + } + + // Empty alias ARN — falls through to key.ARN. + ctx2 := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx2.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "x", + Config: &schemas.AliasConfig{ + ModelID: "x", + BedrockAliasCfg: &schemas.BedrockAliasCfg{ + InferenceProfileARN: schemas.NewEnvVar(""), + }, + }, + }) + if got := resolveBedrockARN(ctx2, key); got != keyARN { + t.Errorf("empty alias ARN should fall through to key ARN: got %q, want %q", got, keyARN) + } +} + func TestResolveBedrockRegion(t *testing.T) { configuredRegion := "ap-southeast-1" key := schemas.Key{ @@ -119,7 +205,7 @@ func TestResolveBedrockRegion(t *testing.T) { } for _, tc := range cases { t.Run(tc.desc, func(t *testing.T) { - got := resolveBedrockRegion(tc.key, tc.model) + got := resolveBedrockRegion(nil, tc.key, tc.model) assert.Equal(t, tc.wantRegion, got) }) } diff --git a/core/providers/bedrock/responses.go b/core/providers/bedrock/responses.go index e081f6bacc..4132021b1c 100644 --- a/core/providers/bedrock/responses.go +++ b/core/providers/bedrock/responses.go @@ -1921,7 +1921,7 @@ func (request *BedrockConverseRequest) ToBifrostResponsesRequest(ctx *schemas.Bi continue } bifrostReq.Params.Tools = append(bifrostReq.Params.Tools, schemas.ResponsesTool{Type: toolType}) - } else if tool.CachePoint != nil && !schemas.IsNovaModel(bifrostReq.Model) { + } else if tool.CachePoint != nil && !schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { // add cache control to last tool in tools array if len(bifrostReq.Params.Tools) > 0 { bifrostReq.Params.Tools[len(bifrostReq.Params.Tools)-1].CacheControl = &schemas.CacheControl{ @@ -2018,7 +2018,7 @@ func (request *BedrockConverseRequest) ToBifrostResponsesRequest(ctx *schemas.Bi if request.InferenceConfig != nil && request.InferenceConfig.MaxTokens != nil { defaultMaxTokens = *request.InferenceConfig.MaxTokens } - if schemas.IsAnthropicModel(bifrostReq.Model) { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { minBudgetTokens = anthropic.MinimumReasoningMaxTokens } effort := providerUtils.GetReasoningEffortFromBudgetTokens(maxTokens, minBudgetTokens, defaultMaxTokens) @@ -2151,7 +2151,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. // map bifrost messages to bedrock messages using the new conversion method if bifrostReq.Input != nil { input := bifrostReq.Input - if schemas.IsAnthropicModel(bifrostReq.Model) && ctx.Value(schemas.BifrostContextKeySupportsAssistantPrefill) == false { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) && ctx.Value(schemas.BifrostContextKeySupportsAssistantPrefill) == false { trimmed := len(input) for trimmed > 0 && input[trimmed-1].Role != nil && *input[trimmed-1].Role == schemas.ResponsesInputMessageRoleAssistant { trimmed-- @@ -2180,7 +2180,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. // Trim trailing whitespace from the last assistant message text blocks // (only for Anthropic models which use text-based prefill) lastMsgIndex := len(bedrockReq.Messages) - 1 - if schemas.IsAnthropicModel(bifrostReq.Model) && lastMsgIndex >= 0 && bedrockReq.Messages[lastMsgIndex].Role == BedrockMessageRoleAssistant { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) && lastMsgIndex >= 0 && bedrockReq.Messages[lastMsgIndex].Role == BedrockMessageRoleAssistant { blocks := bedrockReq.Messages[lastMsgIndex].Content for j := len(blocks) - 1; j >= 0; j-- { if blocks[j].Text != nil { @@ -2217,15 +2217,15 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. // setting it to default max tokens tokenBudget = anthropic.MinimumReasoningMaxTokens } - if schemas.IsAnthropicModel(bifrostReq.Model) && tokenBudget < anthropic.MinimumReasoningMaxTokens { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) && tokenBudget < anthropic.MinimumReasoningMaxTokens { return nil, fmt.Errorf("reasoning.max_tokens must be >= %d for anthropic", anthropic.MinimumReasoningMaxTokens) } - if schemas.IsAnthropicModel(bifrostReq.Model) { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { bedrockReq.AdditionalModelRequestFields.Set("thinking", map[string]any{ "type": "enabled", "budget_tokens": tokenBudget, }) - } else if schemas.IsNovaModel(bifrostReq.Model) { + } else if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { minBudgetTokens := MinimumReasoningMaxTokens modelDefaultMaxTokens := providerUtils.GetMaxOutputTokensOrDefault(bifrostReq.Model, DefaultCompletionMaxTokens) defaultMaxTokens := modelDefaultMaxTokens @@ -2258,7 +2258,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. } } else { if bifrostReq.Params.Reasoning.Effort != nil && *bifrostReq.Params.Reasoning.Effort != "none" { - if schemas.IsNovaModel(bifrostReq.Model) { + if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { effort := *bifrostReq.Params.Reasoning.Effort typeStr := "enabled" switch effort { @@ -2283,7 +2283,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. } bedrockReq.AdditionalModelRequestFields.Set("reasoningConfig", config) - } else if schemas.IsAnthropicModel(bifrostReq.Model) { + } else if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { if anthropic.SupportsAdaptiveThinking(bifrostReq.Model) { // Opus 4.6+: adaptive thinking + output_config.effort effort := anthropic.MapBifrostEffortToAnthropic(*bifrostReq.Params.Reasoning.Effort) @@ -2338,11 +2338,11 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. }) } } else { - if schemas.IsAnthropicModel(bifrostReq.Model) { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { bedrockReq.AdditionalModelRequestFields.Set("thinking", map[string]any{ "type": "disabled", }) - } else if schemas.IsNovaModel(bifrostReq.Model) { + } else if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { bedrockReq.AdditionalModelRequestFields.Set("reasoningConfig", map[string]any{ "type": "disabled", }) @@ -2449,7 +2449,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. } bedrockTools = append(bedrockTools, bedrockTool) - if tool.CacheControl != nil && !schemas.IsNovaModel(bifrostReq.Model) { + if tool.CacheControl != nil && !schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { bedrockTools = append(bedrockTools, BedrockTool{ CachePoint: &BedrockCachePoint{ Type: BedrockCachePointTypeDefault, @@ -2480,7 +2480,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. // behavior. See per-model support matrix at // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html // (mirrors the gate in convertToolConfigFromFiltered for ChatCompletions). - if bedrockToolChoice != nil && bedrockToolChoice.Tool != nil && schemas.IsLlamaModel(bifrostReq.Model) { + if bedrockToolChoice != nil && bedrockToolChoice.Tool != nil && schemas.IsLlamaModelFamily(ctx, bifrostReq.Model) { bedrockToolChoice = nil } if bedrockToolChoice != nil { @@ -2510,7 +2510,7 @@ func ToBedrockResponsesRequest(ctx *schemas.BifrostContext, bifrostReq *schemas. thinkingEnabled := bifrostReq.Params.Reasoning != nil && (bifrostReq.Params.Reasoning.MaxTokens != nil || (bifrostReq.Params.Reasoning.Effort != nil && *bifrostReq.Params.Reasoning.Effort != "none")) - if !schemas.IsLlamaModel(bifrostReq.Model) && !thinkingEnabled { + if !schemas.IsLlamaModelFamily(ctx, bifrostReq.Model) && !thinkingEnabled { bedrockReq.ToolConfig.ToolChoice = &BedrockToolChoice{ Tool: &BedrockToolChoiceTool{ Name: responsesStructuredOutputTool.ToolSpec.Name, diff --git a/core/providers/bedrock/utils.go b/core/providers/bedrock/utils.go index 8ec4ee51cf..bf86916acf 100644 --- a/core/providers/bedrock/utils.go +++ b/core/providers/bedrock/utils.go @@ -41,6 +41,43 @@ func parseBedrockRegionAndModel(model string) (region, bareModel string) { return "", model } +// resolveBedrockRegion returns the AWS region to use for a request. +// Priority: model-string region prefix > alias-level Region > key-level +// BedrockKeyConfig.Region > DefaultBedrockRegion. The model-string prefix +// stays highest since it's the most explicit signal — when an admin types a +// region into their model ID they expect that to win. +func resolveBedrockRegion(ctx *schemas.BifrostContext, key schemas.Key, model string) string { + if region, _ := parseBedrockRegionAndModel(model); region != "" { + return region + } + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.Region != nil { + if v := ra.Config.Region.GetValue(); v != "" { + return v + } + } + if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.Region != nil && key.BedrockKeyConfig.Region.GetValue() != "" { + return key.BedrockKeyConfig.Region.GetValue() + } + return DefaultBedrockRegion +} + +// resolveBedrockARN returns the inference-profile / resource ARN prepended +// to the Bedrock URL path. Priority: alias-level BedrockAliasCfg +// InferenceProfileARN > key-level BedrockKeyConfig.ARN. Returns empty when +// neither is set, in which case getModelPathAndRegion emits the bare model +// path. +func resolveBedrockARN(ctx *schemas.BifrostContext, key schemas.Key) string { + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.BedrockAliasCfg != nil && ra.Config.BedrockAliasCfg.InferenceProfileARN != nil { + if v := ra.Config.BedrockAliasCfg.InferenceProfileARN.GetValue(); v != "" { + return v + } + } + if key.BedrockKeyConfig != nil && key.BedrockKeyConfig.ARN != nil { + return key.BedrockKeyConfig.ARN.GetValue() + } + return "" +} + var ( invalidCharRegex = regexp.MustCompile(`[^a-zA-Z0-9\s\-\(\)\[\]]`) multiSpaceRegex = regexp.MustCompile(`\s{2,}`) @@ -157,7 +194,7 @@ func bedrockAliasToolName(ctx context.Context, name string) string { } alias := hash + "_" + semanticName - if bifrostCtx, ok := ctx.(*schemas.BifrostContext); ok && alias != name { + if bifrostCtx, ok := ctx.(*schemas.BifrostContext); ok && bifrostCtx != nil && alias != name { aliases, _ := bifrostCtx.Value(bedrockToolNameAliasKey{}).(map[string]string) if aliases == nil { aliases = make(map[string]string) @@ -170,7 +207,7 @@ func bedrockAliasToolName(ctx context.Context, name string) string { // bedrockRestoreToolName maps a Bedrock wire-name alias back to the caller's tool name. func bedrockRestoreToolName(ctx context.Context, name string) string { - if bifrostCtx, ok := ctx.(*schemas.BifrostContext); ok { + if bifrostCtx, ok := ctx.(*schemas.BifrostContext); ok && bifrostCtx != nil { if aliases, _ := bifrostCtx.Value(bedrockToolNameAliasKey{}).(map[string]string); aliases != nil { if original, ok := aliases[name]; ok { return original @@ -254,7 +291,7 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr // setting it to default max tokens tokenBudget = anthropic.MinimumReasoningMaxTokens } - if schemas.IsAnthropicModel(bifrostReq.Model) { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { if tokenBudget < anthropic.MinimumReasoningMaxTokens { return fmt.Errorf("reasoning.max_tokens must be >= %d for anthropic", anthropic.MinimumReasoningMaxTokens) } @@ -262,7 +299,7 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr "type": "enabled", "budget_tokens": tokenBudget, }) - } else if schemas.IsNovaModel(bifrostReq.Model) { + } else if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { minBudgetTokens := MinimumReasoningMaxTokens modelDefaultMaxTokens := providerUtils.GetMaxOutputTokensOrDefault(bifrostReq.Model, DefaultCompletionMaxTokens) defaultMaxTokens := modelDefaultMaxTokens @@ -319,7 +356,7 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr } } } - if schemas.IsNovaModel(bifrostReq.Model) { + if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { effort := *bifrostReq.Params.Reasoning.Effort typeStr := "enabled" switch effort { @@ -343,7 +380,7 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr } bedrockReq.AdditionalModelRequestFields.Set("reasoningConfig", config) - } else if schemas.IsAnthropicModel(bifrostReq.Model) { + } else if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { if anthropic.SupportsAdaptiveThinking(bifrostReq.Model) { // Opus 4.6+: adaptive thinking + output_config.effort effort := anthropic.MapBifrostEffortToAnthropic(*bifrostReq.Params.Reasoning.Effort) @@ -371,11 +408,11 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr } } } else { - if schemas.IsAnthropicModel(bifrostReq.Model) { + if schemas.IsAnthropicModelFamily(ctx, bifrostReq.Model) { bedrockReq.AdditionalModelRequestFields.Set("thinking", map[string]any{ "type": "disabled", }) - } else if schemas.IsNovaModel(bifrostReq.Model) { + } else if schemas.IsNovaModelFamily(ctx, bifrostReq.Model) { bedrockReq.AdditionalModelRequestFields.Set("reasoningConfig", map[string]any{ "type": "disabled", }) @@ -408,7 +445,7 @@ func convertChatParameters(ctx *schemas.BifrostContext, bifrostReq *schemas.Bifr thinkingEnabled := bifrostReq.Params.Reasoning != nil && (bifrostReq.Params.Reasoning.MaxTokens != nil || (bifrostReq.Params.Reasoning.Effort != nil && *bifrostReq.Params.Reasoning.Effort != "none")) - if !schemas.IsLlamaModel(bifrostReq.Model) && !thinkingEnabled { + if !schemas.IsLlamaModelFamily(ctx, bifrostReq.Model) && !thinkingEnabled { bedrockReq.ToolConfig.ToolChoice = &BedrockToolChoice{ Tool: &BedrockToolChoiceTool{ Name: responseFormatTool.ToolSpec.Name, @@ -1677,7 +1714,12 @@ func convertToolConfig(model string, params *schemas.ChatParameters) *BedrockToo // pre-filtered tool set. convertChatParameters uses this to avoid filtering // twice (once here, once in collectBedrockServerTools). The public // convertToolConfig entry point is a thin wrapper preserved for tests. -func convertToolConfigFromFiltered(ctx context.Context, model string, params *schemas.ChatParameters, filtered []schemas.ChatTool) *BedrockToolConfig { +// +// ctx is the BifrostContext (not context.Context) so the family gates inside +// this function can consult the resolved alias and honor explicit +// AliasConfig.ModelFamily overrides. Test paths may pass nil — family +// detection then falls back to substring matching on model. +func convertToolConfigFromFiltered(ctx *schemas.BifrostContext, model string, params *schemas.ChatParameters, filtered []schemas.ChatTool) *BedrockToolConfig { if params == nil { return nil } @@ -1717,7 +1759,7 @@ func convertToolConfigFromFiltered(ctx context.Context, model string, params *sc } bedrockTools = append(bedrockTools, bedrockTool) - if tool.CacheControl != nil && !schemas.IsNovaModel(model) { + if tool.CacheControl != nil && !schemas.IsNovaModelFamily(ctx, model) { bedrockTools = append(bedrockTools, BedrockTool{ CachePoint: &BedrockCachePoint{ Type: BedrockCachePointTypeDefault, @@ -1774,7 +1816,7 @@ func convertToolConfigFromFiltered(ctx context.Context, model string, params *sc // behavior. See per-model support matrix at // https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html // (mirrors the synthetic-tool gate in convertChatParameters). - if toolChoice != nil && toolChoice.Tool != nil && schemas.IsLlamaModel(model) { + if toolChoice != nil && toolChoice.Tool != nil && schemas.IsLlamaModelFamily(ctx, model) { toolChoice = nil } if toolChoice != nil { diff --git a/core/schemas/account.go b/core/schemas/account.go index 6a6e1bb519..a23f57ce74 100644 --- a/core/schemas/account.go +++ b/core/schemas/account.go @@ -156,6 +156,7 @@ const ( ModelFamilyMistral ModelFamily = "mistral" ModelFamilyCohere ModelFamily = "cohere" ModelFamilyGemini ModelFamily = "gemini" + ModelFamilyLlama ModelFamily = "llama" ModelFamilyNova ModelFamily = "nova" ModelFamilyTitan ModelFamily = "titan" ) @@ -167,7 +168,8 @@ func (mf *ModelFamily) IsValid() bool { } switch *mf { case ModelFamilyAnthropic, ModelFamilyOpenAI, ModelFamilyMistral, - ModelFamilyCohere, ModelFamilyGemini, ModelFamilyNova, ModelFamilyTitan: + ModelFamilyCohere, ModelFamilyGemini, ModelFamilyLlama, + ModelFamilyNova, ModelFamilyTitan: return true } return false @@ -359,8 +361,14 @@ func ResolveFamily(ctx *BifrostContext, fallbackModel string) ModelFamily { return ModelFamilyMistral case IsGeminiModel(s): return ModelFamilyGemini + case IsLlamaModel(s): + return ModelFamilyLlama case IsNovaModel(s): return ModelFamilyNova + case IsTitanModel(s): + return ModelFamilyTitan + case IsCohereModel(s): + return ModelFamilyCohere } } return "" @@ -376,6 +384,40 @@ func IsAnthropicModelFamily(ctx *BifrostContext, model string) bool { return ResolveFamily(ctx, model) == ModelFamilyAnthropic } +// IsMistralModelFamily reports whether the current attempt resolves to the +// Mistral model family. See IsAnthropicModelFamily for usage notes. +func IsMistralModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyMistral +} + +// IsLlamaModelFamily reports whether the current attempt resolves to the +// Llama model family. Used by Bedrock to gate tool_choice handling — AWS +// Bedrock Converse rejects toolConfig.toolChoice.tool on Meta Llama variants. +func IsLlamaModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyLlama +} + +// IsNovaModelFamily reports whether the current attempt resolves to the +// Amazon Nova model family. Used by Bedrock to gate cache-point insertion +// and tool shaping that differs from Anthropic. +func IsNovaModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyNova +} + +// IsCohereModelFamily reports whether the current attempt resolves to the +// Cohere model family. Used by Bedrock to pick the Cohere request/response +// shape for embeddings (vs. the Titan envelope). +func IsCohereModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyCohere +} + +// IsTitanModelFamily reports whether the current attempt resolves to the +// Amazon Titan model family. Used by Bedrock to pick the Titan embedding +// request/response envelope. +func IsTitanModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyTitan +} + // 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/utils.go b/core/schemas/utils.go index 6ab1b9a3a7..ac55a928fb 100644 --- a/core/schemas/utils.go +++ b/core/schemas/utils.go @@ -1433,6 +1433,19 @@ func IsImagenModel(model string) bool { return strings.Contains(strings.ToLower(model), "imagen") } +// IsCohereModel checks if the model is a Cohere model. Matches the Bedrock +// identifier prefix ("cohere.embed-*", "cohere.command-*") which is the wire +// shape that flows through alias resolution. +func IsCohereModel(model string) bool { + return strings.Contains(model, "cohere") +} + +// IsTitanModel checks if the model is an Amazon Titan model. Matches the +// Bedrock identifier prefix ("amazon.titan-*"). +func IsTitanModel(model string) bool { + return strings.Contains(model, "titan") +} + // List of grok reasoning models var grokReasoningModels = []string{ "grok-3", diff --git a/transports/bifrost-http/integrations/bedrock.go b/transports/bifrost-http/integrations/bedrock.go index efe2045261..89b784e440 100644 --- a/transports/bifrost-http/integrations/bedrock.go +++ b/transports/bifrost-http/integrations/bedrock.go @@ -255,7 +255,7 @@ func createBedrockInvokeRouteConfig(pathPrefix string, handlerStore lib.HandlerS return bedrock.ToBedrockInvokeMessagesResponse(ctx, resp) }, EmbeddingResponseConverter: func(ctx *schemas.BifrostContext, resp *schemas.BifrostEmbeddingResponse) (interface{}, error) { - return bedrock.ToBedrockEmbeddingInvokeResponse(resp) + return bedrock.ToBedrockEmbeddingInvokeResponse(ctx, resp) }, ImageGenerationResponseConverter: func(ctx *schemas.BifrostContext, resp *schemas.BifrostImageGenerationResponse) (interface{}, error) { return bedrock.ToBedrockInvokeImagesResponse(ctx, resp) From 7f2150fe30640b3f1c862df8da9cf044aeb83ea4 Mon Sep 17 00:00:00 2001 From: Pratham Mishra <99235987+Pratham-Mishra04@users.noreply.github.com> Date: Tue, 9 Jun 2026 16:00:45 +0530 Subject: [PATCH 010/108] feat: add per-alias Vertex project/region overrides and context-aware model family resolution (#4183) ## Summary Introduces per-alias overrides for Vertex AI's `project_id`, `project_number`, and `region` configuration values, allowing a single Vertex credential to serve requests across multiple GCP projects and regions. This is particularly useful when different model families (e.g., Anthropic Claude on `us-east5`, Gemini on `us-central1`) are deployed in separate GCP projects or regions. ## Changes - Added `resolveVertexProjectID`, `resolveVertexProjectNumber`, and `resolveVertexRegion` helper functions in `utils.go` that check for alias-level overrides in `BifrostContext` before falling back to key-level configuration values. - Replaced all direct `key.VertexKeyConfig.*` field accesses throughout the Vertex provider with calls to these resolver functions, covering all operations: chat completion, streaming, embeddings, responses, cached content, image/video generation, reranking, token counting, and passthrough. - Replaced model-family detection calls (`IsAnthropicModel`, `IsGeminiModel`, `IsGemmaModel`, `IsImagenModel`, `IsVeoModel`, `IsMistralModel`) with context-aware variants (`IsAnthropicModelFamily`, `IsGeminiModelFamily`, etc.) so that alias-level `ModelFamily` overrides are respected when routing requests to the correct Vertex endpoint and request format. - Added `ModelFamilyGemma`, `ModelFamilyImagen`, and `ModelFamilyVeo` as first-class `ModelFamily` constants and registered them in `ResolveFamily`, with Imagen and Veo checked before Gemini to avoid substring-match conflicts. - Added context-aware `IsGeminiModelFamily`, `IsGemmaModelFamily`, `IsImagenModelFamily`, and `IsVeoModelFamily` helpers to `account.go`. - Added unit tests covering alias override precedence, empty-alias fallthrough to key-level values, and nil-context handling for all three resolver functions. ## Type of change - [ ] Bug fix - [x] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [ ] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./core/providers/vertex/... -run TestResolveVertex go test ./core/... ``` Configure a Vertex key alias with `VertexAliasCfg.ProjectID` or `AliasConfig.Region` set to a value different from the key-level config and verify that requests are routed to the alias-specified project/region rather than the key-level defaults. ## Breaking changes - [ ] Yes - [x] No ## Related issues ## Security considerations No new secrets or auth mechanisms are introduced. Alias-level project/region values follow the same `EnvVar` resolution path as key-level values, so secrets can still be sourced from environment variables rather than being hardcoded. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable ## Summary by CodeRabbit * **New Features** * Extended support for additional Google Vertex AI model families. * Added flexible configuration resolution with alias-level parameter overrides for the Vertex provider. * **Tests** * Added comprehensive tests validating configuration override behavior for Vertex parameters. --- core/providers/vertex/cachedcontents.go | 40 +++---- core/providers/vertex/utils.go | 49 ++++++++ core/providers/vertex/utils_test.go | 143 ++++++++++++++++++++++ core/providers/vertex/vertex.go | 150 ++++++++++++------------ core/schemas/account.go | 43 ++++++- 5 files changed, 329 insertions(+), 96 deletions(-) diff --git a/core/providers/vertex/cachedcontents.go b/core/providers/vertex/cachedcontents.go index 73bf8d1dd8..05adb24bcb 100644 --- a/core/providers/vertex/cachedcontents.go +++ b/core/providers/vertex/cachedcontents.go @@ -115,13 +115,13 @@ func (provider *VertexProvider) CachedContentCreate(ctx *schemas.BifrostContext, return nil, providerUtils.NewBifrostOperationError("model is required for cached content create", nil) } - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { - return nil, providerUtils.NewConfigurationError("project_id is not set in vertex key config") + return nil, providerUtils.NewConfigurationError("project_id is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { - return nil, providerUtils.NewConfigurationError("region is not set in vertex key config") + return nil, providerUtils.NewConfigurationError("region is not set") } model := expandVertexModelPath(request.Model, projectID, region) @@ -210,13 +210,13 @@ func (provider *VertexProvider) CachedContentCreate(ctx *schemas.BifrostContext, } func (provider *VertexProvider) cachedContentListByKey(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostCachedContentListRequest) (*schemas.BifrostCachedContentListResponse, time.Duration, *schemas.BifrostError) { - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { - return nil, 0, providerUtils.NewConfigurationError("project_id is not set in vertex key config") + return nil, 0, providerUtils.NewConfigurationError("project_id is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { - return nil, 0, providerUtils.NewConfigurationError("region is not set in vertex key config") + return nil, 0, providerUtils.NewConfigurationError("region is not set") } req := fasthttp.AcquireRequest() @@ -292,13 +292,13 @@ func (provider *VertexProvider) CachedContentList(ctx *schemas.BifrostContext, k } func (provider *VertexProvider) cachedContentRetrieveByKey(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostCachedContentRetrieveRequest) (*schemas.BifrostCachedContentRetrieveResponse, time.Duration, *schemas.BifrostError) { - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { - return nil, 0, providerUtils.NewConfigurationError("project_id is not set in vertex key config") + return nil, 0, providerUtils.NewConfigurationError("project_id is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { - return nil, 0, providerUtils.NewConfigurationError("region is not set in vertex key config") + return nil, 0, providerUtils.NewConfigurationError("region is not set") } req := fasthttp.AcquireRequest() @@ -372,13 +372,13 @@ func (provider *VertexProvider) CachedContentRetrieve(ctx *schemas.BifrostContex } func (provider *VertexProvider) cachedContentUpdateByKey(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostCachedContentUpdateRequest) (*schemas.BifrostCachedContentUpdateResponse, time.Duration, *schemas.BifrostError) { - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { - return nil, 0, providerUtils.NewConfigurationError("project_id is not set in vertex key config") + return nil, 0, providerUtils.NewConfigurationError("project_id is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { - return nil, 0, providerUtils.NewConfigurationError("region is not set in vertex key config") + return nil, 0, providerUtils.NewConfigurationError("region is not set") } body := vertexCachedContent{} @@ -482,13 +482,13 @@ func (provider *VertexProvider) CachedContentUpdate(ctx *schemas.BifrostContext, } func (provider *VertexProvider) cachedContentDeleteByKey(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostCachedContentDeleteRequest) (*schemas.BifrostCachedContentDeleteResponse, time.Duration, *schemas.BifrostError) { - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { - return nil, 0, providerUtils.NewConfigurationError("project_id is not set in vertex key config") + return nil, 0, providerUtils.NewConfigurationError("project_id is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { - return nil, 0, providerUtils.NewConfigurationError("region is not set in vertex key config") + return nil, 0, providerUtils.NewConfigurationError("region is not set") } req := fasthttp.AcquireRequest() diff --git a/core/providers/vertex/utils.go b/core/providers/vertex/utils.go index ef806ab750..0fd3b3fcc1 100644 --- a/core/providers/vertex/utils.go +++ b/core/providers/vertex/utils.go @@ -10,6 +10,55 @@ import ( schemas "github.com/maximhq/bifrost/core/schemas" ) +// resolveVertexProjectID returns the GCP project ID for the current attempt. +// Priority: alias-level VertexAliasCfg.ProjectID > key-level +// VertexKeyConfig.ProjectID. Per-alias override lets one Vertex credential +// span deployments across distinct GCP projects (e.g. Anthropic models in +// one project, Gemini in another). +func resolveVertexProjectID(ctx *schemas.BifrostContext, key schemas.Key) string { + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.VertexAliasCfg != nil && ra.Config.VertexAliasCfg.ProjectID != nil { + if v := ra.Config.VertexAliasCfg.ProjectID.GetValue(); v != "" { + return v + } + } + if key.VertexKeyConfig != nil { + return key.VertexKeyConfig.ProjectID.GetValue() + } + return "" +} + +// resolveVertexProjectNumber returns the GCP project number for the current +// attempt. Same precedence as resolveVertexProjectID. +func resolveVertexProjectNumber(ctx *schemas.BifrostContext, key schemas.Key) string { + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.VertexAliasCfg != nil && ra.Config.VertexAliasCfg.ProjectNumber != nil { + if v := ra.Config.VertexAliasCfg.ProjectNumber.GetValue(); v != "" { + return v + } + } + if key.VertexKeyConfig != nil { + return key.VertexKeyConfig.ProjectNumber.GetValue() + } + return "" +} + +// resolveVertexRegion returns the Vertex region for the current attempt. +// Priority: alias-level AliasConfig.Region (top-level, shared with other +// providers) > key-level VertexKeyConfig.Region. Different Vertex model +// families publish in different regions (Anthropic on us-east5, Gemini on +// us-central1, …), so per-alias overrides let one credential reach all of +// them. +func resolveVertexRegion(ctx *schemas.BifrostContext, key schemas.Key) string { + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.Region != nil { + if v := ra.Config.Region.GetValue(); v != "" { + return v + } + } + if key.VertexKeyConfig != nil { + return key.VertexKeyConfig.Region.GetValue() + } + return "" +} + // getRequestBodyForAnthropicResponses serializes a BifrostResponsesRequest into the Anthropic wire format for Vertex AI. // Compared to the native Anthropic path, it strips model/region fields, remaps tool versions, injects beta headers // into the request body (rather than HTTP headers), and pins the Anthropic API version to DefaultVertexAnthropicVersion. diff --git a/core/providers/vertex/utils_test.go b/core/providers/vertex/utils_test.go index cc2b209890..acf6a557ad 100644 --- a/core/providers/vertex/utils_test.go +++ b/core/providers/vertex/utils_test.go @@ -354,3 +354,146 @@ func TestVertexRegionToPool(t *testing.T) { }) } } + +// TestResolveVertexProjectID_AliasOverride verifies the per-alias ProjectID +// override lets one Vertex credential serve deployments across distinct GCP +// projects. +func TestResolveVertexProjectID_AliasOverride(t *testing.T) { + keyProject := "key-level-project" + aliasProject := "alias-level-project" + key := schemas.Key{ + VertexKeyConfig: &schemas.VertexKeyConfig{ + ProjectID: *schemas.NewEnvVar(keyProject), + }, + } + + if got := resolveVertexProjectID(nil, key); got != keyProject { + t.Errorf("nil ctx: got %q, want key-level %q", got, keyProject) + } + + ctx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + if got := resolveVertexProjectID(ctx, key); got != keyProject { + t.Errorf("empty ctx: got %q, want key-level %q", got, keyProject) + } + + ctx.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "best-claude", + Config: &schemas.AliasConfig{ + ModelID: "claude-sonnet-4-5", + VertexAliasCfg: &schemas.VertexAliasCfg{ + ProjectID: schemas.NewEnvVar(aliasProject), + }, + }, + }) + if got := resolveVertexProjectID(ctx, key); got != aliasProject { + t.Errorf("alias override should win: got %q, want %q", got, aliasProject) + } + + // Empty alias ProjectID falls through to key-level. + ctx2 := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx2.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "x", + Config: &schemas.AliasConfig{ + ModelID: "x", + VertexAliasCfg: &schemas.VertexAliasCfg{ + ProjectID: schemas.NewEnvVar(""), + }, + }, + }) + if got := resolveVertexProjectID(ctx2, key); got != keyProject { + t.Errorf("empty alias ProjectID should fall through: got %q, want %q", got, keyProject) + } +} + +// TestResolveVertexRegion_AliasOverride verifies the top-level +// AliasConfig.Region override for Vertex. +func TestResolveVertexRegion_AliasOverride(t *testing.T) { + keyRegion := "us-central1" + aliasRegion := "us-east5" + key := schemas.Key{ + VertexKeyConfig: &schemas.VertexKeyConfig{ + Region: *schemas.NewEnvVar(keyRegion), + }, + } + + if got := resolveVertexRegion(nil, key); got != keyRegion { + t.Errorf("nil ctx: got %q, want %q", got, keyRegion) + } + + ctx0 := schemas.NewBifrostContext(nil, schemas.NoDeadline) + if got := resolveVertexRegion(ctx0, key); got != keyRegion { + t.Errorf("empty ctx: got %q, want key-level %q", got, keyRegion) + } + + ctx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "best-claude", + Config: &schemas.AliasConfig{ + ModelID: "claude-sonnet-4-5", + Region: schemas.NewEnvVar(aliasRegion), + }, + }) + if got := resolveVertexRegion(ctx, key); got != aliasRegion { + t.Errorf("alias Region should win: got %q, want %q", got, aliasRegion) + } + + ctx2 := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx2.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "x", + Config: &schemas.AliasConfig{ + ModelID: "x", + Region: schemas.NewEnvVar(""), + }, + }) + if got := resolveVertexRegion(ctx2, key); got != keyRegion { + t.Errorf("empty alias Region should fall through: got %q, want %q", got, keyRegion) + } +} + +// TestResolveVertexProjectNumber_AliasOverride mirrors the ProjectID test. +func TestResolveVertexProjectNumber_AliasOverride(t *testing.T) { + keyNumber := "111111" + aliasNumber := "222222" + key := schemas.Key{ + VertexKeyConfig: &schemas.VertexKeyConfig{ + ProjectNumber: *schemas.NewEnvVar(keyNumber), + }, + } + + if got := resolveVertexProjectNumber(nil, key); got != keyNumber { + t.Errorf("nil ctx: got %q, want %q", got, keyNumber) + } + + ctx0 := schemas.NewBifrostContext(nil, schemas.NoDeadline) + if got := resolveVertexProjectNumber(ctx0, key); got != keyNumber { + t.Errorf("empty ctx: got %q, want key-level %q", got, keyNumber) + } + + ctx := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "x", + Config: &schemas.AliasConfig{ + ModelID: "x", + VertexAliasCfg: &schemas.VertexAliasCfg{ + ProjectNumber: schemas.NewEnvVar(aliasNumber), + }, + }, + }) + if got := resolveVertexProjectNumber(ctx, key); got != aliasNumber { + t.Errorf("alias ProjectNumber should win: got %q, want %q", got, aliasNumber) + } + + ctx2 := schemas.NewBifrostContext(nil, schemas.NoDeadline) + ctx2.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "x", + Config: &schemas.AliasConfig{ + ModelID: "x", + VertexAliasCfg: &schemas.VertexAliasCfg{ + ProjectNumber: schemas.NewEnvVar(""), + }, + }, + }) + if got := resolveVertexProjectNumber(ctx2, key); got != keyNumber { + t.Errorf("empty alias ProjectNumber should fall through: got %q, want %q", got, keyNumber) + } +} diff --git a/core/providers/vertex/vertex.go b/core/providers/vertex/vertex.go index 4656da7ac1..eac3552a94 100644 --- a/core/providers/vertex/vertex.go +++ b/core/providers/vertex/vertex.go @@ -197,7 +197,7 @@ func (provider *VertexProvider) GetProviderKey() schemas.ModelProvider { // 1. If deployments or allowedModels are configured, return those (no API call needed) // 2. Otherwise, fetch from the publishers.models.list API endpoint (Model Garden) func (provider *VertexProvider) listModelsByKey(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostListModelsRequest) (*schemas.BifrostListModelsResponse, *schemas.BifrostError) { - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } @@ -426,7 +426,7 @@ func (provider *VertexProvider) ChatCompletion(ctx *schemas.BifrostContext, key var extraParams map[string]interface{} var err error - if schemas.IsAnthropicModel(request.Model) { + if schemas.IsAnthropicModelFamily(ctx, request.Model) { // Anthropic-on-Vertex doesn't accept URL-source document blocks. // Inline any URL documents to base64 before the converter runs. if err := inlineDocumentURLs(ctx, request); err != nil { @@ -467,7 +467,7 @@ func (provider *VertexProvider) ChatCompletion(ctx *schemas.BifrostContext, key if err != nil { return nil, fmt.Errorf("failed to delete model field: %w", err) } - } else if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModel(request.Model) { + } else if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModelFamily(ctx, request.Model) { reqBody, err := gemini.ToGeminiChatCompletionRequest(request) if err != nil { return nil, err @@ -508,25 +508,25 @@ func (provider *VertexProvider) ChatCompletion(ctx *schemas.BifrostContext, key if bifrostErr != nil { return nil, bifrostErr } - if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModel(request.Model) { + if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModelFamily(ctx, request.Model) { if rawBody, ok := ctx.Value(schemas.BifrostContextKeyUseRawRequestBody).(bool); ok && rawBody { jsonBody = gemini.NormalizeRawGenerateContentRequestForCompatibility(jsonBody) } jsonBody = stripVertexGeminiUnsupportedFieldsRaw(jsonBody) } - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } // Remap unsupported tool versions for Vertex (handles raw passthrough bodies) - if schemas.IsAnthropicModel(request.Model) && jsonBody != nil { + if schemas.IsAnthropicModelFamily(ctx, request.Model) && jsonBody != nil { remappedBody, remapErr := anthropic.RemapRawToolVersionsForProvider(jsonBody, schemas.Vertex, request.Model) if remapErr != nil { return nil, providerUtils.NewBifrostOperationError(remapErr.Error(), nil) @@ -547,7 +547,7 @@ func (provider *VertexProvider) ChatCompletion(ctx *schemas.BifrostContext, key var completeURL string if schemas.IsAllDigitsASCII(request.Model) { // Custom Fine-tuned models use OpenAPI endpoint - projectNumber := key.VertexKeyConfig.ProjectNumber.GetValue() + projectNumber := resolveVertexProjectNumber(ctx, key) if projectNumber == "" { return nil, providerUtils.NewConfigurationError("project number is not set for fine-tuned models") } @@ -555,13 +555,13 @@ func (provider *VertexProvider) ChatCompletion(ctx *schemas.BifrostContext, key authQuery = fmt.Sprintf("key=%s", url.QueryEscape(key.Value.GetValue())) } completeURL = getVertexEndpointURL(region, "v1beta1", projectNumber, request.Model, ":generateContent") - } else if schemas.IsAnthropicModel(request.Model) { + } else if schemas.IsAnthropicModelFamily(ctx, request.Model) { // Claude models use Anthropic publisher — model-aware host for multi-region support completeURL = getVertexModelAwarePublisherModelURL(region, "v1", projectID, "anthropic", request.Model, ":rawPredict") - } else if schemas.IsMistralModel(request.Model) { + } else if schemas.IsMistralModelFamily(ctx, request.Model) { // Mistral models use mistralai publisher with rawPredict completeURL = getVertexPublisherModelURL(region, "v1", projectID, "mistralai", request.Model, ":rawPredict") - } else if schemas.IsGeminiModel(request.Model) || schemas.IsGemmaModel(request.Model) { + } else if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsGemmaModelFamily(ctx, request.Model) { // Gemini models support api key if key.Value.GetValue() != "" { authQuery = fmt.Sprintf("key=%s", url.QueryEscape(key.Value.GetValue())) @@ -584,7 +584,7 @@ func (provider *VertexProvider) ChatCompletion(ctx *schemas.BifrostContext, key req.Header.SetMethod(http.MethodPost) req.Header.SetContentType("application/json") - if (schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model)) && + if (schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model)) && request.Params != nil && request.Params.ServiceTier != nil { if v := vertexServiceTierHeaderValue(region, request.Model, *request.Params.ServiceTier); v != "" { req.Header.Set(VertexServiceTierHeader, v) @@ -653,7 +653,7 @@ func (provider *VertexProvider) ChatCompletion(ctx *schemas.BifrostContext, key }, nil } - if schemas.IsAnthropicModel(request.Model) { + if schemas.IsAnthropicModelFamily(ctx, request.Model) { // Create response object from pool anthropicResponse := anthropic.AcquireAnthropicMessageResponse() defer anthropic.ReleaseAnthropicMessageResponse(anthropicResponse) @@ -682,7 +682,7 @@ func (provider *VertexProvider) ChatCompletion(ctx *schemas.BifrostContext, key } return response, nil - } else if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModel(request.Model) { + } else if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModelFamily(ctx, request.Model) { geminiResponse := gemini.GenerateContentResponse{} rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, &geminiResponse, jsonBody, providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse)) @@ -734,17 +734,17 @@ func (provider *VertexProvider) ChatCompletion(ctx *schemas.BifrostContext, key // Returns a channel of BifrostStreamChunk objects for streaming results or an error if the request fails. func (provider *VertexProvider) ChatCompletionStream(ctx *schemas.BifrostContext, postHookRunner schemas.PostHookRunner, postHookSpanFinalizer func(context.Context), key schemas.Key, request *schemas.BifrostChatRequest) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { providerName := provider.GetProviderKey() - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } - if schemas.IsAnthropicModel(request.Model) { + if schemas.IsAnthropicModelFamily(ctx, request.Model) { // Use Anthropic-style streaming for Claude models jsonData, bifrostErr := providerUtils.CheckContextAndGetRequestBody( ctx, @@ -859,7 +859,7 @@ func (provider *VertexProvider) ChatCompletionStream(ctx *schemas.BifrostContext provider.logger, postHookSpanFinalizer, ) - } else if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModel(request.Model) { + } else if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModelFamily(ctx, request.Model) { // Use Gemini-style streaming for Gemini models jsonData, bifrostErr := providerUtils.CheckContextAndGetRequestBody( ctx, @@ -888,7 +888,7 @@ func (provider *VertexProvider) ChatCompletionStream(ctx *schemas.BifrostContext } // For custom/fine-tuned models, validate projectNumber is set - projectNumber := key.VertexKeyConfig.ProjectNumber.GetValue() + projectNumber := resolveVertexProjectNumber(ctx, key) if schemas.IsAllDigitsASCII(request.Model) && projectNumber == "" { return nil, providerUtils.NewConfigurationError("project number is not set for fine-tuned models") } @@ -909,7 +909,7 @@ func (provider *VertexProvider) ChatCompletionStream(ctx *schemas.BifrostContext "Cache-Control": "no-cache", } - if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) { + if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) { if _, overridden := provider.networkConfig.ExtraHeaders[VertexServiceTierHeader]; !overridden { if request.Params != nil && request.Params.ServiceTier != nil { if v := vertexServiceTierHeaderValue(region, request.Model, *request.Params.ServiceTier); v != "" { @@ -955,7 +955,7 @@ func (provider *VertexProvider) ChatCompletionStream(ctx *schemas.BifrostContext authQuery := "" // Determine the URL based on model type var completeURL string - if schemas.IsMistralModel(request.Model) { + if schemas.IsMistralModelFamily(ctx, request.Model) { // Mistral models use mistralai publisher with streamRawPredict completeURL = getVertexPublisherModelURL(region, "v1", projectID, "mistralai", request.Model, ":streamRawPredict") } else { @@ -1009,17 +1009,17 @@ func (provider *VertexProvider) ChatCompletionStream(ctx *schemas.BifrostContext // Responses performs a responses request to the Vertex API. func (provider *VertexProvider) Responses(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostResponsesRequest) (*schemas.BifrostResponsesResponse, *schemas.BifrostError) { - if schemas.IsAnthropicModel(request.Model) { + if schemas.IsAnthropicModelFamily(ctx, request.Model) { jsonBody, bifrostErr := getRequestBodyForAnthropicResponses(ctx, request, request.Model, false, false, provider.networkConfig.BetaHeaderOverrides, provider.networkConfig.ExtraHeaders, provider.sendBackRawRequest, provider.sendBackRawResponse) if bifrostErr != nil { return nil, bifrostErr } - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } @@ -1128,7 +1128,7 @@ func (provider *VertexProvider) Responses(ctx *schemas.BifrostContext, key schem } return response, nil - } else if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModel(request.Model) { + } else if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModelFamily(ctx, request.Model) { jsonBody, bifrostErr := providerUtils.CheckContextAndGetRequestBody( ctx, request, @@ -1153,12 +1153,12 @@ func (provider *VertexProvider) Responses(ctx *schemas.BifrostContext, key schem } jsonBody = stripVertexGeminiUnsupportedFieldsRaw(jsonBody) - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } @@ -1169,7 +1169,7 @@ func (provider *VertexProvider) Responses(ctx *schemas.BifrostContext, key schem } // For custom/fine-tuned models, validate projectNumber is set - projectNumber := key.VertexKeyConfig.ProjectNumber.GetValue() + projectNumber := resolveVertexProjectNumber(ctx, key) if schemas.IsAllDigitsASCII(request.Model) && projectNumber == "" { return nil, providerUtils.NewConfigurationError("project number is not set for fine-tuned models") } @@ -1189,7 +1189,7 @@ func (provider *VertexProvider) Responses(ctx *schemas.BifrostContext, key schem req.Header.SetMethod(http.MethodPost) req.Header.SetContentType("application/json") - if (schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model)) && + if (schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model)) && request.Params != nil && request.Params.ServiceTier != nil { if v := vertexServiceTierHeaderValue(region, request.Model, *request.Params.ServiceTier); v != "" { req.Header.Set(VertexServiceTierHeader, v) @@ -1290,13 +1290,13 @@ func (provider *VertexProvider) Responses(ctx *schemas.BifrostContext, key schem // ResponsesStream performs a streaming responses request to the Vertex API. func (provider *VertexProvider) ResponsesStream(ctx *schemas.BifrostContext, postHookRunner schemas.PostHookRunner, postHookSpanFinalizer func(context.Context), key schemas.Key, request *schemas.BifrostResponsesRequest) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { - if schemas.IsAnthropicModel(request.Model) { - region := key.VertexKeyConfig.Region.GetValue() + if schemas.IsAnthropicModelFamily(ctx, request.Model) { + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } @@ -1344,13 +1344,13 @@ func (provider *VertexProvider) ResponsesStream(ctx *schemas.BifrostContext, pos provider.logger, postHookSpanFinalizer, ) - } else if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModel(request.Model) { - region := key.VertexKeyConfig.Region.GetValue() + } else if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModelFamily(ctx, request.Model) { + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } @@ -1387,7 +1387,7 @@ func (provider *VertexProvider) ResponsesStream(ctx *schemas.BifrostContext, pos } // For custom/fine-tuned models, validate projectNumber is set - projectNumber := key.VertexKeyConfig.ProjectNumber.GetValue() + projectNumber := resolveVertexProjectNumber(ctx, key) if schemas.IsAllDigitsASCII(request.Model) && projectNumber == "" { return nil, providerUtils.NewConfigurationError("project number is not set for fine-tuned models") } @@ -1407,7 +1407,7 @@ func (provider *VertexProvider) ResponsesStream(ctx *schemas.BifrostContext, pos "Cache-Control": "no-cache", } - if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) { + if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) { if _, overridden := provider.networkConfig.ExtraHeaders[VertexServiceTierHeader]; !overridden { if request.Params != nil && request.Params.ServiceTier != nil { if v := vertexServiceTierHeaderValue(region, request.Model, *request.Params.ServiceTier); v != "" { @@ -1463,12 +1463,12 @@ func (provider *VertexProvider) ResponsesStream(ctx *schemas.BifrostContext, pos // All Vertex AI embedding models use the same response format regardless of the model type. // Returns a BifrostResponse containing the embedding(s) and any error that occurred. func (provider *VertexProvider) Embedding(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostEmbeddingRequest) (*schemas.BifrostEmbeddingResponse, *schemas.BifrostError) { - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } @@ -1485,7 +1485,7 @@ func (provider *VertexProvider) Embedding(ctx *schemas.BifrostContext, key schem } // For custom/fine-tuned models, validate projectNumber is set - projectNumber := key.VertexKeyConfig.ProjectNumber.GetValue() + projectNumber := resolveVertexProjectNumber(ctx, key) if schemas.IsAllDigitsASCII(request.Model) && projectNumber == "" { return nil, providerUtils.NewConfigurationError("project number is not set for fine-tuned models") } @@ -1617,7 +1617,7 @@ func (provider *VertexProvider) Speech(ctx *schemas.BifrostContext, key schemas. // Rerank performs a rerank request using Vertex Discovery Engine ranking API. func (provider *VertexProvider) Rerank(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostRerankRequest) (*schemas.BifrostRerankResponse, *schemas.BifrostError) { - projectID := strings.TrimSpace(key.VertexKeyConfig.ProjectID.GetValue()) + projectID := strings.TrimSpace(resolveVertexProjectID(ctx, key)) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } @@ -1771,7 +1771,7 @@ func (provider *VertexProvider) TranscriptionStream(ctx *schemas.BifrostContext, func (provider *VertexProvider) ImageGeneration(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostImageGenerationRequest) (*schemas.BifrostImageGenerationResponse, *schemas.BifrostError) { // Validate model type before processing - if !schemas.IsGeminiModel(request.Model) && !schemas.IsAllDigitsASCII(request.Model) && !schemas.IsImagenModel(request.Model) { + if !schemas.IsGeminiModelFamily(ctx, request.Model) && !schemas.IsAllDigitsASCII(request.Model) && !schemas.IsImagenModelFamily(ctx, request.Model) { return nil, providerUtils.NewConfigurationError(fmt.Sprintf("image generation is only supported for Gemini and Imagen models, got: %s", request.Model)) } @@ -1783,7 +1783,7 @@ func (provider *VertexProvider) ImageGeneration(ctx *schemas.BifrostContext, key var extraParams map[string]interface{} var err error - if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) { + if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) { reqBody := gemini.ToGeminiImageGenerationRequest(request) if reqBody == nil { return nil, fmt.Errorf("image generation input is not provided") @@ -1796,7 +1796,7 @@ func (provider *VertexProvider) ImageGeneration(ctx *schemas.BifrostContext, key if err != nil { return nil, fmt.Errorf("failed to marshal request body: %w", err) } - } else if schemas.IsImagenModel(request.Model) { + } else if schemas.IsImagenModelFamily(ctx, request.Model) { reqBody := gemini.ToImagenImageGenerationRequest(request) if reqBody == nil { return nil, fmt.Errorf("image generation input is not provided") @@ -1821,12 +1821,12 @@ func (provider *VertexProvider) ImageGeneration(ctx *schemas.BifrostContext, key return nil, bifrostErr } - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } @@ -1837,7 +1837,7 @@ func (provider *VertexProvider) ImageGeneration(ctx *schemas.BifrostContext, key var completeURL string if schemas.IsAllDigitsASCII(request.Model) { // Custom Fine-tuned models use OpenAPI endpoint - projectNumber := key.VertexKeyConfig.ProjectNumber.GetValue() + projectNumber := resolveVertexProjectNumber(ctx, key) if projectNumber == "" { return nil, providerUtils.NewConfigurationError("project number is not set for fine-tuned models") } @@ -1846,13 +1846,13 @@ func (provider *VertexProvider) ImageGeneration(ctx *schemas.BifrostContext, key } completeURL = getVertexEndpointURL(region, "v1beta1", projectNumber, request.Model, ":generateContent") - } else if schemas.IsImagenModel(request.Model) { + } else if schemas.IsImagenModelFamily(ctx, request.Model) { // Imagen models are published models, use publishers/google/models path if value := key.Value.GetValue(); value != "" { authQuery = fmt.Sprintf("key=%s", url.QueryEscape(value)) } completeURL = getVertexPublisherModelURL(region, "v1", projectID, "google", gemini.NormalizeModelName(request.Model), ":predict") - } else if schemas.IsGeminiModel(request.Model) { + } else if schemas.IsGeminiModelFamily(ctx, request.Model) { if value := key.Value.GetValue(); value != "" { authQuery = fmt.Sprintf("key=%s", url.QueryEscape(value)) } @@ -1932,7 +1932,7 @@ func (provider *VertexProvider) ImageGeneration(ctx *schemas.BifrostContext, key }, nil } - if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) { + if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) { geminiResponse := gemini.GenerateContentResponse{} rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, &geminiResponse, jsonBody, providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse)) @@ -1991,7 +1991,7 @@ func (provider *VertexProvider) ImageGenerationStream(ctx *schemas.BifrostContex // Returns a BifrostResponse containing the images and any error that occurred. func (provider *VertexProvider) ImageEdit(ctx *schemas.BifrostContext, key schemas.Key, request *schemas.BifrostImageEditRequest) (*schemas.BifrostImageGenerationResponse, *schemas.BifrostError) { // Validate model type before processing - if !schemas.IsGeminiModel(request.Model) && !schemas.IsAllDigitsASCII(request.Model) && !schemas.IsImagenModel(request.Model) { + if !schemas.IsGeminiModelFamily(ctx, request.Model) && !schemas.IsAllDigitsASCII(request.Model) && !schemas.IsImagenModelFamily(ctx, request.Model) { return nil, providerUtils.NewConfigurationError(fmt.Sprintf("image edit is only supported for Gemini and Imagen models, got: %s", request.Model)) } @@ -2003,7 +2003,7 @@ func (provider *VertexProvider) ImageEdit(ctx *schemas.BifrostContext, key schem var extraParams map[string]interface{} var err error - if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) { + if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) { reqBody := gemini.ToGeminiImageEditRequest(request) if reqBody == nil { return nil, fmt.Errorf("image edit input is not provided") @@ -2016,7 +2016,7 @@ func (provider *VertexProvider) ImageEdit(ctx *schemas.BifrostContext, key schem if err != nil { return nil, fmt.Errorf("failed to marshal request body: %w", err) } - } else if schemas.IsImagenModel(request.Model) { + } else if schemas.IsImagenModelFamily(ctx, request.Model) { reqBody := gemini.ToImagenImageEditRequest(request) if reqBody == nil { return nil, fmt.Errorf("image edit input is not provided") @@ -2041,12 +2041,12 @@ func (provider *VertexProvider) ImageEdit(ctx *schemas.BifrostContext, key schem return nil, bifrostErr } - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } @@ -2058,14 +2058,14 @@ func (provider *VertexProvider) ImageEdit(ctx *schemas.BifrostContext, key schem var completeURL string if schemas.IsAllDigitsASCII(request.Model) { - projectNumber := key.VertexKeyConfig.ProjectNumber.GetValue() + projectNumber := resolveVertexProjectNumber(ctx, key) if projectNumber == "" { return nil, providerUtils.NewConfigurationError("project number is not set for fine-tuned models") } completeURL = getVertexEndpointURL(region, "v1beta1", projectNumber, gemini.NormalizeModelName(request.Model), ":generateContent") - } else if schemas.IsImagenModel(request.Model) { + } else if schemas.IsImagenModelFamily(ctx, request.Model) { completeURL = getVertexPublisherModelURL(region, "v1", projectID, "google", gemini.NormalizeModelName(request.Model), ":predict") - } else if schemas.IsGeminiModel(request.Model) { + } else if schemas.IsGeminiModelFamily(ctx, request.Model) { completeURL = getVertexPublisherModelURL(region, "v1", projectID, "google", gemini.NormalizeModelName(request.Model), ":generateContent") } @@ -2140,7 +2140,7 @@ func (provider *VertexProvider) ImageEdit(ctx *schemas.BifrostContext, key schem }, nil } - if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) { + if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) { geminiResponse := gemini.GenerateContentResponse{} rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, &geminiResponse, jsonBody, providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse)) @@ -2206,7 +2206,7 @@ func (provider *VertexProvider) VideoGeneration(ctx *schemas.BifrostContext, key providerName := provider.GetProviderKey() // Only Gemini models support video generation in Vertex - if !schemas.IsVeoModel(bifrostReq.Model) && !schemas.IsAllDigitsASCII(bifrostReq.Model) { + if !schemas.IsVeoModelFamily(ctx, bifrostReq.Model) && !schemas.IsAllDigitsASCII(bifrostReq.Model) { return nil, providerUtils.NewConfigurationError(fmt.Sprintf("video generation is only supported for Veo models in Vertex, got: %s", bifrostReq.Model)) } @@ -2222,12 +2222,12 @@ func (provider *VertexProvider) VideoGeneration(ctx *schemas.BifrostContext, key return nil, bifrostErr } - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } @@ -2239,7 +2239,7 @@ func (provider *VertexProvider) VideoGeneration(ctx *schemas.BifrostContext, key } // For custom/fine-tuned models, validate projectNumber is set - projectNumber := key.VertexKeyConfig.ProjectNumber.GetValue() + projectNumber := resolveVertexProjectNumber(ctx, key) if schemas.IsAllDigitsASCII(bifrostReq.Model) && projectNumber == "" { return nil, providerUtils.NewConfigurationError("project number is not set for fine-tuned models") } @@ -2328,7 +2328,7 @@ func (provider *VertexProvider) VideoRetrieve(ctx *schemas.BifrostContext, key s sendBackRawResponse := providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse) sendBackRawRequest := providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest) - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } @@ -2682,7 +2682,7 @@ func (provider *VertexProvider) CountTokens(ctx *schemas.BifrostContext, key sch bifrostErr *schemas.BifrostError ) - if schemas.IsAnthropicModel(request.Model) { + if schemas.IsAnthropicModelFamily(ctx, request.Model) { jsonBody, bifrostErr = getRequestBodyForAnthropicResponses(ctx, request, request.Model, false, true, provider.networkConfig.BetaHeaderOverrides, provider.networkConfig.ExtraHeaders, provider.sendBackRawRequest, provider.sendBackRawResponse) if bifrostErr != nil { return nil, bifrostErr @@ -2713,12 +2713,12 @@ func (provider *VertexProvider) CountTokens(ctx *schemas.BifrostContext, key sch } } - projectID := key.VertexKeyConfig.ProjectID.GetValue() + projectID := resolveVertexProjectID(ctx, key) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } - region := key.VertexKeyConfig.Region.GetValue() + region := resolveVertexRegion(ctx, key) if region == "" { return nil, providerUtils.NewConfigurationError("region is not set in key config") } @@ -2726,17 +2726,17 @@ func (provider *VertexProvider) CountTokens(ctx *schemas.BifrostContext, key sch authQuery := "" var completeURL string - if schemas.IsAnthropicModel(request.Model) { + if schemas.IsAnthropicModelFamily(ctx, request.Model) { // Use model-aware host based on request.Model, but URL path uses "count-tokens" effectiveRegion := getVertexEffectiveRegion(region, request.Model) baseURL := getVertexModelAwareAPIBaseURL(region, "v1", request.Model) completeURL = fmt.Sprintf("%s/projects/%s/locations/%s/publishers/%s/models/%s%s", baseURL, projectID, effectiveRegion, "anthropic", "count-tokens", ":rawPredict") - } else if schemas.IsGeminiModel(request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModel(request.Model) { + } else if schemas.IsGeminiModelFamily(ctx, request.Model) || schemas.IsAllDigitsASCII(request.Model) || schemas.IsGemmaModelFamily(ctx, request.Model) { if key.Value.GetValue() != "" { authQuery = fmt.Sprintf("key=%s", url.QueryEscape(key.Value.GetValue())) } - projectNumber := key.VertexKeyConfig.ProjectNumber.GetValue() + projectNumber := resolveVertexProjectNumber(ctx, key) if schemas.IsAllDigitsASCII(request.Model) && projectNumber == "" { return nil, providerUtils.NewConfigurationError("project number is not set for fine-tuned models") } @@ -2816,7 +2816,7 @@ func (provider *VertexProvider) CountTokens(ctx *schemas.BifrostContext, key sch }, nil } - if schemas.IsAnthropicModel(request.Model) { + if schemas.IsAnthropicModelFamily(ctx, request.Model) { anthropicResponse := &anthropic.AnthropicCountTokensResponse{} rawRequest, rawResponse, bifrostErr := providerUtils.HandleProviderResponse(responseBody, anthropicResponse, jsonBody, providerUtils.ShouldSendBackRawRequest(ctx, provider.sendBackRawRequest), providerUtils.ShouldSendBackRawResponse(ctx, provider.sendBackRawResponse)) @@ -2916,12 +2916,12 @@ func (provider *VertexProvider) Passthrough( key schemas.Key, req *schemas.BifrostPassthroughRequest, ) (*schemas.BifrostPassthroughResponse, *schemas.BifrostError) { - projectID := strings.TrimSpace(key.VertexKeyConfig.ProjectID.GetValue()) + projectID := strings.TrimSpace(resolveVertexProjectID(ctx, key)) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } - keyRegion := key.VertexKeyConfig.Region.GetValue() + keyRegion := resolveVertexRegion(ctx, key) if keyRegion == "" { keyRegion = "global" } @@ -3065,12 +3065,12 @@ func (provider *VertexProvider) PassthroughStream( key schemas.Key, req *schemas.BifrostPassthroughRequest, ) (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) { - projectID := strings.TrimSpace(key.VertexKeyConfig.ProjectID.GetValue()) + projectID := strings.TrimSpace(resolveVertexProjectID(ctx, key)) if projectID == "" { return nil, providerUtils.NewConfigurationError("project ID is not set") } - keyRegion := key.VertexKeyConfig.Region.GetValue() + keyRegion := resolveVertexRegion(ctx, key) if keyRegion == "" { keyRegion = "global" } diff --git a/core/schemas/account.go b/core/schemas/account.go index a23f57ce74..c81117d597 100644 --- a/core/schemas/account.go +++ b/core/schemas/account.go @@ -156,7 +156,10 @@ const ( ModelFamilyMistral ModelFamily = "mistral" ModelFamilyCohere ModelFamily = "cohere" ModelFamilyGemini ModelFamily = "gemini" + ModelFamilyGemma ModelFamily = "gemma" ModelFamilyLlama ModelFamily = "llama" + ModelFamilyImagen ModelFamily = "imagen" + ModelFamilyVeo ModelFamily = "veo" ModelFamilyNova ModelFamily = "nova" ModelFamilyTitan ModelFamily = "titan" ) @@ -168,7 +171,8 @@ func (mf *ModelFamily) IsValid() bool { } switch *mf { case ModelFamilyAnthropic, ModelFamilyOpenAI, ModelFamilyMistral, - ModelFamilyCohere, ModelFamilyGemini, ModelFamilyLlama, + ModelFamilyCohere, ModelFamilyGemini, ModelFamilyGemma, + ModelFamilyLlama, ModelFamilyImagen, ModelFamilyVeo, ModelFamilyNova, ModelFamilyTitan: return true } @@ -359,8 +363,18 @@ func ResolveFamily(ctx *BifrostContext, fallbackModel string) ModelFamily { return ModelFamilyAnthropic case IsMistralModel(s): return ModelFamilyMistral + // Imagen and Veo are checked before Gemini as a defensive ordering: + // they are distinct Google model families whose names do not contain + // "gemini", so they could never be mis-classified here, but keeping + // them first makes the intent explicit. + case IsImagenModel(s): + return ModelFamilyImagen + case IsVeoModel(s): + return ModelFamilyVeo case IsGeminiModel(s): return ModelFamilyGemini + case IsGemmaModel(s): + return ModelFamilyGemma case IsLlamaModel(s): return ModelFamilyLlama case IsNovaModel(s): @@ -418,6 +432,33 @@ func IsTitanModelFamily(ctx *BifrostContext, model string) bool { return ResolveFamily(ctx, model) == ModelFamilyTitan } +// IsGeminiModelFamily reports whether the current attempt resolves to the +// Google Gemini model family. Used by Vertex to pick Gemini-shaped request +// transforms and the publishers/google URL prefix. +func IsGeminiModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyGemini +} + +// IsGemmaModelFamily reports whether the current attempt resolves to the +// Gemma model family. Vertex routes Gemma via the publishers/google path +// alongside Gemini. +func IsGemmaModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyGemma +} + +// IsImagenModelFamily reports whether the current attempt resolves to the +// Imagen model family. Used by Vertex for the :predict endpoint and Imagen- +// specific request shaping. +func IsImagenModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyImagen +} + +// IsVeoModelFamily reports whether the current attempt resolves to the Veo +// model family. Used by Vertex for video-generation request shaping. +func IsVeoModelFamily(ctx *BifrostContext, model string) bool { + return ResolveFamily(ctx, model) == ModelFamilyVeo +} + // 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 { From 6d6f96cbb91b878b43b41d361522c5e7889a5011 Mon Sep 17 00:00:00 2001 From: Pratham Mishra <99235987+Pratham-Mishra04@users.noreply.github.com> Date: Tue, 9 Jun 2026 16:02:46 +0530 Subject: [PATCH 011/108] feat: add per-alias `ReplicateAliasCfg.UseDeploymentsEndpoint` override with provider-scoped alias sub-config validation (#4184) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds per-alias control over whether Replicate requests are routed to the deployments endpoint or the predictions endpoint. Previously this was only configurable at the key level via `ReplicateKeyConfig.UseDeploymentsEndpoint`. This change introduces a `ReplicateAliasCfg.UseDeploymentsEndpoint` field that, when set, takes precedence over the key-level setting — allowing a single Replicate API token to route some aliases through the deployments endpoint (e.g. production-pinned models) while others use the predictions endpoint (e.g. experimental versioned models). Additionally, `KeyAliases.Validate` now accepts the owning key's provider and rejects provider-specific sub-configs (Azure, Vertex, Bedrock, Replicate) when attached to a key belonging to a different provider. ## Changes - `useDeploymentsEndpoint` now accepts a `*schemas.BifrostContext` and checks for a `ReplicateAliasCfg.UseDeploymentsEndpoint` override on the resolved alias before falling back to the key-level config. All call sites updated accordingly. - `KeyAliases.Validate` signature changed from `Validate()` to `Validate(providerKey ModelProvider)`. It now returns an error if a provider-specific alias sub-config (e.g. `AzureAliasCfg`) is attached to a key that does not belong to that provider. - Alias validation in `processProvider`, `processAuthoritativeProvider`, and the HTTP handler create/update paths now resolves the effective base provider (accounting for custom provider configs) before calling `Validate`. - The redundant `Validate` call in the GORM `BeforeSave` hook was removed since validation is enforced at the handler layer. - Tests added for `useDeploymentsEndpoint` covering nil context, key-level fallback, alias override true/false, and alias present but without a `ReplicateAliasCfg`. Existing `TestKeyAliasesValidate` extended with provider-mismatch cases for all four provider sub-configs. ## Type of change - [ ] Bug fix - [x] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./core/providers/replicate/... go test ./core/schemas/... go test ./transports/bifrost-http/... ``` To exercise the alias override end-to-end, configure a Replicate key with `use_deployments_endpoint: false` at the key level, then define an alias with `replicate_alias_cfg.use_deployments_endpoint: true`. Requests routed through that alias should target the deployments endpoint while requests without the alias continue to use the predictions endpoint. To verify provider-mismatch validation, attach a `replicate_alias_cfg` to an Azure key in the config file or via the HTTP API and confirm a `400 Bad Request` is returned with a descriptive error. ## Breaking changes - [x] Yes `KeyAliases.Validate()` now requires a `ModelProvider` argument. Any code calling `Validate()` directly must be updated to pass the owning key's provider. ## Security considerations None beyond standard input validation. The new provider-mismatch check prevents misconfigured aliases from silently routing requests to unintended endpoints. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable --- core/providers/replicate/replicate.go | 41 +++++--- core/providers/replicate/replicate_test.go | 1 + .../use_deployments_endpoint_test.go | 76 +++++++++++++++ core/schemas/account.go | 22 ++++- core/schemas/account_test.go | 96 +++++++++++++++---- framework/configstore/tables/key.go | 3 - .../bifrost-http/handlers/provider_keys.go | 4 +- transports/bifrost-http/lib/config.go | 20 +++- 8 files changed, 219 insertions(+), 44 deletions(-) create mode 100644 core/providers/replicate/use_deployments_endpoint_test.go diff --git a/core/providers/replicate/replicate.go b/core/providers/replicate/replicate.go index aa6b75c9df..ee4292ae54 100644 --- a/core/providers/replicate/replicate.go +++ b/core/providers/replicate/replicate.go @@ -90,9 +90,20 @@ const ( pollingInterval = 2 * time.Second ) -// useDeploymentsEndpoint returns whether the key uses the deployments endpoint. -// Nil ReplicateKeyConfig is treated as false (default models/predictions behavior). -func useDeploymentsEndpoint(key schemas.Key) bool { +// useDeploymentsEndpoint returns whether the request should target the +// Replicate deployments endpoint vs the predictions endpoint. +// +// Priority: per-alias ReplicateAliasCfg.UseDeploymentsEndpoint (when set) > +// key-level ReplicateKeyConfig.UseDeploymentsEndpoint. The override lets one +// Replicate API token route some aliases through the deployments endpoint +// (e.g. production-pinned models) while others use the predictions endpoint +// (e.g. experimental versioned models). +// +// Nil ReplicateKeyConfig and missing alias both default to false (predictions). +func useDeploymentsEndpoint(ctx *schemas.BifrostContext, key schemas.Key) bool { + if ra := schemas.GetResolvedAlias(ctx); ra != nil && ra.Config != nil && ra.Config.ReplicateAliasCfg != nil && ra.Config.ReplicateAliasCfg.UseDeploymentsEndpoint != nil { + return *ra.Config.ReplicateAliasCfg.UseDeploymentsEndpoint + } return key.ReplicateKeyConfig != nil && key.ReplicateKeyConfig.UseDeploymentsEndpoint } @@ -284,7 +295,7 @@ func (provider *ReplicateProvider) listDeploymentsByKey(ctx *schemas.BifrostCont client := provider.client extraHeaders := provider.networkConfig.ExtraHeaders - if !useDeploymentsEndpoint(key) { + if !useDeploymentsEndpoint(ctx, key) { return ToBifrostListModelsResponse( &ReplicateDeploymentListResponse{}, providerName, @@ -439,7 +450,7 @@ func (provider *ReplicateProvider) TextCompletion(ctx *schemas.BifrostContext, k request.Model, provider.customProviderConfig, schemas.TextCompletionRequest, - useDeploymentsEndpoint(key), + useDeploymentsEndpoint(ctx, key), ) // create prediction @@ -531,7 +542,7 @@ func (provider *ReplicateProvider) TextCompletionStream(ctx *schemas.BifrostCont request.Model, provider.customProviderConfig, schemas.TextCompletionStreamRequest, - useDeploymentsEndpoint(key), + useDeploymentsEndpoint(ctx, key), ) startTime := time.Now() @@ -779,7 +790,7 @@ func (provider *ReplicateProvider) ChatCompletion(ctx *schemas.BifrostContext, k request.Model, provider.customProviderConfig, schemas.ChatCompletionRequest, - useDeploymentsEndpoint(key), + useDeploymentsEndpoint(ctx, key), ) // create prediction @@ -871,7 +882,7 @@ func (provider *ReplicateProvider) ChatCompletionStream(ctx *schemas.BifrostCont request.Model, provider.customProviderConfig, schemas.ChatCompletionStreamRequest, - useDeploymentsEndpoint(key), + useDeploymentsEndpoint(ctx, key), ) startTime := time.Now() @@ -1136,7 +1147,7 @@ func (provider *ReplicateProvider) Responses(ctx *schemas.BifrostContext, key sc request.Model, provider.customProviderConfig, schemas.ResponsesRequest, - useDeploymentsEndpoint(key), + useDeploymentsEndpoint(ctx, key), ) // create prediction @@ -1223,7 +1234,7 @@ func (provider *ReplicateProvider) ResponsesStream(ctx *schemas.BifrostContext, request.Model, provider.customProviderConfig, schemas.ResponsesStreamRequest, - useDeploymentsEndpoint(key), + useDeploymentsEndpoint(ctx, key), ) startTime := time.Now() @@ -1745,7 +1756,7 @@ func (provider *ReplicateProvider) ImageGeneration(ctx *schemas.BifrostContext, request.Model, provider.customProviderConfig, schemas.ImageGenerationRequest, - useDeploymentsEndpoint(key), + useDeploymentsEndpoint(ctx, key), ) // Create prediction with appropriate mode @@ -1839,7 +1850,7 @@ func (provider *ReplicateProvider) ImageGenerationStream(ctx *schemas.BifrostCon request.Model, provider.customProviderConfig, schemas.ImageGenerationStreamRequest, - useDeploymentsEndpoint(key), + useDeploymentsEndpoint(ctx, key), ) startTime := time.Now() // Create prediction @@ -2150,7 +2161,7 @@ func (provider *ReplicateProvider) ImageEdit(ctx *schemas.BifrostContext, key sc request.Model, provider.customProviderConfig, schemas.ImageEditRequest, - useDeploymentsEndpoint(key), + useDeploymentsEndpoint(ctx, key), ) // Create prediction with appropriate mode @@ -2244,7 +2255,7 @@ func (provider *ReplicateProvider) ImageEditStream(ctx *schemas.BifrostContext, request.Model, provider.customProviderConfig, schemas.ImageEditStreamRequest, - useDeploymentsEndpoint(key), + useDeploymentsEndpoint(ctx, key), ) startTime := time.Now() @@ -2538,7 +2549,7 @@ func (provider *ReplicateProvider) VideoGeneration(ctx *schemas.BifrostContext, request.Model, provider.customProviderConfig, schemas.VideoGenerationRequest, - useDeploymentsEndpoint(key), + useDeploymentsEndpoint(ctx, key), ) // Create prediction with appropriate mode diff --git a/core/providers/replicate/replicate_test.go b/core/providers/replicate/replicate_test.go index 855b9cb690..6d3f4bea55 100644 --- a/core/providers/replicate/replicate_test.go +++ b/core/providers/replicate/replicate_test.go @@ -1438,3 +1438,4 @@ func TestReplicateToBifrostResponsesResponse(t *testing.T) { }) } } + diff --git a/core/providers/replicate/use_deployments_endpoint_test.go b/core/providers/replicate/use_deployments_endpoint_test.go new file mode 100644 index 0000000000..f85b7ae4ff --- /dev/null +++ b/core/providers/replicate/use_deployments_endpoint_test.go @@ -0,0 +1,76 @@ +package replicate + +import ( + "context" + "testing" + + "github.com/maximhq/bifrost/core/schemas" +) + +// TestUseDeploymentsEndpoint_AliasOverride verifies the per-alias +// ReplicateAliasCfg.UseDeploymentsEndpoint override resolves correctly: +// alias value wins when set, else falls through to key-level config. +func TestUseDeploymentsEndpoint_AliasOverride(t *testing.T) { + keyDeployments := schemas.Key{ + ReplicateKeyConfig: &schemas.ReplicateKeyConfig{UseDeploymentsEndpoint: true}, + } + keyPredictions := schemas.Key{ + ReplicateKeyConfig: &schemas.ReplicateKeyConfig{UseDeploymentsEndpoint: false}, + } + + // No alias in ctx — falls back to key-level setting. + if got := useDeploymentsEndpoint(nil, keyDeployments); !got { + t.Errorf("nil ctx + key=deployments: want true, got false") + } + if got := useDeploymentsEndpoint(nil, keyPredictions); got { + t.Errorf("nil ctx + key=predictions: want false, got true") + } + if got := useDeploymentsEndpoint(nil, schemas.Key{}); got { + t.Errorf("nil ctx + nil ReplicateKeyConfig: want false, got true") + } + + // Alias override true wins over key=false. + ctxOverrideTrue := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + defer ctxOverrideTrue.Cancel() + trueVal := true + ctxOverrideTrue.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "prod-llm", + Config: &schemas.AliasConfig{ + ModelID: "owner/name:version", + ReplicateAliasCfg: &schemas.ReplicateAliasCfg{ + UseDeploymentsEndpoint: &trueVal, + }, + }, + }) + if got := useDeploymentsEndpoint(ctxOverrideTrue, keyPredictions); !got { + t.Errorf("alias=true should override key=false: got false") + } + + // Alias override false wins over key=true. + ctxOverrideFalse := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + defer ctxOverrideFalse.Cancel() + falseVal := false + ctxOverrideFalse.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "experimental-llm", + Config: &schemas.AliasConfig{ + ModelID: "owner/name:version", + ReplicateAliasCfg: &schemas.ReplicateAliasCfg{ + UseDeploymentsEndpoint: &falseVal, + }, + }, + }) + if got := useDeploymentsEndpoint(ctxOverrideFalse, keyDeployments); got { + t.Errorf("alias=false should override key=true: got true") + } + + // Alias present but ReplicateAliasCfg unset — falls through to key. + ctxNoCfg := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + defer ctxNoCfg.Cancel() + ctxNoCfg.SetValue(schemas.BifrostContextKeyResolvedAlias, &schemas.ResolvedAlias{ + Key: "x", + Config: &schemas.AliasConfig{ModelID: "x"}, + }) + if got := useDeploymentsEndpoint(ctxNoCfg, keyDeployments); !got { + t.Errorf("no alias cfg + key=deployments: want true, got false") + } +} diff --git a/core/schemas/account.go b/core/schemas/account.go index c81117d597..2c43c851e2 100644 --- a/core/schemas/account.go +++ b/core/schemas/account.go @@ -261,7 +261,15 @@ func (ac AliasConfig) MarshalJSON() ([]byte, error) { // ModelName / ModelFamily / provider sub-configs are populated explicitly. type KeyAliases map[string]AliasConfig -func (ka KeyAliases) Validate() error { +// Validate checks that every entry in the alias map is well-formed and that +// any provider-specific sub-configs (AzureAliasCfg, VertexAliasCfg, +// BedrockAliasCfg, ReplicateAliasCfg) are only set when the owning Key +// actually belongs to that provider. Catches misconfigurations like an +// AzureAliasCfg attached to a Bedrock key. +// +// providerKey is the provider this Key is registered under (e.g. schemas.Azure +// for keys in the azure provider config). +func (ka KeyAliases) Validate(providerKey ModelProvider) error { seen := make(map[string]struct{}, len(ka)) for from, ac := range ka { if strings.TrimSpace(from) == "" { @@ -282,6 +290,18 @@ func (ka KeyAliases) Validate() error { if ac.ModelFamily != nil && !ac.ModelFamily.IsValid() { return fmt.Errorf("alias %q: invalid model_family %q", from, *ac.ModelFamily) } + if ac.AzureAliasCfg != nil && providerKey != Azure { + return fmt.Errorf("alias %q: azure sub-config is only valid on Azure keys (got provider %q)", from, providerKey) + } + if ac.VertexAliasCfg != nil && providerKey != Vertex { + return fmt.Errorf("alias %q: vertex sub-config is only valid on Vertex keys (got provider %q)", from, providerKey) + } + if ac.BedrockAliasCfg != nil && providerKey != Bedrock { + return fmt.Errorf("alias %q: bedrock sub-config is only valid on Bedrock keys (got provider %q)", from, providerKey) + } + if ac.ReplicateAliasCfg != nil && providerKey != Replicate { + return fmt.Errorf("alias %q: replicate sub-config is only valid on Replicate keys (got provider %q)", from, providerKey) + } normalized := strings.ToLower(from) if _, ok := seen[normalized]; ok { return fmt.Errorf("duplicate alias source %q (case-insensitive)", from) diff --git a/core/schemas/account_test.go b/core/schemas/account_test.go index 39634863a4..ecd9653d01 100644 --- a/core/schemas/account_test.go +++ b/core/schemas/account_test.go @@ -125,9 +125,9 @@ func TestKeyAliasesMarshalLegacyShapeWhenOnlyModelIDSet(t *testing.T) { func TestKeyAliasesMarshalRichShapeWhenAnyExtraFieldSet(t *testing.T) { cases := map[string]struct { - ac AliasConfig - wantKey string - wantValue any + ac AliasConfig + wantKey string + wantValue any }{ "with_model_name": {AliasConfig{ModelID: "x", ModelName: Ptr("canonical")}, "model_name", "canonical"}, "with_model_family": {AliasConfig{ModelID: "x", ModelFamily: Ptr(ModelFamilyAnthropic)}, "model_family", "anthropic"}, @@ -222,29 +222,38 @@ func TestKeyAliasesResolveConfig(t *testing.T) { func TestKeyAliasesValidate(t *testing.T) { madeUp := ModelFamily("made-up") + azureCfg := &AzureAliasCfg{APIVersion: Ptr("2024-08-01-preview")} + bedrockCfg := &BedrockAliasCfg{InferenceProfileARN: NewEnvVar("arn:aws:bedrock:...")} + vertexCfg := &VertexAliasCfg{ProjectID: NewEnvVar("my-gcp-project")} + replicateCfg := &ReplicateAliasCfg{UseDeploymentsEndpoint: Ptr(true)} cases := []struct { - name string - ka KeyAliases - wantErr string + name string + provider ModelProvider + ka KeyAliases + wantErr string }{ { - name: "ok", - ka: KeyAliases{"k": {ModelID: "v"}}, + name: "ok", + provider: OpenAI, + ka: KeyAliases{"k": {ModelID: "v"}}, }, { - name: "empty source", - ka: KeyAliases{"": {ModelID: "v"}}, - wantErr: "alias source cannot be empty", + name: "empty source", + provider: OpenAI, + ka: KeyAliases{"": {ModelID: "v"}}, + wantErr: "alias source cannot be empty", }, { - name: "empty model id", - ka: KeyAliases{"k": {ModelID: ""}}, - wantErr: "model_id cannot be empty", + name: "empty model id", + provider: OpenAI, + ka: KeyAliases{"k": {ModelID: ""}}, + wantErr: "model_id cannot be empty", }, { - name: "whitespace source", - ka: KeyAliases{" k ": {ModelID: "v"}}, - wantErr: "leading or trailing whitespace", + name: "whitespace source", + provider: OpenAI, + ka: KeyAliases{" k ": {ModelID: "v"}}, + wantErr: "leading or trailing whitespace", }, { name: "whitespace model_id", @@ -262,14 +271,59 @@ func TestKeyAliasesValidate(t *testing.T) { wantErr: "duplicate alias source", }, { - name: "invalid family", - ka: KeyAliases{"k": {ModelID: "v", ModelFamily: &madeUp}}, - wantErr: "invalid model_family", + name: "invalid family", + provider: OpenAI, + ka: KeyAliases{"k": {ModelID: "v", ModelFamily: &madeUp}}, + wantErr: "invalid model_family", + }, + { + name: "azure sub-config on azure key — ok", + provider: Azure, + ka: KeyAliases{"k": {ModelID: "v", AzureAliasCfg: azureCfg}}, + }, + { + name: "azure sub-config on non-azure key — error", + provider: Bedrock, + ka: KeyAliases{"k": {ModelID: "v", AzureAliasCfg: azureCfg}}, + wantErr: "azure sub-config is only valid on Azure keys", + }, + { + name: "bedrock sub-config on bedrock key — ok", + provider: Bedrock, + ka: KeyAliases{"k": {ModelID: "v", BedrockAliasCfg: bedrockCfg}}, + }, + { + name: "bedrock sub-config on azure key — error", + provider: Azure, + ka: KeyAliases{"k": {ModelID: "v", BedrockAliasCfg: bedrockCfg}}, + wantErr: "bedrock sub-config is only valid on Bedrock keys", + }, + { + name: "vertex sub-config on vertex key — ok", + provider: Vertex, + ka: KeyAliases{"k": {ModelID: "v", VertexAliasCfg: vertexCfg}}, + }, + { + name: "vertex sub-config on openai key — error", + provider: OpenAI, + ka: KeyAliases{"k": {ModelID: "v", VertexAliasCfg: vertexCfg}}, + wantErr: "vertex sub-config is only valid on Vertex keys", + }, + { + name: "replicate sub-config on replicate key — ok", + provider: Replicate, + ka: KeyAliases{"k": {ModelID: "v", ReplicateAliasCfg: replicateCfg}}, + }, + { + name: "replicate sub-config on bedrock key — error", + provider: Bedrock, + ka: KeyAliases{"k": {ModelID: "v", ReplicateAliasCfg: replicateCfg}}, + wantErr: "replicate sub-config is only valid on Replicate keys", }, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { - err := c.ka.Validate() + err := c.ka.Validate(c.provider) if c.wantErr == "" { if err != nil { t.Fatalf("want ok, got %v", err) diff --git a/framework/configstore/tables/key.go b/framework/configstore/tables/key.go index 51df1def45..325f5caed7 100644 --- a/framework/configstore/tables/key.go +++ b/framework/configstore/tables/key.go @@ -276,9 +276,6 @@ func (k *TableKey) BeforeSave(tx *gorm.DB) error { } if k.Aliases != nil { - if err := k.Aliases.Validate(); err != nil { - return err - } data, err := sonic.Marshal(k.Aliases) if err != nil { return err diff --git a/transports/bifrost-http/handlers/provider_keys.go b/transports/bifrost-http/handlers/provider_keys.go index 2e04417bee..cb8614c1eb 100644 --- a/transports/bifrost-http/handlers/provider_keys.go +++ b/transports/bifrost-http/handlers/provider_keys.go @@ -113,7 +113,7 @@ func (h *ProviderHandler) createProviderKey(ctx *fasthttp.RequestCtx) { return } - if err := key.Aliases.Validate(); err != nil { + if err := key.Aliases.Validate(baseProvider); err != nil { SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("Invalid aliases: %v", err)) return } @@ -224,7 +224,7 @@ func (h *ProviderHandler) updateProviderKey(ctx *fasthttp.RequestCtx) { return } - if err := mergedKey.Aliases.Validate(); err != nil { + if err := mergedKey.Aliases.Validate(baseProvider); err != nil { SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("Invalid aliases: %v", err)) return } diff --git a/transports/bifrost-http/lib/config.go b/transports/bifrost-http/lib/config.go index 685a70eba1..4864c5159d 100644 --- a/transports/bifrost-http/lib/config.go +++ b/transports/bifrost-http/lib/config.go @@ -1240,12 +1240,21 @@ func processProvider( ) error { provider := schemas.ModelProvider(strings.ToLower(providerName)) + if err := ValidateCustomProvider(providerCfgInFile, provider); err != nil { + return err + } + + baseProvider := provider + if providerCfgInFile.CustomProviderConfig != nil && providerCfgInFile.CustomProviderConfig.BaseProviderType != "" { + baseProvider = providerCfgInFile.CustomProviderConfig.BaseProviderType + } + // Process environment variables in keys (including key-level configs) for i, providerKeyInFile := range providerCfgInFile.Keys { if providerKeyInFile.ID == "" { providerCfgInFile.Keys[i].ID = uuid.NewString() } - if err := providerKeyInFile.Aliases.Validate(); err != nil { + if err := providerKeyInFile.Aliases.Validate(baseProvider); err != nil { return fmt.Errorf("invalid aliases for key %q in provider %s: %w", providerKeyInFile.Name, provider, err) } } @@ -1269,11 +1278,18 @@ func processAuthoritativeProvider( providers map[schemas.ModelProvider]configstore.ProviderConfig, ) error { provider := schemas.ModelProvider(strings.ToLower(providerName)) + if err := ValidateCustomProvider(providerCfgInFile, provider); err != nil { + return err + } + baseProvider := provider + if providerCfgInFile.CustomProviderConfig != nil && providerCfgInFile.CustomProviderConfig.BaseProviderType != "" { + baseProvider = providerCfgInFile.CustomProviderConfig.BaseProviderType + } for i, providerKeyInFile := range providerCfgInFile.Keys { if providerKeyInFile.ID == "" { providerCfgInFile.Keys[i].ID = uuid.NewString() } - if err := providerKeyInFile.Aliases.Validate(); err != nil { + if err := providerKeyInFile.Aliases.Validate(baseProvider); err != nil { return fmt.Errorf("invalid aliases for key %q in provider %s: %w", providerKeyInFile.Name, provider, err) } } From f2444cf113527b4c0619ccbc3072560689f6f564 Mon Sep 17 00:00:00 2001 From: Pratham Mishra <99235987+Pratham-Mishra04@users.noreply.github.com> Date: Tue, 9 Jun 2026 16:05:46 +0530 Subject: [PATCH 012/108] feat: replace flat aliases table with rich `DeploymentsTable` supporting per-deployment model family, canonical name, and provider overrides (#4185) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Replaces the flat "Aliases" key-value table (mapping request model name → string deployment ID) with a richer "Deployments" table that supports per-deployment metadata and provider-specific overrides. This enables cost/pricing logs and family-based routing to work correctly for custom deployments. ## Changes - Introduced a new `DeploymentsTable` component that renders each deployment as a collapsible row. Expanding a row exposes fields for canonical model name, model family, description, and provider-specific overrides (Azure API version, endpoint, Anthropic version; Vertex project ID/number/region; Bedrock region and inference profile ARN; Replicate deployments endpoint toggle). - Replaced the `normalizeAliasesValue` helper and `HeadersTable`-based aliases editor in `apiKeysFormFragment.tsx` with the new `DeploymentsTable`. The form label and description were updated from "Aliases" to "Deployments" to reflect the richer semantics. - Added `AliasConfig` and `ModelFamily` types to `config.ts`, mirroring the Go `schemas.AliasConfig` struct (with embedded provider sub-configs flattened to top-level fields on the wire). - Added `aliasConfigSchema` and `modelFamilySchema` Zod schemas to `schemas.ts`. The alias schema uses `z.preprocess` to accept the legacy `string` wire shape emitted by the Go server for simple aliases, coercing it to `{ model_id: string }` so hydrated state passes validation without a migration. - Updated `KeySchema` in `providerForm.ts` to use `z.record(z.string(), aliasConfigSchema)` and updated the validation error message to reflect the new requirement. - Rewrote `isValidAliases` in `validation.ts` to validate the rich `Record` shape, checking that every entry has a non-empty deployment name and a non-empty `model_id`. - Updated `ModelProviderKey` in `config.ts` to type `aliases` as `Record` instead of `Record`. - The `DeploymentsTable` includes a draft row at the bottom for adding new entries. The draft is committed automatically when both the deployment name and model ID are filled. Rename collision detection is case-insensitive and stable row IDs are used to preserve expanded/pending state across renames. ## Type of change - [ ] Bug fix - [x] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [ ] Core (Go) - [ ] Transports (HTTP) - [ ] Providers/Integrations - [ ] Plugins - [x] UI (React) - [ ] Docs ## How to test ```sh cd ui pnpm i pnpm build ``` 1. Navigate to a provider key configuration form. 2. Verify the "Deployments" table renders in place of the old "Aliases" table. 3. Add a new deployment by filling in the deployment name and model ID in the draft row — confirm it commits automatically when both fields are populated. 4. Expand a committed row and verify the canonical model name, model family, description, and provider-specific override fields are visible and editable. 5. For Azure, Vertex, Bedrock, and Replicate providers, confirm the correct provider-specific section appears in the expanded panel. 6. Rename a deployment to an existing name and confirm the inline collision error appears and the row is not committed. 7. Load an existing config that uses the legacy `Record` alias format and confirm it hydrates correctly into the new table without validation errors. ## Screenshots/Recordings Before: A simple two-column key/value table labeled "Aliases" with a plain text input for the deployment ID. After: A collapsible table labeled "Deployments" where each row can be expanded to reveal canonical model name, model family, description, and provider-specific override fields. ## Breaking changes - [x] Yes - [ ] No The `aliases` field type changes from `Record` to `Record` in the UI type system and form schema. Existing configs using the legacy string format are handled transparently via the `aliasConfigSchema` preprocessor and the `normalize` function in `DeploymentsTable`, so no manual migration is required for stored configs. Any code outside this diff that directly constructs or reads `ModelProviderKey.aliases` as `Record` will need to be updated. ## Related issues ## Security considerations No new secrets or auth surfaces introduced. Provider-specific override fields (endpoint, credentials) use the existing `EnvVarInput` component, which supports environment variable references and redaction consistent with the rest of the form. ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [ ] I added/updated tests where appropriate - [ ] I updated documentation where needed - [ ] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable --- .../fragments/apiKeysFormFragment.tsx | 58 +- .../providers/fragments/deploymentsTable.tsx | 584 ++++++++++++++++++ ui/lib/schemas/providerForm.ts | 6 +- ui/lib/types/config.ts | 53 +- ui/lib/types/schemas.ts | 47 +- ui/lib/utils/validation.ts | 37 +- 6 files changed, 708 insertions(+), 77 deletions(-) create mode 100644 ui/app/workspace/providers/fragments/deploymentsTable.tsx diff --git a/ui/app/workspace/providers/fragments/apiKeysFormFragment.tsx b/ui/app/workspace/providers/fragments/apiKeysFormFragment.tsx index d841c27071..ec2c2d531d 100644 --- a/ui/app/workspace/providers/fragments/apiKeysFormFragment.tsx +++ b/ui/app/workspace/providers/fragments/apiKeysFormFragment.tsx @@ -1,6 +1,5 @@ import { EnvVarInput } from "@/components/ui/envVarInput"; import { FormControl, FormDescription, FormField, FormItem, FormLabel, FormMessage } from "@/components/ui/form"; -import { HeadersTable, type CellRenderParams } from "@/components/ui/headersTable"; import { Input } from "@/components/ui/input"; import { ModelMultiselect } from "@/components/ui/modelMultiselect"; import { Separator } from "@/components/ui/separator"; @@ -12,36 +11,11 @@ import { isRedacted } from "@/lib/utils/validation"; import { Info } from "lucide-react"; import { useEffect, useState } from "react"; import { Control, UseFormReturn } from "react-hook-form"; +import { DeploymentsTable } from "./deploymentsTable"; // Providers that support batch APIs const BATCH_SUPPORTED_PROVIDERS = ["openai", "bedrock", "anthropic", "gemini", "azure"]; -/** Normalize form value (object or legacy JSON string) for the alias map editor. */ -function normalizeAliasesValue(v: Record | string | undefined | null): Record { - if (v == null) { - return {}; - } - if (typeof v === "string") { - const t = v.trim(); - if (!t) { - return {}; - } - try { - const p = JSON.parse(t) as unknown; - if (typeof p === "object" && p !== null && !Array.isArray(p)) { - return Object.fromEntries(Object.entries(p as Record).map(([k, val]) => [k, String(val ?? "")])); - } - } catch { - return {}; - } - return {}; - } - if (typeof v === "object" && !Array.isArray(v)) { - return Object.fromEntries(Object.entries(v).map(([k, val]) => [k, typeof val === "string" ? val : String(val ?? "")])); - } - return {}; -} - interface Props { control: Control; providerName: string; @@ -346,34 +320,22 @@ export function ApiKeyFormFragment({ control, providerName, form }: Props) { control={control} name={`key.aliases`} render={({ field }) => ( - - Aliases (Optional) + + Deployments (Optional) - Map each request model name to the provider's identifier (deployment name, inference profile ID, fine-tuned endpoint - ID, etc.) or just a custom name, e.g. "claude-sonnet-4-5" -> "custom-claude-4.5-sonnet". + Map a request model name to the provider's identifier (deployment name, inference profile ID, fine-tuned endpoint + ID, etc.). Expand a row to set the canonical model name, model family, and provider-specific overrides — these power + cost/pricing logs and family-based routing. -
- + { form.clearErrors("key.aliases"); field.onChange(Object.keys(next).length > 0 ? next : {}); }} - keyPlaceholder="Request model name" - valuePlaceholder="Deployment / profile / resource ID" - renderValueInput={({ value: cellValue, onChange, placeholder, disabled }: CellRenderParams) => ( - - )} />
diff --git a/ui/app/workspace/providers/fragments/deploymentsTable.tsx b/ui/app/workspace/providers/fragments/deploymentsTable.tsx new file mode 100644 index 0000000000..e7420dab27 --- /dev/null +++ b/ui/app/workspace/providers/fragments/deploymentsTable.tsx @@ -0,0 +1,584 @@ +import { Button } from "@/components/ui/button"; +import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; +import { EnvVarInput } from "@/components/ui/envVarInput"; +import { Input } from "@/components/ui/input"; +import { ModelMultiselect } from "@/components/ui/modelMultiselect"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Switch } from "@/components/ui/switch"; +import { Textarea } from "@/components/ui/textarea"; +import { AliasConfig, ModelFamily, ModelFamilyValues } from "@/lib/types/config"; +import { EnvVar } from "@/lib/types/schemas"; +import { cn } from "@/lib/utils"; +import { ChevronDown, ChevronRight, Trash } from "lucide-react"; +import { useMemo, useRef, useState } from "react"; + +type DeploymentsValue = Record | undefined | null; + +interface Props { + value: DeploymentsValue; + onChange: (next: Record) => void; + providerName: string; + disabled?: boolean; +} + +interface Row { + name: string; + config: AliasConfig; +} + +// Normalize legacy shapes (Record from older configs or stringified JSON) +// into the rich Record the component operates on. +function normalize(value: DeploymentsValue): Record { + if (value == null) { + return {}; + } + if (typeof value === "string") { + try { + const parsed = JSON.parse(value); + return normalize(parsed); + } catch { + return {}; + } + } + if (typeof value !== "object" || Array.isArray(value)) { + return {}; + } + const out: Record = {}; + for (const [k, v] of Object.entries(value)) { + if (typeof v === "string") { + out[k] = { model_id: v }; + } else if (v && typeof v === "object") { + const cfg = v as Partial; + out[k] = { ...cfg, model_id: typeof cfg.model_id === "string" ? cfg.model_id : "" }; + } + } + return out; +} + +const emptyEnvVar: EnvVar = { value: "", env_var: "", from_env: false }; +const isEmptyEnvVar = (v: EnvVar | undefined): boolean => !v || (!v.value && !v.env_var); + +function FieldRow({ label, hint, children }: { label: string; hint?: string; children: React.ReactNode }) { + return ( +
+ + {children} + {hint &&

{hint}

} +
+ ); +} + +function SectionHeader({ title, description }: { title: string; description?: string }) { + return ( +
+

{title}

+ {description &&

{description}

} +
+ ); +} + +function EnvVarField({ + value, + onChange, + placeholder, + disabled, +}: { + value: EnvVar | undefined; + onChange: (next: EnvVar | undefined) => void; + placeholder?: string; + disabled?: boolean; +}) { + return ( + onChange(isEmptyEnvVar(next) ? undefined : next)} + placeholder={placeholder} + disabled={disabled} + /> + ); +} + +function StringField({ + value, + onChange, + placeholder, + disabled, +}: { + value: string | undefined; + onChange: (next: string | undefined) => void; + placeholder?: string; + disabled?: boolean; +}) { + return ( + onChange(e.target.value === "" ? undefined : e.target.value)} + placeholder={placeholder} + disabled={disabled} + /> + ); +} + +interface ProviderSectionProps { + config: AliasConfig; + onChange: (patch: Partial) => void; + disabled?: boolean; +} + +function AzureSection({ config, onChange, disabled }: ProviderSectionProps) { + return ( +
+ + + onChange({ api_version: v })} + placeholder="2024-10-21" + disabled={disabled} + /> + + + onChange({ anthropic_version: v })} + placeholder="2023-06-01" + disabled={disabled} + /> + + + onChange({ endpoint: v })} + placeholder="https://your-resource.openai.azure.com or env.AZURE_ENDPOINT" + disabled={disabled} + /> + +
+ ); +} + +function VertexSection({ config, onChange, disabled }: ProviderSectionProps) { + return ( +
+ + + onChange({ project_id: v })} + placeholder="gcp-project-id or env.VERTEX_PROJECT_ID" + disabled={disabled} + /> + + + onChange({ project_number: v })} + placeholder="123456789 or env.VERTEX_PROJECT_NUMBER" + disabled={disabled} + /> + + + onChange({ region: v })} + placeholder="us-central1 or env.VERTEX_REGION" + disabled={disabled} + /> + +
+ ); +} + +function BedrockSection({ config, onChange, disabled }: ProviderSectionProps) { + return ( +
+ + + onChange({ region: v })} + placeholder="us-east-1 or env.BEDROCK_REGION" + disabled={disabled} + /> + + + onChange({ inference_profile_arn: v })} + placeholder="arn:aws:bedrock:us-east-1:123:inference-profile/... or env.BEDROCK_PROFILE_ARN" + disabled={disabled} + /> + +
+ ); +} + +function ReplicateSection({ config, onChange, disabled }: ProviderSectionProps) { + return ( +
+ +
+
+ +

+ Route through Replicate's deployments endpoint instead of the models endpoint. +

+
+ onChange({ use_deployments_endpoint: checked ? true : undefined })} + disabled={disabled} + /> +
+
+ ); +} + +function ProviderSection({ providerName, ...props }: ProviderSectionProps & { providerName: string }) { + switch (providerName) { + case "azure": + return ; + case "vertex": + return ; + case "bedrock": + return ; + case "replicate": + return ; + default: + return null; + } +} + +function ExpandedConfigPanel({ + config, + onChange, + providerName, + disabled, +}: { + config: AliasConfig; + onChange: (patch: Partial) => void; + providerName: string; + disabled?: boolean; +}) { + return ( +
+
+ + onChange({ model_name: v })} + placeholder="e.g. claude-sonnet-4-5" + disabled={disabled} + /> + + + + + +