From 7307fde8399cbed658f51507b7c24df20b1464a8 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Fri, 12 Jun 2026 12:09:35 +0100 Subject: [PATCH] Studio: Add custom provider option to Connections (#6112) * feat: add custom connection * Fix custom provider handling for PR #6112 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix custom provider connection test for PR #6112 --------- Co-authored-by: wasimysaid Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/providers.py | 17 +++ studio/backend/models/providers.py | 3 + studio/backend/routes/providers.py | 31 +++- .../test_external_provider_usage_chunk.py | 139 ++++++++++++++++++ .../src/features/chat/api/providers-api.ts | 2 + .../features/chat/chat-providers-dialog.tsx | 32 +++- .../src/features/chat/external-providers.ts | 15 +- 7 files changed, 229 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py index f7a2da2833..5b72373c03 100644 --- a/studio/backend/core/inference/providers.py +++ b/studio/backend/core/inference/providers.py @@ -249,6 +249,23 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { # Surfaced via the frontend's CUSTOM_PROVIDER_PRESETS, not the dropdown. "hidden": True, }, + "custom": { + "display_name": "Custom", + # User-supplied via provider_base_url. + "base_url": "", + "default_models": [], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "notes": ( + "User-supplied OpenAI-compatible server. Routed to " + "/v1/chat/completions; /models is optional." + ), + # Surfaced by the frontend's generic Custom option, not the dropdown. + "hidden": True, + }, "ollama": { "display_name": "Ollama", "base_url": "http://localhost:11434/v1", diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py index 8f90e19df6..5a75246c07 100644 --- a/studio/backend/models/providers.py +++ b/studio/backend/models/providers.py @@ -108,6 +108,9 @@ class ProviderTestRequest(BaseModel): base_url: Optional[str] = Field( None, description = "Custom base URL (overrides registry default)" ) + model_id: Optional[str] = Field( + None, description = "Model ID for providers that need a chat probe" + ) class ProviderTestResult(BaseModel): diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index 64d1c3eab3..5a55c9b0bb 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -190,7 +190,8 @@ async def test_provider( """ Test connectivity to an external provider. - Makes a lightweight GET /models call to verify the API key works. + Makes a lightweight GET /models call to verify the API key works. Generic + custom endpoints use a chat-completions probe because /models is optional. encrypted_api_key is decrypted server-side and never stored. """ info = get_provider_info(payload.provider_type) @@ -212,6 +213,14 @@ async def test_provider( ) base_url = payload.base_url or info["base_url"] + if payload.provider_type == "custom": + if not base_url: + return ProviderTestResult( + success = False, + message = "Connection failed: Base URL is required for custom providers.", + models_count = None, + ) + client = ExternalProviderClient( provider_type = payload.provider_type, base_url = base_url, @@ -220,6 +229,26 @@ async def test_provider( ) try: + if payload.provider_type == "custom": + model_id = (payload.model_id or "").strip() + if not model_id: + return ProviderTestResult( + success = False, + message = "Connection failed: add a model ID to test custom providers.", + models_count = None, + ) + await client.chat_completion( + messages = [{"role": "user", "content": "ping"}], + model = model_id, + temperature = 0.0, + top_p = 1.0, + max_tokens = 1, + ) + return ProviderTestResult( + success = True, + message = "Connected successfully. Chat completions endpoint responded.", + models_count = None, + ) if info.get("model_list_mode") == "curated": await client.verify_models_endpoint_lightweight() return ProviderTestResult( diff --git a/studio/backend/tests/test_external_provider_usage_chunk.py b/studio/backend/tests/test_external_provider_usage_chunk.py index ebfd6a8f50..3ec9133718 100644 --- a/studio/backend/tests/test_external_provider_usage_chunk.py +++ b/studio/backend/tests/test_external_provider_usage_chunk.py @@ -144,6 +144,14 @@ def _make_openai_client() -> ExternalProviderClient: ) +def _make_custom_client() -> ExternalProviderClient: + return ExternalProviderClient( + provider_type = "custom", + base_url = "http://custom.example/v1", + api_key = "", + ) + + def _anthropic_sse(events: list[dict]) -> bytes: chunks: list[str] = [] for event in events: @@ -180,6 +188,137 @@ def _usage_chunks(lines: list[str]) -> list[dict]: return out +def test_custom_provider_registry_is_hidden(): + from core.inference.providers import get_provider_info, list_available_providers + + info = get_provider_info("custom") + assert info is not None + assert info["hidden"] is True + assert "custom" not in {p["provider_type"] for p in list_available_providers()} + + +def test_custom_provider_uses_chat_completions_without_auth_key(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["headers"] = dict(request.headers) + captured["body"] = json.loads(request.content.decode("utf-8")) + return httpx.Response( + 200, + content = b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\ndata: [DONE]\n\n', + headers = {"content-type": "text/event-stream"}, + ) + + _mock_http_client(monkeypatch, handler) + + async def run(): + client = _make_custom_client() + lines = await _collect( + client.stream_chat_completion( + messages = [{"role": "user", "content": "ping"}], + model = "Qwen/Qwen3-0.6B", + temperature = 0.7, + top_p = 0.95, + max_tokens = 64, + ) + ) + await client.close() + return lines + + lines = _drive(run()) + assert captured["url"] == "http://custom.example/v1/chat/completions" + assert "authorization" not in {k.lower() for k in captured["headers"]} + assert captured["body"]["model"] == "Qwen/Qwen3-0.6B" + assert any("ok" in line for line in lines) + + +def test_custom_provider_test_endpoint_probes_chat_completion(monkeypatch): + import importlib.util + import sys + from pathlib import Path + + module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py" + spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path) + assert spec is not None + assert spec.loader is not None + providers_route = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = providers_route + spec.loader.exec_module(providers_route) + + captured: dict = {} + + class _FakeClient: + def __init__(self, **kwargs): + captured["init"] = kwargs + + async def chat_completion(self, **kwargs): + captured["chat_completion"] = kwargs + return {"choices": [{"message": {"content": "ok"}}]} + + async def list_models(self): + raise AssertionError("custom provider test must not call /models") + + async def close(self): + captured["closed"] = True + + monkeypatch.setattr(providers_route, "ExternalProviderClient", _FakeClient) + + async def run(): + return await providers_route.test_provider( + providers_route.ProviderTestRequest( + provider_type = "custom", + base_url = "http://custom.example/v1", + model_id = "Qwen/Qwen3-0.6B", + ), + current_subject = "unsloth", + ) + + result = _drive(run()) + assert result.success is True + assert result.models_count is None + assert captured["init"]["provider_type"] == "custom" + assert captured["chat_completion"]["model"] == "Qwen/Qwen3-0.6B" + assert captured["chat_completion"]["max_tokens"] == 1 + assert captured["closed"] is True + + +def test_custom_provider_test_endpoint_requires_model_id(monkeypatch): + import importlib.util + import sys + from pathlib import Path + + module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py" + spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path) + assert spec is not None + assert spec.loader is not None + providers_route = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = providers_route + spec.loader.exec_module(providers_route) + + class _FakeClient: + def __init__(self, **kwargs): + pass + + async def close(self): + pass + + monkeypatch.setattr(providers_route, "ExternalProviderClient", _FakeClient) + + async def run(): + return await providers_route.test_provider( + providers_route.ProviderTestRequest( + provider_type = "custom", + base_url = "http://custom.example/v1", + ), + current_subject = "unsloth", + ) + + result = _drive(run()) + assert result.success is False + assert "model ID" in result.message + + def test_anthropic_stream_emits_usage_chunk_before_done(monkeypatch): sse_events = [ { diff --git a/studio/frontend/src/features/chat/api/providers-api.ts b/studio/frontend/src/features/chat/api/providers-api.ts index d52e5e855e..4ad996d54f 100644 --- a/studio/frontend/src/features/chat/api/providers-api.ts +++ b/studio/frontend/src/features/chat/api/providers-api.ts @@ -197,6 +197,7 @@ export async function testProviderConnection(payload: { providerType: string; apiKey: string; baseUrl?: string | null; + modelId?: string | null; }): Promise { return withApiKeyEncryptionRetry(payload.apiKey, async (encryptedApiKey) => { const response = await authFetch("/api/providers/test", { @@ -206,6 +207,7 @@ export async function testProviderConnection(payload: { provider_type: payload.providerType, encrypted_api_key: encryptedApiKey, base_url: payload.baseUrl ?? null, + model_id: payload.modelId ?? null, }), }); return parseJsonOrThrow(response); diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index 0594b55083..da19cd749b 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -59,6 +59,7 @@ import { getExternalProviderApiKey, isCustomProviderType, LEGACY_CUSTOM_PROVIDER_TYPE, + CUSTOM_PROVIDER_DISPLAY_NAME, removeExternalProviderApiKey, setExternalProviderApiKey, supportsProviderPromptCaching, @@ -176,7 +177,8 @@ function shouldAppendOpenAiVersionPath(providerType: string): boolean { return ( providerType === "ollama" || providerType === "llama_cpp" || - providerType === "vllm" + providerType === "vllm" || + providerType === LEGACY_CUSTOM_PROVIDER_TYPE ); } @@ -230,7 +232,9 @@ export function ChatProvidersSettings({ const [mutatingProvider, setMutatingProvider] = useState(false); const [manualModelIds, setManualModelIds] = useState(""); const [modelSearchQuery, setModelSearchQuery] = useState(""); - const [customProviderName, setCustomProviderName] = useState("Custom"); + const [customProviderName, setCustomProviderName] = useState( + CUSTOM_PROVIDER_DISPLAY_NAME, + ); const [isReasoningModel, setIsReasoningModel] = useState(false); const reduceMotion = useReducedMotion(); const connectionsEnabled = useExternalProvidersStore( @@ -939,8 +943,12 @@ export function ChatProvidersSettings({ async function testProvider(provider: ExternalProviderConfig) { const savedKey = getExternalProviderApiKey(provider.id).trim(); - // Local OpenAI-compat presets skip API keys — run the connection check. - if (!savedKey && !supportsRemoteModelCatalog(provider.providerType)) { + // Hosted registry providers require keys. Local OpenAI-compatible presets + // may be keyless. + if ( + !savedKey && + !supportsRemoteModelCatalog(provider.providerType) + ) { if (isCustomProviderType(provider.providerType)) { await editProvider(provider); toast.info(CUSTOM_PROVIDER_MISSING_KEY_MESSAGE); @@ -957,6 +965,10 @@ export function ChatProvidersSettings({ provider.providerType, apiKey: savedKey, baseUrl: provider.baseUrl || null, + modelId: + provider.providerType === LEGACY_CUSTOM_PROVIDER_TYPE + ? (provider.models[0] ?? null) + : null, }); if (result.success) { toast.success(result.message); @@ -1062,6 +1074,16 @@ export function ChatProvidersSettings({ ))} + + + + {CUSTOM_PROVIDER_DISPLAY_NAME} + + @@ -1144,7 +1166,7 @@ export function ChatProvidersSettings({ onChange={(event) => setCustomProviderName(event.target.value) } - placeholder="Custom" + placeholder={CUSTOM_PROVIDER_DISPLAY_NAME} className="h-9 text-sm" /> diff --git a/studio/frontend/src/features/chat/external-providers.ts b/studio/frontend/src/features/chat/external-providers.ts index 300d384cdf..bc718abbba 100644 --- a/studio/frontend/src/features/chat/external-providers.ts +++ b/studio/frontend/src/features/chat/external-providers.ts @@ -107,6 +107,7 @@ export function providerTypeSupportsVision( export const CUSTOM_BACKEND_PROVIDER_TYPE = "openai"; export const LEGACY_CUSTOM_PROVIDER_TYPE = "custom"; +export const CUSTOM_PROVIDER_DISPLAY_NAME = "Custom"; export const CUSTOM_PROVIDER_PRESETS = [ { @@ -130,7 +131,7 @@ export const CUSTOM_PROVIDER_PRESETS = [ ] as const; const CUSTOM_PROVIDER_LABELS: Record = { - [LEGACY_CUSTOM_PROVIDER_TYPE]: "Custom", + [LEGACY_CUSTOM_PROVIDER_TYPE]: CUSTOM_PROVIDER_DISPLAY_NAME, ...Object.fromEntries( CUSTOM_PROVIDER_PRESETS.map((preset) => [ preset.providerType, @@ -166,8 +167,9 @@ export function isCustomProviderType( return providerType in CUSTOM_PROVIDER_LABELS; } -/** Local OpenAI-compat presets that expose GET /v1/models (no API key). */ +/** OpenAI-compat custom types that may expose GET /v1/models. */ const REMOTE_MODEL_CATALOG_CUSTOM_PROVIDER_TYPES = new Set([ + LEGACY_CUSTOM_PROVIDER_TYPE, "ollama", "vllm", "llama_cpp", @@ -189,7 +191,7 @@ export function customPresetSkipsApiKeyField( return providerType === "ollama" || providerType === "llama_cpp"; } -/** Catalog load plus optional manual model IDs (OpenRouter + local presets). */ +/** Catalog load plus optional manual model IDs. */ export function allowsManualModelIdsWithCatalog( providerType: string | null | undefined, ): boolean { @@ -201,7 +203,7 @@ export function allowsManualModelIdsWithCatalog( export function customProviderDisplayName( providerType: string | null | undefined, ): string { - if (!providerType) return "Custom"; + if (!providerType) return CUSTOM_PROVIDER_DISPLAY_NAME; return CUSTOM_PROVIDER_LABELS[providerType] ?? providerType; } @@ -246,6 +248,11 @@ export function toExternalBackendProviderType( if (providerType === "vllm") return "vllm"; if (providerType === "ollama") return "ollama"; if (providerType === "llama_cpp") return "llama_cpp"; + // Generic custom servers are OpenAI-compatible, but should still use the + // chat-completions backend path instead of OpenAI's Responses API route. + if (providerType === LEGACY_CUSTOM_PROVIDER_TYPE) { + return LEGACY_CUSTOM_PROVIDER_TYPE; + } return isCustomProviderType(providerType) ? CUSTOM_BACKEND_PROVIDER_TYPE : providerType;