fix: base provider level request level restriction

TejasGhatte · TejasGhatte · commit 6d784ea093a3 · 2025-09-01T12:54:14.000+05:30
diff --git a/core/schemas/provider.go b/core/schemas/provider.go
@@ -123,6 +123,117 @@ func (ar *AllowedRequests) IsOperationAllowed(operation Operation) bool {
 	}
 }
 
+// ProviderCapabilities defines what operations each provider supports
+type ProviderCapabilities struct {
+	TextCompletion       bool
+	ChatCompletion       bool
+	ChatCompletionStream bool
+	Embedding            bool
+	Speech               bool
+	SpeechStream         bool
+	Transcription        bool
+	TranscriptionStream  bool
+}
+
+// GetProviderCapabilities returns the capabilities for a given provider
+func GetProviderCapabilities(provider ModelProvider) ProviderCapabilities {
+	switch provider {
+	case OpenAI:
+		return ProviderCapabilities{
+			TextCompletion:       false,
+			ChatCompletion:       true,
+			ChatCompletionStream: true,
+			Embedding:            true,
+			Speech:               true,
+			SpeechStream:         true,
+			Transcription:        true,
+			TranscriptionStream:  true,
+		}
+	case Anthropic:
+		return ProviderCapabilities{
+			TextCompletion:       true,
+			ChatCompletion:       true,
+			ChatCompletionStream: true,
+			Embedding:            false,
+			Speech:               false,
+			SpeechStream:         false,
+			Transcription:        false,
+			TranscriptionStream:  false,
+		}
+	case Cohere:
+		return ProviderCapabilities{
+			TextCompletion:       false,
+			ChatCompletion:       true,
+			ChatCompletionStream: true,
+			Embedding:            true,
+			Speech:               false,
+			SpeechStream:         false,
+			Transcription:        false,
+			TranscriptionStream:  false,
+		}
+	case Bedrock:
+		return ProviderCapabilities{
+			TextCompletion:       true,
+			ChatCompletion:       true,
+			ChatCompletionStream: true,
+			Embedding:            true,
+			Speech:               false,
+			SpeechStream:         false,
+			Transcription:        false,
+			TranscriptionStream:  false,
+		}
+	default:
+		// For unknown providers, assume all capabilities (fallback)
+		return ProviderCapabilities{
+			TextCompletion:       true,
+			ChatCompletion:       true,
+			ChatCompletionStream: true,
+			Embedding:            true,
+			Speech:               true,
+			SpeechStream:         true,
+			Transcription:        true,
+			TranscriptionStream:  true,
+		}
+	}
+}
+
+// ValidateAllowedRequests checks if the allowed requests are valid for the given base provider
+// and silently sets any unsupported operations to false instead of returning an error
+func ValidateAllowedRequests(allowedRequests *AllowedRequests, baseProvider ModelProvider) {
+	if allowedRequests == nil {
+		// nil means all operations allowed, which is always valid
+		return
+	}
+
+	capabilities := GetProviderCapabilities(baseProvider)
+
+	// Silently set unsupported operations to false instead of returning an error
+	if allowedRequests.TextCompletion && !capabilities.TextCompletion {
+		allowedRequests.TextCompletion = false
+	}
+	if allowedRequests.ChatCompletion && !capabilities.ChatCompletion {
+		allowedRequests.ChatCompletion = false
+	}
+	if allowedRequests.ChatCompletionStream && !capabilities.ChatCompletionStream {
+		allowedRequests.ChatCompletionStream = false
+	}
+	if allowedRequests.Embedding && !capabilities.Embedding {
+		allowedRequests.Embedding = false
+	}
+	if allowedRequests.Speech && !capabilities.Speech {
+		allowedRequests.Speech = false
+	}
+	if allowedRequests.SpeechStream && !capabilities.SpeechStream {
+		allowedRequests.SpeechStream = false
+	}
+	if allowedRequests.Transcription && !capabilities.Transcription {
+		allowedRequests.Transcription = false
+	}
+	if allowedRequests.TranscriptionStream && !capabilities.TranscriptionStream {
+		allowedRequests.TranscriptionStream = false
+	}
+}
+
 type CustomProviderConfig struct {
 	CustomProviderKey string           `json:"-"`                  // Custom provider key, internally set by Bifrost
 	BaseProviderType  ModelProvider    `json:"base_provider_type"` // Base provider type
diff --git a/docs/features/custom-providers.mdx b/docs/features/custom-providers.mdx
@@ -107,6 +107,73 @@ curl --location 'http://localhost:8080/api/providers' \
 }
 ```
 
+</Tab>
+
+<Tab title="Using GO SDK">
+
+```go
+type MyAccount struct{}
+
+OpenAICustom = schemas.ModelProvider("openai-custom")
+
+func (a *MyAccount) GetConfiguredProviders() ([]schemas.ModelProvider, error) {
+    return []schemas.ModelProvider{schemas.OpenAI, schemas.Anthropic, OpenAICustom}, nil
+}
+
+func (a *MyAccount) GetKeysForProvider(ctx *context.Context, provider schemas.ModelProvider) ([]schemas.Key, error) {
+    switch provider {
+    case schemas.OpenAI:
+        return []schemas.Key{{
+            Value:  os.Getenv("OPENAI_API_KEY"),
+            Models: []string{},
+            Weight: 1.0,
+        }}, nil
+    case schemas.Anthropic:
+        return []schemas.Key{{
+            Value:  os.Getenv("ANTHROPIC_API_KEY"),
+            Models: []string{},
+            Weight: 1.0,
+        }}, nil
+    case OpenAICustom:
+        return []schemas.Key{{
+            Value:  os.Getenv("PROVIDER_API_KEY"),
+            Models: []string{},
+            Weight: 1.0,
+        }}, nil
+    }
+    return nil, fmt.Errorf("provider %s not supported", provider)
+}
+
+func (a *MyAccount) GetConfigForProvider(provider schemas.ModelProvider) (*schemas.ProviderConfig, error) {
+    switch provider {
+        case OpenAICustom:
+            // Return config with base provider and allowed requests
+            return &schemas.ProviderConfig{
+                    NetworkConfig:            schemas.DefaultNetworkConfig,
+                    ConcurrencyAndBufferSize: schemas.DefaultConcurrencyAndBufferSize,
+				    BaseProviderType: schemas.OpenAI,
+                    AllowedRequests: &schemas.AllowedRequests{
+                        TextCompletion:       false,
+                        ChatCompletion:       true,
+                        ChatCompletionStream: true,
+                        Embedding:            true,
+                        Speech:               false,
+                        SpeechStream:         false,
+                        Transcription:        false,
+                        TranscriptionStream:  false,
+                    },
+            }, nil
+        default: 
+            // Return default config for other providers
+            return &schemas.ProviderConfig{
+                    NetworkConfig:            schemas.DefaultNetworkConfig,
+                    ConcurrencyAndBufferSize: schemas.DefaultConcurrencyAndBufferSize,
+            }, nil     
+    }
+}
+```
+
+
 </Tab>
 
 </Tabs>
diff --git a/transports/bifrost-http/handlers/providers.go b/transports/bifrost-http/handlers/providers.go
@@ -178,6 +178,11 @@ func (h *ProviderHandler) AddProvider(ctx *fasthttp.RequestCtx) {
 			return
 		}
 
+		// Validate allowed requests against base provider capabilities
+		if req.CustomProviderConfig.AllowedRequests != nil {
+			schemas.ValidateAllowedRequests(req.CustomProviderConfig.AllowedRequests, req.CustomProviderConfig.BaseProviderType)
+		}
+
 		// CustomProviderKey is internally set by Bifrost, no validation needed
 		baseProvider = req.CustomProviderConfig.BaseProviderType
 	}
diff --git a/transports/bifrost-http/lib/config.go b/transports/bifrost-http/lib/config.go
@@ -1841,6 +1841,12 @@ func ValidateCustomProvider(config configstore.ProviderConfig, provider schemas.
 	if !bifrost.IsSupportedBaseProvider(cpc.BaseProviderType) {
 		return fmt.Errorf("custom provider validation failed: unsupported base_provider_type: %s", cpc.BaseProviderType)
 	}
+
+	// Validate allowed requests against base provider capabilities
+	if cpc.AllowedRequests != nil {
+		schemas.ValidateAllowedRequests(cpc.AllowedRequests, cpc.BaseProviderType)
+	}
+
 	return nil
 }
 
diff --git a/ui/app/providers/configure/provider-form.tsx b/ui/app/providers/configure/provider-form.tsx
@@ -17,7 +17,7 @@ import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { TagInput } from "@/components/ui/tag-input";
 import { Textarea } from "@/components/ui/textarea";
 import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip";
-import { DEFAULT_ALLOWED_REQUESTS, DEFAULT_NETWORK_CONFIG, DEFAULT_PERFORMANCE_CONFIG } from "@/lib/constants/config";
+import { DEFAULT_NETWORK_CONFIG, DEFAULT_PERFORMANCE_CONFIG, getProviderDefaultAllowedRequests } from "@/lib/constants/config";
 import { ProviderIconType, RenderProviderIcon } from "@/lib/constants/icons";
 import { PROVIDER_LABELS, PROVIDERS as Providers } from "@/lib/constants/logs";
 import { getErrorMessage, useCreateProviderMutation, useUpdateProviderMutation } from "@/lib/store";
@@ -111,11 +111,12 @@ const createInitialState = (provider?: ProviderResponse | null, defaultProvider?
 
 	// Check if this is a custom provider
 	const isCustomProvider = provider && !Providers.includes(provider.name as any);
+	const baseProviderType = provider?.custom_provider_config?.base_provider_type || "";
 
 	return {
 		selectedProvider: providerName,
 		customProviderName: isCustomProvider ? provider.name : "",
-		baseProviderType: provider?.custom_provider_config?.base_provider_type || "",
+		baseProviderType,
 		keys: isNewProvider && keysRequired ? [createDefaultKey()] : !isNewProvider && keysRequired && provider?.keys ? provider.keys : [],
 		networkConfig: provider?.network_config || DEFAULT_NETWORK_CONFIG,
 		performanceConfig: provider?.concurrency_and_buffer_size || DEFAULT_PERFORMANCE_CONFIG,
@@ -126,7 +127,7 @@ const createInitialState = (provider?: ProviderResponse | null, defaultProvider?
 			password: "",
 		},
 		sendBackRawResponse: provider?.send_back_raw_response || false,
-		allowedRequests: provider?.custom_provider_config?.allowed_requests || DEFAULT_ALLOWED_REQUESTS,
+		allowedRequests: provider?.custom_provider_config?.allowed_requests || getProviderDefaultAllowedRequests(baseProviderType),
 	};
 };
 
@@ -189,6 +190,18 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 	const isCustomProvider =
 		selectedProvider === "custom" || !!customProviderName || !!baseProviderType || !Providers.includes(selectedProvider as any);
 
+	// Update allowed requests when base provider type changes for custom providers
+	useEffect(() => {
+		if (baseProviderType && isCustomProvider && !isEditingExisting) {
+			const newAllowedRequests = getProviderDefaultAllowedRequests(baseProviderType);
+			setFormData((prev) => ({
+				...prev,
+				allowedRequests: newAllowedRequests,
+				isDirty: true,
+			}));
+		}
+	}, [baseProviderType, isCustomProvider, isEditingExisting]);
+
 	const performanceValid =
 		performanceConfig.concurrency > 0 && performanceConfig.buffer_size > 0 && performanceConfig.concurrency < performanceConfig.buffer_size;
 
@@ -410,10 +423,15 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 						/^[a-z0-9_-]+$/,
 						"Custom provider name must be lowercase alphanumeric and may include ‘-’ or ‘_’ (no spaces)",
 					),
-					Validator.custom(
-						!allProviders.some((p) => p.name === customProviderName.trim() && p.name !== (provider?.name || "")),
-						"A provider with this name already exists",
-					),
+					// Only check for duplicate names when creating new providers, not when updating existing ones
+					...(provider
+						? []
+						: [
+								Validator.custom(
+									!allProviders.some((p) => p.name === customProviderName.trim()),
+									"A provider with this name already exists",
+								),
+							]),
 					Validator.required(baseProviderType, "Base provider type is required for custom providers"),
 					Validator.custom(
 						!Providers.includes(customProviderName.trim() as any),
@@ -646,7 +664,7 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 			password: "",
 		},
 		sendBackRawResponse: false,
-		allowedRequests: DEFAULT_ALLOWED_REQUESTS,
+		allowedRequests: getProviderDefaultAllowedRequests("openai"),
 	});
 
 	const [selectedTab, setSelectedTab] = useState(tabs[0]?.id || "api-keys");
@@ -843,7 +861,10 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 										{/* Allowed Requests Configuration */}
 										<div className="space-y-2">
 											<div className="text-sm font-medium">Allowed Request Types</div>
-											<p className="text-muted-foreground text-xs">Select which request types this custom provider can handle</p>
+											<p className="text-muted-foreground text-xs">
+												Select which request types this custom provider can handle. Automatically disabled for features not supported by the
+												underlying provider.
+											</p>
 
 											<div className="grid grid-cols-2 gap-4">
 												<div className="space-y-3">
@@ -853,6 +874,7 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 															size="md"
 															checked={allowedRequests.text_completion}
 															onCheckedChange={(checked) => updateAllowedRequest("text_completion", checked)}
+															disabled={!getProviderDefaultAllowedRequests(baseProviderType || "openai").text_completion}
 														/>
 													</div>
 													<div className="flex items-center justify-between">
@@ -861,6 +883,7 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 															size="md"
 															checked={allowedRequests.chat_completion}
 															onCheckedChange={(checked) => updateAllowedRequest("chat_completion", checked)}
+															disabled={!getProviderDefaultAllowedRequests(baseProviderType || "openai").chat_completion}
 														/>
 													</div>
 													<div className="flex items-center justify-between">
@@ -869,6 +892,7 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 															size="md"
 															checked={allowedRequests.chat_completion_stream}
 															onCheckedChange={(checked) => updateAllowedRequest("chat_completion_stream", checked)}
+															disabled={!getProviderDefaultAllowedRequests(baseProviderType || "openai").chat_completion_stream}
 														/>
 													</div>
 													<div className="flex items-center justify-between">
@@ -877,6 +901,7 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 															size="md"
 															checked={allowedRequests.embedding}
 															onCheckedChange={(checked) => updateAllowedRequest("embedding", checked)}
+															disabled={!getProviderDefaultAllowedRequests(baseProviderType || "openai").embedding}
 														/>
 													</div>
 												</div>
@@ -887,6 +912,7 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 															size="md"
 															checked={allowedRequests.speech}
 															onCheckedChange={(checked) => updateAllowedRequest("speech", checked)}
+															disabled={!getProviderDefaultAllowedRequests(baseProviderType || "openai").speech}
 														/>
 													</div>
 													<div className="flex items-center justify-between">
@@ -895,6 +921,7 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 															size="md"
 															checked={allowedRequests.speech_stream}
 															onCheckedChange={(checked) => updateAllowedRequest("speech_stream", checked)}
+															disabled={!getProviderDefaultAllowedRequests(baseProviderType || "openai").speech_stream}
 														/>
 													</div>
 													<div className="flex items-center justify-between">
@@ -903,6 +930,7 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 															size="md"
 															checked={allowedRequests.transcription}
 															onCheckedChange={(checked) => updateAllowedRequest("transcription", checked)}
+															disabled={!getProviderDefaultAllowedRequests(baseProviderType || "openai").transcription}
 														/>
 													</div>
 													<div className="flex items-center justify-between">
@@ -911,6 +939,7 @@ export default function ProviderForm({ provider, onSave, onCancel, existingProvi
 															size="md"
 															checked={allowedRequests.transcription_stream}
 															onCheckedChange={(checked) => updateAllowedRequest("transcription_stream", checked)}
+															disabled={!getProviderDefaultAllowedRequests(baseProviderType || "openai").transcription_stream}
 														/>
 													</div>
 												</div>
diff --git a/ui/lib/constants/config.ts b/ui/lib/constants/config.ts

Original file line number	Diff line number	Diff line change
`@@ -178,6 +178,11 @@ func (h ProviderHandler) AddProvider(ctx fasthttp.RequestCtx) {`
`178`	`178`	`return`
`179`	`179`	`}`
`180`	`180`
	`181`	`+ // Validate allowed requests against base provider capabilities`
	`182`	`+ if req.CustomProviderConfig.AllowedRequests != nil {`
	`183`	`+ schemas.ValidateAllowedRequests(req.CustomProviderConfig.AllowedRequests, req.CustomProviderConfig.BaseProviderType)`
	`184`	`+ }`
	`185`	`+`
`181`	`186`	`// CustomProviderKey is internally set by Bifrost, no validation needed`
`182`	`187`	`baseProvider = req.CustomProviderConfig.BaseProviderType`
`183`	`188`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1841,6 +1841,12 @@ func ValidateCustomProvider(config configstore.ProviderConfig, provider schemas.`
`1841`	`1841`	`if !bifrost.IsSupportedBaseProvider(cpc.BaseProviderType) {`
`1842`	`1842`	`return fmt.Errorf("custom provider validation failed: unsupported base_provider_type: %s", cpc.BaseProviderType)`
`1843`	`1843`	`}`
	`1844`	`+`
	`1845`	`+ // Validate allowed requests against base provider capabilities`
	`1846`	`+ if cpc.AllowedRequests != nil {`
	`1847`	`+ schemas.ValidateAllowedRequests(cpc.AllowedRequests, cpc.BaseProviderType)`
	`1848`	`+ }`
	`1849`	`+`
`1844`	`1850`	`return nil`
`1845`	`1851`	`}`
`1846`	`1852`