feat: add per-alias Replicate use deployments override and provider-scoped alias sub-config validation

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

core/providers/replicate/replicate_test.go

@@ -1438,3 +1438,4 @@ func TestReplicateToBifrostResponsesResponse(t *testing.T) {
 		})
 	}
 }
+

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")
+	}
+}

core/schemas/account.go

@@ -265,7 +265,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) == "" {
@@ -286,6 +294,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)

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)

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