diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py new file mode 100644 index 0000000000..cb4b391864 --- /dev/null +++ b/studio/backend/core/inference/external_provider.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Async HTTP client for proxying chat completions to external LLM providers. + +All target providers (OpenAI, Mistral, Google Gemini, Cohere, Together AI, +Fireworks AI, Perplexity) expose OpenAI-compatible /v1/chat/completions +endpoints, so a single client handles all of them. +""" + +import logging +from typing import Any, AsyncGenerator, Optional + +import httpx + +logger = logging.getLogger(__name__) + + +class ExternalProviderClient: + """Async proxy for OpenAI-compatible external LLM APIs.""" + + def __init__( + self, + provider_type: str, + base_url: str, + api_key: str, + timeout: float = 120.0, + ): + self.provider_type = provider_type + self.base_url = base_url.rstrip("/") + self.api_key = api_key + self._client = httpx.AsyncClient(timeout=httpx.Timeout(timeout, connect=10.0)) + + def _auth_headers(self) -> dict[str, str]: + """Build authentication headers. All supported providers use Bearer tokens.""" + return { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } + + async def stream_chat_completion( + self, + messages: list[dict[str, Any]], + model: str, + temperature: float = 0.7, + top_p: float = 0.95, + max_tokens: Optional[int] = None, + presence_penalty: float = 0.0, + stream: bool = True, + ) -> AsyncGenerator[str, None]: + """ + Yield raw SSE lines from the external provider. + + Each yielded string is a complete SSE line (e.g. 'data: {...}' or + 'data: [DONE]'). The caller wraps these into a StreamingResponse. + """ + body: dict[str, Any] = { + "model": model, + "messages": messages, + "stream": stream, + "temperature": temperature, + "top_p": top_p, + "presence_penalty": presence_penalty, + } + if max_tokens is not None: + body["max_tokens"] = max_tokens + + url = f"{self.base_url}/chat/completions" + logger.info( + "Proxying chat completion to %s (provider=%s, model=%s)", + url, + self.provider_type, + model, + ) + + try: + async with self._client.stream( + "POST", + url, + json=body, + headers=self._auth_headers(), + ) as response: + if response.status_code != 200: + error_body = await response.aread() + error_text = error_body.decode("utf-8", errors="replace") + logger.error( + "External provider returned %d: %s", + response.status_code, + error_text[:500], + ) + # Yield an error in SSE format so the frontend can display it + yield _error_sse_line( + response.status_code, error_text, self.provider_type + ) + return + + async for line in response.aiter_lines(): + if line.strip(): + yield line + + except httpx.ConnectError as exc: + logger.error("Connection error to %s: %s", self.provider_type, exc) + yield _error_sse_line( + 502, f"Failed to connect to {self.provider_type}: {exc}", self.provider_type + ) + except httpx.ReadTimeout as exc: + logger.error("Read timeout from %s: %s", self.provider_type, exc) + yield _error_sse_line( + 504, f"Timeout waiting for {self.provider_type} response", self.provider_type + ) + except httpx.HTTPError as exc: + logger.error("HTTP error from %s: %s", self.provider_type, exc) + yield _error_sse_line( + 502, f"Error communicating with {self.provider_type}: {exc}", self.provider_type + ) + + async def chat_completion( + self, + messages: list[dict[str, Any]], + model: str, + temperature: float = 0.7, + top_p: float = 0.95, + max_tokens: Optional[int] = None, + presence_penalty: float = 0.0, + ) -> dict[str, Any]: + """Non-streaming chat completion. Returns the full response dict.""" + body: dict[str, Any] = { + "model": model, + "messages": messages, + "stream": False, + "temperature": temperature, + "top_p": top_p, + "presence_penalty": presence_penalty, + } + if max_tokens is not None: + body["max_tokens"] = max_tokens + + response = await self._client.post( + f"{self.base_url}/chat/completions", + json=body, + headers=self._auth_headers(), + ) + response.raise_for_status() + return response.json() + + async def list_models(self) -> list[dict[str, Any]]: + """ + Call GET /models on the provider to discover available models. + + Returns a list of model dicts with at least 'id' and optionally + 'created', 'owned_by', etc. + """ + try: + response = await self._client.get( + f"{self.base_url}/models", + headers=self._auth_headers(), + ) + response.raise_for_status() + data = response.json() + # OpenAI format: {"data": [{"id": "...", ...}, ...]} + models = data.get("data", []) + return models + except httpx.HTTPError as exc: + logger.error("Failed to list models from %s: %s", self.provider_type, exc) + raise + + async def close(self) -> None: + """Close the underlying HTTP client.""" + await self._client.aclose() + + +def _error_sse_line(status_code: int, message: str, provider_type: str) -> str: + """Format an error as an SSE data line in OpenAI error format.""" + import json + + error_obj = { + "error": { + "message": message, + "type": "provider_error", + "code": str(status_code), + "provider": provider_type, + } + } + return f"data: {json.dumps(error_obj)}" diff --git a/studio/backend/core/inference/key_exchange.py b/studio/backend/core/inference/key_exchange.py new file mode 100644 index 0000000000..c92bd9d9d9 --- /dev/null +++ b/studio/backend/core/inference/key_exchange.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +RSA key pair for encrypting API keys in transit. + +The frontend encrypts API keys with the server's public key before +including them in requests. The backend decrypts with its private key +before forwarding to external providers. + +The key pair is generated at server startup and lives only in memory — +it is regenerated on each restart. The frontend fetches the public key +via GET /api/providers/public-key on load. +""" + +import base64 +import logging + +from cryptography.hazmat.primitives.asymmetric import rsa, padding +from cryptography.hazmat.primitives import serialization, hashes + +logger = logging.getLogger(__name__) + +_private_key: rsa.RSAPrivateKey | None = None +_public_key_pem: str | None = None + + +def init_key_pair() -> None: + """Generate an RSA-2048 key pair. Called once at server startup.""" + global _private_key, _public_key_pem + _private_key = rsa.generate_private_key( + public_exponent=65537, + key_size=2048, + ) + _public_key_pem = _private_key.public_key().public_bytes( + serialization.Encoding.PEM, + serialization.PublicFormat.SubjectPublicKeyInfo, + ).decode("utf-8") + logger.info("RSA key pair generated for API key encryption") + + +def get_public_key_pem() -> str: + """Return the PEM-encoded public key for the frontend.""" + if _public_key_pem is None: + raise RuntimeError("Key pair not initialized. Call init_key_pair() first.") + return _public_key_pem + + +def decrypt_api_key(encrypted_b64: str) -> str: + """ + Decrypt an API key that was encrypted with the public key. + + Args: + encrypted_b64: Base64-encoded RSA-OAEP ciphertext. + + Returns: + The plaintext API key string. + """ + if _private_key is None: + raise RuntimeError("Key pair not initialized. Call init_key_pair() first.") + + ciphertext = base64.b64decode(encrypted_b64) + plaintext = _private_key.decrypt( + ciphertext, + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None, + ), + ) + return plaintext.decode("utf-8") diff --git a/studio/backend/core/inference/providers.py b/studio/backend/core/inference/providers.py new file mode 100644 index 0000000000..0fb86431f5 --- /dev/null +++ b/studio/backend/core/inference/providers.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Static registry of supported external LLM providers. + +All providers expose OpenAI-compatible /v1/chat/completions endpoints +with Bearer token authentication and SSE streaming support. +""" + +from typing import Any + +PROVIDER_REGISTRY: dict[str, dict[str, Any]] = { + "openai": { + "display_name": "OpenAI", + "base_url": "https://api.openai.com/v1", + "default_models": [ + "gpt-4o", + "gpt-4o-mini", + "gpt-4.1", + "gpt-4.1-mini", + "gpt-4.1-nano", + "o3-mini", + ], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + }, + "mistral": { + "display_name": "Mistral AI", + "base_url": "https://api.mistral.ai/v1", + "default_models": [ + "mistral-large-latest", + "mistral-small-latest", + "codestral-latest", + ], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + }, + "google": { + "display_name": "Google Gemini", + "base_url": "https://generativelanguage.googleapis.com/v1beta/openai", + "default_models": [ + "gemini-2.5-flash", + "gemini-2.5-pro", + "gemini-2.5-flash-lite", + ], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "notes": "OpenAI compatibility layer (beta). Uses native Google API key as Bearer token.", + }, + "cohere": { + "display_name": "Cohere", + "base_url": "https://api.cohere.ai/compatibility/v1", + "default_models": [ + "command-r-plus", + "command-r", + "command-a", + ], + "supports_streaming": True, + "supports_vision": False, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "notes": "OpenAI compatibility layer over native Cohere API.", + }, + "together": { + "display_name": "Together AI", + "base_url": "https://api.together.xyz/v1", + "default_models": [ + "deepseek-ai/DeepSeek-R1", + "deepseek-ai/DeepSeek-V3", + "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", + "Qwen/Qwen3-235B-A22B", + ], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + }, + "fireworks": { + "display_name": "Fireworks AI", + "base_url": "https://api.fireworks.ai/inference/v1", + "default_models": [ + "accounts/fireworks/models/deepseek-v3-0324", + "accounts/fireworks/models/llama4-maverick-instruct-basic", + "accounts/fireworks/models/qwen3-30b", + "accounts/fireworks/models/llama-3.3-70b", + ], + "supports_streaming": True, + "supports_vision": True, + "supports_tool_calling": True, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "notes": "Model IDs use 'accounts/fireworks/models/' prefix. Usage stats included in streaming.", + }, + "perplexity": { + "display_name": "Perplexity", + "base_url": "https://api.perplexity.ai", + "default_models": [ + "sonar-pro", + "sonar", + "sonar-reasoning", + "sonar-reasoning-pro", + ], + "supports_streaming": True, + "supports_vision": False, + "supports_tool_calling": False, + "auth_header": "Authorization", + "auth_prefix": "Bearer ", + "notes": "Web-grounded responses with built-in search.", + }, +} + + +def get_provider_info(provider_type: str) -> dict[str, Any] | None: + """Return the registry entry for a provider type, or None if unknown.""" + return PROVIDER_REGISTRY.get(provider_type) + + +def get_base_url(provider_type: str) -> str | None: + """Return the default base URL for a provider type.""" + info = PROVIDER_REGISTRY.get(provider_type) + return info["base_url"] if info else None + + +def list_available_providers() -> list[dict[str, Any]]: + """Return all registered providers (for the /registry endpoint).""" + result = [] + for provider_type, info in PROVIDER_REGISTRY.items(): + result.append( + { + "provider_type": provider_type, + "display_name": info["display_name"], + "base_url": info["base_url"], + "default_models": info["default_models"], + "supports_streaming": info["supports_streaming"], + "supports_vision": info.get("supports_vision", False), + "supports_tool_calling": info.get("supports_tool_calling", False), + } + ) + return result diff --git a/studio/backend/main.py b/studio/backend/main.py index 67908d8617..590039b151 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -62,6 +62,7 @@ from routes import ( export_router, inference_router, models_router, + providers_router, training_history_router, training_router, ) @@ -113,6 +114,10 @@ async def lifespan(app: FastAPI): threading.Thread(target = _precache, daemon = True).start() + # Initialize RSA key pair for API key encryption (external providers) + from core.inference.key_exchange import init_key_pair + init_key_pair() + if storage.ensure_default_admin(): bootstrap_pw = storage.get_bootstrap_password() app.state.bootstrap_password = bootstrap_pw @@ -172,6 +177,7 @@ app.include_router(inference_router, prefix = "/api/inference", tags = ["inferen # so external tools (Open WebUI, SillyTavern, etc.) can use the # standard /v1/chat/completions path. app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"]) +app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"]) app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"]) app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"]) app.include_router(export_router, prefix = "/api/export", tags = ["export"]) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index accdcc1290..8f45bca759 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -347,6 +347,28 @@ class ChatCompletionRequest(BaseModel): description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.", ) + # ── External provider routing (x-unsloth extensions) ────────── + provider_id: Optional[str] = Field( + None, + description = "[x-unsloth] Saved provider config ID. If set with encrypted_api_key, routes to external LLM.", + ) + provider_type: Optional[str] = Field( + None, + description = "[x-unsloth] Provider type (e.g. 'openai', 'mistral'). Used if provider_id is not set.", + ) + external_model: Optional[str] = Field( + None, + description = "[x-unsloth] Model ID at the external provider.", + ) + encrypted_api_key: Optional[str] = Field( + None, + description = "[x-unsloth] RSA-encrypted, base64-encoded API key for the external provider.", + ) + provider_base_url: Optional[str] = Field( + None, + description = "[x-unsloth] Override base URL for the external provider.", + ) + # ── Streaming response chunks ──────────────────────────────────── diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py new file mode 100644 index 0000000000..572951cffe --- /dev/null +++ b/studio/backend/models/providers.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Pydantic schemas for the external LLM providers API. +""" + +from typing import Optional + +from pydantic import BaseModel, Field + + +# ── Registry (static provider info) ─────────────────────────────── + + +class ProviderRegistryEntry(BaseModel): + """A supported provider type with its default configuration.""" + + provider_type: str = Field(..., description="Provider identifier (e.g. 'openai', 'mistral')") + display_name: str = Field(..., description="Human-readable provider name") + base_url: str = Field(..., description="Default API base URL") + default_models: list[str] = Field( + default_factory=list, description="Well-known model IDs for this provider" + ) + supports_streaming: bool = Field(True, description="Whether this provider supports SSE streaming") + supports_vision: bool = Field(False, description="Whether this provider supports vision/image input") + supports_tool_calling: bool = Field(False, description="Whether this provider supports tool/function calling") + + +# ── Provider config CRUD ────────────────────────────────────────── + + +class ProviderCreate(BaseModel): + """Request to create a saved provider configuration.""" + + provider_type: str = Field(..., description="Provider type from the registry") + display_name: str = Field(..., description="User-chosen label (e.g. 'My OpenAI Key')") + base_url: Optional[str] = Field( + None, + description="Custom base URL (overrides registry default). Omit to use the default.", + ) + + +class ProviderUpdate(BaseModel): + """Request to update a saved provider configuration.""" + + display_name: Optional[str] = Field(None, description="New display name") + base_url: Optional[str] = Field(None, description="New base URL") + is_enabled: Optional[bool] = Field(None, description="Enable or disable this provider") + + +class ProviderResponse(BaseModel): + """A saved provider configuration (returned by list/get endpoints).""" + + id: str = Field(..., description="Unique provider config ID") + provider_type: str = Field(..., description="Provider type (e.g. 'openai')") + display_name: str = Field(..., description="User-chosen label") + base_url: str = Field(..., description="API base URL") + is_enabled: bool = Field(True, description="Whether this provider is enabled") + created_at: str = Field(..., description="ISO 8601 creation timestamp") + updated_at: str = Field(..., description="ISO 8601 last-update timestamp") + + +# ── Model listing ───────────────────────────────────────────────── + + +class ProviderModelInfo(BaseModel): + """A model available from an external provider.""" + + id: str = Field(..., description="Model ID as expected by the provider API") + display_name: str = Field("", description="Human-readable model name") + context_length: Optional[int] = Field(None, description="Maximum context length in tokens") + owned_by: Optional[str] = Field(None, description="Model owner/organization") + + +class ProviderModelsRequest(BaseModel): + """Request to list models from an external provider.""" + + provider_type: str = Field(..., description="Provider type from the registry") + encrypted_api_key: str = Field(..., description="RSA-encrypted, base64-encoded API key") + base_url: Optional[str] = Field( + None, description="Custom base URL (overrides registry default)" + ) + + +# ── Connection testing ──────────────────────────────────────────── + + +class ProviderTestRequest(BaseModel): + """Request to test connectivity to an external provider.""" + + provider_type: str = Field(..., description="Provider type from the registry") + encrypted_api_key: str = Field(..., description="RSA-encrypted, base64-encoded API key") + base_url: Optional[str] = Field( + None, description="Custom base URL (overrides registry default)" + ) + + +class ProviderTestResult(BaseModel): + """Result of a provider connectivity test.""" + + success: bool = Field(..., description="Whether the test succeeded") + message: str = Field(..., description="Human-readable result message") + models_count: Optional[int] = Field( + None, description="Number of models found (if test succeeded)" + ) diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 186ba82fe0..569f6eff3d 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -15,3 +15,5 @@ huggingface-hub==0.36.2 structlog>=24.1.0 diceware ddgs +cryptography>=42.0.0 +httpx>=0.27.0 diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py index e79f6553f9..872269c906 100644 --- a/studio/backend/routes/__init__.py +++ b/studio/backend/routes/__init__.py @@ -13,6 +13,7 @@ from routes.auth import router as auth_router from routes.data_recipe import router as data_recipe_router from routes.export import router as export_router from routes.training_history import router as training_history_router +from routes.providers import router as providers_router __all__ = [ "training_router", @@ -23,4 +24,5 @@ __all__ = [ "data_recipe_router", "export_router", "training_history_router", + "providers_router", ] diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 6f44a3c69f..954541640a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -82,6 +82,11 @@ from models.inference import ( ) from auth.authentication import get_current_subject +from core.inference.key_exchange import decrypt_api_key +from core.inference.providers import get_provider_info, get_base_url +from core.inference.external_provider import ExternalProviderClient +from storage import providers_db + import io import wave import base64 @@ -840,6 +845,109 @@ def _extract_content_parts( return system_prompt, chat_messages, first_image_b64 +# ── External provider proxy ────────────────────────────────────── + + +async def _proxy_to_external_provider( + payload: ChatCompletionRequest, + request: Request, +) -> StreamingResponse: + """ + Proxy a chat completion request to an external LLM provider. + + Resolves provider config (from DB or registry), decrypts the API key, + and streams the response back in OpenAI SSE format. + """ + # Resolve provider type and base URL + provider_type = payload.provider_type + base_url = payload.provider_base_url + + if payload.provider_id: + config = providers_db.get_provider(payload.provider_id) + if config is None: + raise HTTPException( + status_code = 404, + detail = f"Provider config not found: {payload.provider_id}", + ) + if not config["is_enabled"]: + raise HTTPException( + status_code = 400, + detail = f"Provider '{config['display_name']}' is disabled.", + ) + provider_type = provider_type or config["provider_type"] + base_url = base_url or config["base_url"] + + if not provider_type: + raise HTTPException( + status_code = 400, + detail = "Either provider_id or provider_type is required for external provider routing.", + ) + + # Fall back to registry default base URL + if not base_url: + base_url = get_base_url(provider_type) + if not base_url: + raise HTTPException( + status_code = 400, + detail = f"Unknown provider type: {provider_type}", + ) + + # Decrypt the API key + try: + api_key = decrypt_api_key(payload.encrypted_api_key) + except Exception as exc: + logger.warning("external_provider.decrypt_failed", error = str(exc)) + raise HTTPException( + status_code = 400, + detail = "Failed to decrypt API key. The server key may have changed — try refreshing the page.", + ) + + model = payload.external_model or payload.model + if model == "default": + raise HTTPException( + status_code = 400, + detail = "external_model is required when using an external provider.", + ) + + # Extract messages into plain OpenAI format + system_prompt, chat_messages, _ = _extract_content_parts(payload.messages) + if system_prompt: + chat_messages.insert(0, {"role": "system", "content": system_prompt}) + + client = ExternalProviderClient( + provider_type = provider_type, + base_url = base_url, + api_key = api_key, + ) + + async def _stream(): + try: + async for line in client.stream_chat_completion( + messages = chat_messages, + model = model, + temperature = payload.temperature, + top_p = payload.top_p, + max_tokens = payload.max_tokens, + presence_penalty = payload.presence_penalty, + stream = payload.stream, + ): + yield f"{line}\n\n" + yield "data: [DONE]\n\n" + except Exception as exc: + logger.error("external_provider.stream_error", error = str(exc)) + finally: + await client.close() + + return StreamingResponse( + _stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", + }, + ) + + @router.post("/chat/completions") async def openai_chat_completions( payload: ChatCompletionRequest, @@ -859,6 +967,10 @@ async def openai_chat_completions( - GGUF models → llama-server via LlamaCppBackend - Other models → Unsloth/transformers via InferenceBackend """ + # ── External provider routing ──────────────────────────────── + if payload.encrypted_api_key and (payload.provider_id or payload.provider_type): + return await _proxy_to_external_provider(payload, request) + llama_backend = get_llama_cpp_backend() using_gguf = llama_backend.is_loaded diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py new file mode 100644 index 0000000000..0cf891fabc --- /dev/null +++ b/studio/backend/routes/providers.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +API routes for external LLM provider management. + +Provides endpoints for: + - Discovering available provider types (registry) + - CRUD for saved provider configurations (no API keys stored) + - Fetching the RSA public key for API key encryption + - Testing provider connectivity + - Listing models from a provider +""" + +import uuid +import structlog +from fastapi import APIRouter, Depends, HTTPException + +from auth.authentication import get_current_subject +from core.inference.key_exchange import get_public_key_pem, decrypt_api_key +from core.inference.providers import ( + get_base_url, + get_provider_info, + list_available_providers, +) +from core.inference.external_provider import ExternalProviderClient +from models.providers import ( + ProviderCreate, + ProviderModelsRequest, + ProviderModelInfo, + ProviderResponse, + ProviderRegistryEntry, + ProviderTestRequest, + ProviderTestResult, + ProviderUpdate, +) +from storage import providers_db + +logger = structlog.get_logger(__name__) + +router = APIRouter() + + +# ── Public key for API key encryption ───────────────────────────── + + +@router.get("/public-key") +async def get_public_key( + current_subject: str = Depends(get_current_subject), +): + """Return the RSA public key PEM for client-side API key encryption.""" + return {"public_key": get_public_key_pem()} + + +# ── Provider registry (static) ─────────────────────────────────── + + +@router.get("/registry", response_model=list[ProviderRegistryEntry]) +async def list_registry( + current_subject: str = Depends(get_current_subject), +): + """List all supported provider types with their default configurations.""" + return list_available_providers() + + +# ── Provider config CRUD ────────────────────────────────────────── + + +@router.get("/", response_model=list[ProviderResponse]) +async def list_provider_configs( + current_subject: str = Depends(get_current_subject), +): + """List all saved provider configurations.""" + rows = providers_db.list_providers() + return [ + ProviderResponse( + id=row["id"], + provider_type=row["provider_type"], + display_name=row["display_name"], + base_url=row["base_url"], + is_enabled=bool(row["is_enabled"]), + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + for row in rows + ] + + +@router.post("/", response_model=ProviderResponse, status_code=201) +async def create_provider_config( + payload: ProviderCreate, + current_subject: str = Depends(get_current_subject), +): + """Create a new saved provider configuration (no API key stored).""" + info = get_provider_info(payload.provider_type) + if info is None: + raise HTTPException( + status_code=400, + detail=f"Unknown provider type: {payload.provider_type}. " + f"Use GET /api/providers/registry to see available types.", + ) + + provider_id = uuid.uuid4().hex[:16] + base_url = payload.base_url or info["base_url"] + + providers_db.create_provider( + id=provider_id, + provider_type=payload.provider_type, + display_name=payload.display_name, + base_url=base_url, + ) + + row = providers_db.get_provider(provider_id) + return ProviderResponse( + id=row["id"], + provider_type=row["provider_type"], + display_name=row["display_name"], + base_url=row["base_url"], + is_enabled=bool(row["is_enabled"]), + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + + +@router.put("/{provider_id}", response_model=ProviderResponse) +async def update_provider_config( + provider_id: str, + payload: ProviderUpdate, + current_subject: str = Depends(get_current_subject), +): + """Update a saved provider configuration.""" + existing = providers_db.get_provider(provider_id) + if not existing: + raise HTTPException(status_code=404, detail="Provider not found") + + updated = providers_db.update_provider( + id=provider_id, + display_name=payload.display_name, + base_url=payload.base_url, + is_enabled=payload.is_enabled, + ) + if not updated: + raise HTTPException(status_code=400, detail="No fields to update") + + row = providers_db.get_provider(provider_id) + return ProviderResponse( + id=row["id"], + provider_type=row["provider_type"], + display_name=row["display_name"], + base_url=row["base_url"], + is_enabled=bool(row["is_enabled"]), + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + + +@router.delete("/{provider_id}", status_code=204) +async def delete_provider_config( + provider_id: str, + current_subject: str = Depends(get_current_subject), +): + """Delete a saved provider configuration.""" + deleted = providers_db.delete_provider(provider_id) + if not deleted: + raise HTTPException(status_code=404, detail="Provider not found") + + +# ── Test connectivity ───────────────────────────────────────────── + + +@router.post("/test", response_model=ProviderTestResult) +async def test_provider( + payload: ProviderTestRequest, + current_subject: str = Depends(get_current_subject), +): + """ + Test connectivity to an external provider. + + Makes a lightweight GET /models call to verify the API key works. + The encrypted_api_key is decrypted server-side and never stored. + """ + info = get_provider_info(payload.provider_type) + if info is None: + raise HTTPException( + status_code=400, + detail=f"Unknown provider type: {payload.provider_type}", + ) + + try: + api_key = decrypt_api_key(payload.encrypted_api_key) + except Exception as exc: + logger.warning("Failed to decrypt API key: %s", exc) + raise HTTPException( + status_code=400, + detail="Failed to decrypt API key. The public key may have changed — try refreshing the page.", + ) + + base_url = payload.base_url or info["base_url"] + client = ExternalProviderClient( + provider_type=payload.provider_type, + base_url=base_url, + api_key=api_key, + timeout=15.0, + ) + + try: + models = await client.list_models() + return ProviderTestResult( + success=True, + message=f"Connected successfully. Found {len(models)} model(s).", + models_count=len(models), + ) + except Exception as exc: + logger.warning("Provider test failed for %s: %s", payload.provider_type, exc) + return ProviderTestResult( + success=False, + message=f"Connection failed: {exc}", + models_count=None, + ) + finally: + await client.close() + + +# ── List models from provider ───────────────────────────────────── + + +@router.post("/models", response_model=list[ProviderModelInfo]) +async def list_provider_models( + payload: ProviderModelsRequest, + current_subject: str = Depends(get_current_subject), +): + """ + List models available from an external provider. + + The encrypted_api_key is decrypted server-side and never stored. + """ + info = get_provider_info(payload.provider_type) + if info is None: + raise HTTPException( + status_code=400, + detail=f"Unknown provider type: {payload.provider_type}", + ) + + try: + api_key = decrypt_api_key(payload.encrypted_api_key) + except Exception as exc: + logger.warning("Failed to decrypt API key: %s", exc) + raise HTTPException( + status_code=400, + detail="Failed to decrypt API key. The public key may have changed — try refreshing the page.", + ) + + base_url = payload.base_url or info["base_url"] + client = ExternalProviderClient( + provider_type=payload.provider_type, + base_url=base_url, + api_key=api_key, + timeout=15.0, + ) + + try: + models = await client.list_models() + return [ + ProviderModelInfo( + id=m.get("id", ""), + display_name=m.get("id", ""), + context_length=m.get("context_length") or m.get("context_window"), + owned_by=m.get("owned_by"), + ) + for m in models + ] + except Exception as exc: + logger.error("Failed to list models from %s: %s", payload.provider_type, exc) + raise HTTPException( + status_code=502, + detail=f"Failed to list models from {payload.provider_type}: {exc}", + ) + finally: + await client.close() diff --git a/studio/backend/storage/providers_db.py b/studio/backend/storage/providers_db.py new file mode 100644 index 0000000000..079c769f50 --- /dev/null +++ b/studio/backend/storage/providers_db.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +SQLite storage for external LLM provider configurations. + +Follows the same pattern as studio_db.py — module-level functions, +raw sqlite3, WAL mode, per-function connections. + +NOTE: API keys are NOT stored here. They live only in the browser +(localStorage) and are sent encrypted per-request. +""" + +import logging +import sqlite3 +import threading +from datetime import datetime, timezone +from typing import Optional + +logger = logging.getLogger(__name__) + +from utils.paths import studio_db_path, ensure_dir + +_schema_lock = threading.Lock() +_schema_ready = False + + +def _ensure_schema(conn: sqlite3.Connection) -> None: + """Create the llm_providers table if it doesn't exist. Called once per process.""" + conn.execute("PRAGMA journal_mode=WAL") + conn.execute( + """ + CREATE TABLE IF NOT EXISTS llm_providers ( + id TEXT NOT NULL PRIMARY KEY, + provider_type TEXT NOT NULL, + display_name TEXT NOT NULL, + base_url TEXT NOT NULL, + is_enabled INTEGER NOT NULL DEFAULT 1, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ) + """ + ) + + +def get_connection() -> sqlite3.Connection: + """Open studio.db with WAL mode, create table once per process.""" + global _schema_ready + db_path = studio_db_path() + ensure_dir(db_path.parent) + conn = sqlite3.connect(str(db_path)) + conn.row_factory = sqlite3.Row + if not _schema_ready: + with _schema_lock: + if not _schema_ready: + try: + _ensure_schema(conn) + _schema_ready = True + except Exception: + conn.close() + raise + return conn + + +def create_provider( + id: str, + provider_type: str, + display_name: str, + base_url: str, +) -> None: + """Insert a new provider configuration.""" + now = datetime.now(timezone.utc).isoformat() + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO llm_providers (id, provider_type, display_name, base_url, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?) + """, + (id, provider_type, display_name, base_url, now, now), + ) + conn.commit() + finally: + conn.close() + + +def update_provider( + id: str, + display_name: Optional[str] = None, + base_url: Optional[str] = None, + is_enabled: Optional[bool] = None, +) -> bool: + """Update fields on an existing provider. Returns True if a row was updated.""" + updates = [] + params = [] + if display_name is not None: + updates.append("display_name = ?") + params.append(display_name) + if base_url is not None: + updates.append("base_url = ?") + params.append(base_url) + if is_enabled is not None: + updates.append("is_enabled = ?") + params.append(1 if is_enabled else 0) + if not updates: + return False + updates.append("updated_at = ?") + params.append(datetime.now(timezone.utc).isoformat()) + params.append(id) + + conn = get_connection() + try: + cursor = conn.execute( + f"UPDATE llm_providers SET {', '.join(updates)} WHERE id = ?", + params, + ) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + +def delete_provider(id: str) -> bool: + """Delete a provider by ID. Returns True if a row was deleted.""" + conn = get_connection() + try: + cursor = conn.execute("DELETE FROM llm_providers WHERE id = ?", (id,)) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + +def get_provider(id: str) -> Optional[dict]: + """Fetch a single provider by ID.""" + conn = get_connection() + try: + row = conn.execute( + "SELECT * FROM llm_providers WHERE id = ?", (id,) + ).fetchone() + return dict(row) if row else None + finally: + conn.close() + + +def list_providers() -> list[dict]: + """List all provider configurations, ordered by creation time.""" + conn = get_connection() + try: + rows = conn.execute( + "SELECT * FROM llm_providers ORDER BY created_at" + ).fetchall() + return [dict(row) for row in rows] + finally: + conn.close() diff --git a/studio/backend/tests/test_providers_api.py b/studio/backend/tests/test_providers_api.py new file mode 100644 index 0000000000..3d0184a773 --- /dev/null +++ b/studio/backend/tests/test_providers_api.py @@ -0,0 +1,441 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Integration tests for the external providers API. + +Requires a running Unsloth Studio server. Configure via environment variables: + + export STUDIO_TEST_URL="http://localhost:8888" # default + export STUDIO_TEST_USER="unsloth" # default + export STUDIO_TEST_PASSWORD="..." # required — see .bootstrap_password + + # Provider API keys — any left unset will have their tests automatically skipped + export OPENAI_API_KEY="sk-..." + export MISTRAL_API_KEY="..." + export GOOGLE_API_KEY="..." + export COHERE_API_KEY="..." + export TOGETHER_API_KEY="..." + export FIREWORKS_API_KEY="..." + export PERPLEXITY_API_KEY="..." + +Run: + cd studio/backend + pytest tests/test_providers_api.py -v -s +""" + +import base64 +import json +import os + +import pytest +import requests +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import padding + +# ── Configuration ───────────────────────────────────────────────── + +BASE_URL = os.getenv("STUDIO_TEST_URL", "http://localhost:8888") +USERNAME = os.getenv("STUDIO_TEST_USER", "unsloth") +PASSWORD = os.getenv("STUDIO_TEST_PASSWORD", "") + +# Map provider_type → (env var name, model to use for inference test) +_PROVIDER_CONFIGS: dict[str, tuple[str, str]] = { + "openai": ("OPENAI_API_KEY", "gpt-4o-mini"), + "mistral": ("MISTRAL_API_KEY", "mistral-small-latest"), + "google": ("GOOGLE_API_KEY", "gemini-2.5-flash"), + "cohere": ("COHERE_API_KEY", "command-r"), + "together": ("TOGETHER_API_KEY", "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"), + "fireworks": ("FIREWORKS_API_KEY", "accounts/fireworks/models/llama4-maverick-instruct-basic"), + "perplexity": ("PERPLEXITY_API_KEY", "sonar"), +} + +PROVIDER_KEYS: dict[str, str] = { + ptype: os.getenv(env_var, "") + for ptype, (env_var, _) in _PROVIDER_CONFIGS.items() +} + +EXPECTED_PROVIDER_TYPES = set(_PROVIDER_CONFIGS.keys()) + +# ── Helpers ──────────────────────────────────────────────────────── + + +def _url(path: str) -> str: + return f"{BASE_URL}/{path.lstrip('/')}" + + +def _parse_sse_stream(response: requests.Response) -> tuple[str, bool]: + """ + Read a streaming SSE response and return (assembled_text, saw_done). + + Each chunk is a JSON object with choices[0].delta.content. + The stream ends with `data: [DONE]`. + """ + reply_parts: list[str] = [] + saw_done = False + + for raw_line in response.iter_lines(): + if isinstance(raw_line, bytes): + raw_line = raw_line.decode("utf-8") + if not raw_line.startswith("data:"): + continue + data = raw_line[len("data:"):].strip() + if data == "[DONE]": + saw_done = True + break + try: + chunk = json.loads(data) + # Handle both error payloads and normal chunks + if "error" in chunk: + raise RuntimeError(f"Provider error in stream: {chunk['error']}") + delta = chunk.get("choices", [{}])[0].get("delta", {}) + content = delta.get("content") or "" + if content: + reply_parts.append(content) + except (json.JSONDecodeError, IndexError, KeyError): + pass # skip malformed lines + + return "".join(reply_parts), saw_done + + +# ── Session-scoped fixtures ──────────────────────────────────────── + + +@pytest.fixture(scope="session") +def auth_headers() -> dict[str, str]: + """Log in once per test session and return auth headers.""" + assert PASSWORD, ( + "STUDIO_TEST_PASSWORD is not set.\n" + "Run: export STUDIO_TEST_PASSWORD=$(cat studio/backend/.bootstrap_password)" + ) + resp = requests.post( + _url("/api/auth/login"), + json={"username": USERNAME, "password": PASSWORD}, + timeout=10, + ) + assert resp.status_code == 200, f"Login failed ({resp.status_code}): {resp.text}" + token = resp.json()["access_token"] + assert token, "access_token is empty" + return {"Authorization": f"Bearer {token}"} + + +@pytest.fixture(scope="session") +def public_key_pem(auth_headers: dict[str, str]) -> str: + """Fetch RSA public key PEM once per session.""" + resp = requests.get( + _url("/api/providers/public-key"), + headers=auth_headers, + timeout=10, + ) + assert resp.status_code == 200, f"Public key fetch failed: {resp.text}" + pem = resp.json().get("public_key", "") + assert pem.startswith("-----BEGIN PUBLIC KEY-----"), "Not a valid PEM public key" + return pem + + +@pytest.fixture(scope="session") +def encrypt_key(public_key_pem: str): + """ + Return a callable encrypt_key(plaintext: str) -> str (base64 RSA-OAEP ciphertext). + Uses the backend's RSA public key — mirrors what the frontend does. + """ + # Decode PEM → load RSA public key + pem_bytes = public_key_pem.encode("utf-8") + rsa_pub = serialization.load_pem_public_key(pem_bytes) + + def _encrypt(plaintext: str) -> str: + ciphertext = rsa_pub.encrypt( + plaintext.encode("utf-8"), + padding.OAEP( + mgf=padding.MGF1(algorithm=hashes.SHA256()), + algorithm=hashes.SHA256(), + label=None, + ), + ) + return base64.b64encode(ciphertext).decode("utf-8") + + return _encrypt + + +# ── TestAuth ──────────────────────────────────────────────────────── + + +class TestAuth: + def test_login_returns_token(self): + """POST /api/auth/login returns a non-empty access_token.""" + assert PASSWORD, "STUDIO_TEST_PASSWORD not set" + resp = requests.post( + _url("/api/auth/login"), + json={"username": USERNAME, "password": PASSWORD}, + timeout=10, + ) + assert resp.status_code == 200, f"Login failed ({resp.status_code}): {resp.text}" + body = resp.json() + assert body.get("access_token"), "access_token is missing or empty" + assert body.get("token_type") == "bearer" + print(f"\n token_type={body['token_type']}, must_change_password={body.get('must_change_password')}") + + +# ── TestPublicKey ──────────────────────────────────────────────────── + + +class TestPublicKey: + def test_public_key_is_valid_pem(self, auth_headers: dict[str, str], public_key_pem: str): + """GET /api/providers/public-key returns an importable RSA PEM key.""" + pem_bytes = public_key_pem.encode("utf-8") + key = serialization.load_pem_public_key(pem_bytes) + key_size = key.key_size # type: ignore[attr-defined] + assert key_size >= 2048, f"Key size too small: {key_size}" + print(f"\n RSA-{key_size} public key OK") + + +# ── TestRegistry ──────────────────────────────────────────────────── + + +class TestRegistry: + def test_registry_returns_7_providers(self, auth_headers: dict[str, str]): + """GET /api/providers/registry returns all 7 supported providers.""" + resp = requests.get( + _url("/api/providers/registry"), + headers=auth_headers, + timeout=10, + ) + assert resp.status_code == 200, f"Registry failed: {resp.text}" + providers = resp.json() + assert len(providers) == 7, f"Expected 7 providers, got {len(providers)}: {providers}" + print(f"\n {'Provider':<12} {'Base URL'}") + print(f" {'-'*12} {'-'*45}") + for p in providers: + print(f" {p['provider_type']:<12} {p['base_url']}") + + def test_registry_has_expected_types(self, auth_headers: dict[str, str]): + """All 7 provider_type values are present in the registry.""" + resp = requests.get( + _url("/api/providers/registry"), + headers=auth_headers, + timeout=10, + ) + assert resp.status_code == 200 + returned_types = {p["provider_type"] for p in resp.json()} + missing = EXPECTED_PROVIDER_TYPES - returned_types + assert not missing, f"Missing provider types: {missing}" + + def test_registry_entries_have_required_fields(self, auth_headers: dict[str, str]): + """Each registry entry has provider_type, display_name, base_url, default_models.""" + resp = requests.get(_url("/api/providers/registry"), headers=auth_headers, timeout=10) + assert resp.status_code == 200 + for entry in resp.json(): + for field in ("provider_type", "display_name", "base_url", "default_models"): + assert field in entry, f"Missing field '{field}' in entry: {entry}" + assert isinstance(entry["default_models"], list) + assert len(entry["default_models"]) > 0 + + +# ── TestProviderCRUD ──────────────────────────────────────────────── + + +class TestProviderCRUD: + """ + These tests run sequentially within the class and share state via class variables. + They create, read, update, and delete a single test provider config. + """ + + _created_id: str = "" + + def test_create_provider(self, auth_headers: dict[str, str]): + """POST /api/providers/ creates a provider config and returns 201.""" + resp = requests.post( + _url("/api/providers/"), + headers=auth_headers, + json={"provider_type": "openai", "display_name": "Test OpenAI (pytest)"}, + timeout=10, + ) + assert resp.status_code == 201, f"Create failed ({resp.status_code}): {resp.text}" + body = resp.json() + assert body.get("id"), "No id in response" + assert body["provider_type"] == "openai" + assert body["display_name"] == "Test OpenAI (pytest)" + assert body["is_enabled"] is True + TestProviderCRUD._created_id = body["id"] + print(f"\n created id={body['id']}") + + def test_list_includes_created(self, auth_headers: dict[str, str]): + """GET /api/providers/ includes the newly created config.""" + assert TestProviderCRUD._created_id, "No created_id (run test_create_provider first)" + resp = requests.get(_url("/api/providers/"), headers=auth_headers, timeout=10) + assert resp.status_code == 200 + ids = [p["id"] for p in resp.json()] + assert TestProviderCRUD._created_id in ids, ( + f"Created id {TestProviderCRUD._created_id!r} not found in list: {ids}" + ) + print(f"\n found id={TestProviderCRUD._created_id} in list of {len(ids)}") + + def test_update_display_name(self, auth_headers: dict[str, str]): + """PUT /api/providers/{id} updates the display_name.""" + assert TestProviderCRUD._created_id, "No created_id" + new_name = "Test OpenAI (pytest updated)" + resp = requests.put( + _url(f"/api/providers/{TestProviderCRUD._created_id}"), + headers=auth_headers, + json={"display_name": new_name}, + timeout=10, + ) + assert resp.status_code == 200, f"Update failed ({resp.status_code}): {resp.text}" + assert resp.json()["display_name"] == new_name + print(f"\n updated display_name to '{new_name}'") + + def test_delete_provider(self, auth_headers: dict[str, str]): + """DELETE /api/providers/{id} removes the config (204) and it's gone from list.""" + assert TestProviderCRUD._created_id, "No created_id" + resp = requests.delete( + _url(f"/api/providers/{TestProviderCRUD._created_id}"), + headers=auth_headers, + timeout=10, + ) + assert resp.status_code == 204, f"Delete failed ({resp.status_code}): {resp.text}" + + # Confirm gone from list + list_resp = requests.get(_url("/api/providers/"), headers=auth_headers, timeout=10) + ids = [p["id"] for p in list_resp.json()] + assert TestProviderCRUD._created_id not in ids, "Deleted provider still in list" + print(f"\n deleted id={TestProviderCRUD._created_id} confirmed gone") + + +# ── TestProviderInference ──────────────────────────────────────────── + + +# Build parametrize list: (provider_type, model, api_key) for configured providers only +_INFERENCE_PARAMS = [ + pytest.param( + ptype, + model, + PROVIDER_KEYS.get(ptype, ""), + id=ptype, + marks=pytest.mark.skipif( + not PROVIDER_KEYS.get(ptype, ""), + reason=f"no {env_var} set", + ), + ) + for ptype, (env_var, model) in _PROVIDER_CONFIGS.items() +] + + +class TestProviderInference: + """ + Live inference tests — one parametrized set per provider. + Each test is automatically skipped when the provider's API key env var is not set. + """ + + @pytest.mark.parametrize("provider_type,model,api_key", _INFERENCE_PARAMS) + def test_connection( + self, + auth_headers: dict[str, str], + encrypt_key, + provider_type: str, + model: str, + api_key: str, + ): + """POST /api/providers/test → success: true.""" + encrypted = encrypt_key(api_key) + resp = requests.post( + _url("/api/providers/test"), + headers=auth_headers, + json={"provider_type": provider_type, "encrypted_api_key": encrypted}, + timeout=30, + ) + assert resp.status_code == 200, f"Request failed ({resp.status_code}): {resp.text}" + body = resp.json() + assert body["success"] is True, ( + f"Connection test failed for {provider_type}: {body.get('message')}" + ) + print(f"\n [{provider_type}] connection OK — {body['message']}") + + @pytest.mark.parametrize("provider_type,model,api_key", _INFERENCE_PARAMS) + def test_list_models( + self, + auth_headers: dict[str, str], + encrypt_key, + provider_type: str, + model: str, + api_key: str, + ): + """POST /api/providers/models → non-empty list, print first 3.""" + encrypted = encrypt_key(api_key) + resp = requests.post( + _url("/api/providers/models"), + headers=auth_headers, + json={"provider_type": provider_type, "encrypted_api_key": encrypted}, + timeout=30, + ) + assert resp.status_code == 200, f"Request failed ({resp.status_code}): {resp.text}" + models = resp.json() + assert isinstance(models, list), f"Expected list, got {type(models)}" + assert len(models) > 0, f"No models returned for {provider_type}" + preview = [m["id"] for m in models[:3]] + print(f"\n [{provider_type}] {len(models)} models — first 3: {preview}") + + @pytest.mark.parametrize("provider_type,model,api_key", _INFERENCE_PARAMS) + def test_chat_inference( + self, + auth_headers: dict[str, str], + encrypt_key, + provider_type: str, + model: str, + api_key: str, + ): + """POST /v1/chat/completions with provider fields → streamed reply.""" + encrypted = encrypt_key(api_key) + payload = { + "messages": [{"role": "user", "content": "Say hello in one sentence."}], + "stream": True, + "temperature": 0.7, + "max_tokens": 64, + "provider_type": provider_type, + "external_model": model, + "encrypted_api_key": encrypted, + } + with requests.post( + _url("/v1/chat/completions"), + headers={**auth_headers, "Content-Type": "application/json"}, + json=payload, + stream=True, + timeout=60, + ) as resp: + assert resp.status_code == 200, ( + f"Chat completions failed ({resp.status_code}): {resp.text[:500]}" + ) + reply, saw_done = _parse_sse_stream(resp) + + assert reply.strip(), f"Empty reply from {provider_type}/{model}" + assert saw_done, f"Stream did not end with [DONE] for {provider_type}/{model}" + print(f'\n [{provider_type}/{model}] reply: "{reply.strip()}"') + + +# ── TestLocalInferenceUnaffected ──────────────────────────────────── + + +class TestLocalInferenceUnaffected: + def test_chat_without_provider(self, auth_headers: dict[str, str]): + """ + POST /v1/chat/completions without provider fields must not return 422 or 500. + + 200 = a local model is loaded and responded. + 503 = no model loaded (expected in test environment — that's fine). + Any other 4xx/5xx (except 503) = regression in request handling. + """ + resp = requests.post( + _url("/v1/chat/completions"), + headers={**auth_headers, "Content-Type": "application/json"}, + json={ + "messages": [{"role": "user", "content": "Hello"}], + "stream": False, + }, + timeout=15, + ) + allowed = {200, 503} + assert resp.status_code in allowed, ( + f"Unexpected status {resp.status_code} for local inference path: {resp.text[:300]}\n" + f"This likely means the provider fields broke the base request schema." + ) + status_label = "local model responded" if resp.status_code == 200 else "no model loaded (expected)" + print(f"\n status={resp.status_code} ({status_label}) — local path unaffected") diff --git a/studio/docs/external-providers-frontend-spec.md b/studio/docs/external-providers-frontend-spec.md new file mode 100644 index 0000000000..e5fbf2991b --- /dev/null +++ b/studio/docs/external-providers-frontend-spec.md @@ -0,0 +1,327 @@ +# External Providers — Frontend Integration Spec + +## Overview + +The backend proxies chat requests to external LLM providers (OpenAI, Mistral, Gemini, Cohere, Together AI, Fireworks AI, Perplexity). **API keys are never stored on the backend** — the frontend holds them in localStorage and encrypts them before every request using the server's RSA public key. + +--- + +## 1. API Key Encryption + +The server generates an RSA-2048 key pair on startup (rotates on restart). The frontend must: + +1. **Fetch the public key** on app load (and after any 400 "public key may have changed" error) +2. **Encrypt API keys** with RSA-OAEP + SHA-256 before including in any request +3. **Base64-encode** the ciphertext + +```ts +// Fetch once on load, cache in memory +const res = await authFetch("GET", "/api/providers/public-key"); +const pem: string = res.public_key; + +// Helper: PEM string -> ArrayBuffer +function pemToBuffer(pem: string): ArrayBuffer { + const b64 = pem.replace(/-----[^-]+-----/g, "").replace(/\s/g, ""); + const bin = atob(b64); + const buf = new Uint8Array(bin.length); + for (let i = 0; i < bin.length; i++) buf[i] = bin.charCodeAt(i); + return buf.buffer; +} + +// Import the key (do once) +const cryptoKey = await crypto.subtle.importKey( + "spki", + pemToBuffer(pem), + { name: "RSA-OAEP", hash: "SHA-256" }, + false, + ["encrypt"], +); + +// Encrypt before each request +async function encryptApiKey(plaintext: string): Promise { + const encoded = new TextEncoder().encode(plaintext); + const encrypted = await crypto.subtle.encrypt( + { name: "RSA-OAEP" }, + cryptoKey, + encoded, + ); + return btoa(String.fromCharCode(...new Uint8Array(encrypted))); +} +``` + +No npm packages needed — `crypto.subtle` is a native browser API. + +--- + +## 2. Endpoints + +All endpoints require auth (`Authorization: Bearer `). + +### GET `/api/providers/public-key` + +Returns the RSA public key for encrypting API keys. + +```json +// Response +{ "public_key": "-----BEGIN PUBLIC KEY-----\nMIIBI..." } +``` + +### GET `/api/providers/registry` + +Returns all supported provider types with defaults. Use this to populate the "Add Provider" dropdown. + +```json +// Response +[ + { + "provider_type": "openai", + "display_name": "OpenAI", + "base_url": "https://api.openai.com/v1", + "default_models": ["gpt-4o", "gpt-4o-mini", "gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano", "o3-mini"], + "supports_streaming": true, + "supports_vision": true, + "supports_tool_calling": true + } + // ... 6 more providers +] +``` + +**Supported `provider_type` values:** `openai`, `mistral`, `google`, `cohere`, `together`, `fireworks`, `perplexity` + +### GET `/api/providers` + +List saved provider configs. These store display name + base URL — **not** the API key. + +```json +// Response +[ + { + "id": "a1b2c3d4e5f67890", + "provider_type": "openai", + "display_name": "My OpenAI Key", + "base_url": "https://api.openai.com/v1", + "is_enabled": true, + "created_at": "2026-03-29T...", + "updated_at": "2026-03-29T..." + } +] +``` + +### POST `/api/providers` + +Create a saved provider config. No API key sent here. + +```json +// Request +{ + "provider_type": "openai", + "display_name": "My OpenAI Key", + "base_url": null // omit to use registry default +} + +// Response (201) +{ + "id": "a1b2c3d4e5f67890", + "provider_type": "openai", + "display_name": "My OpenAI Key", + "base_url": "https://api.openai.com/v1", + "is_enabled": true, + "created_at": "2026-03-29T...", + "updated_at": "2026-03-29T..." +} +``` + +### PUT `/api/providers/{id}` + +Update a provider config. All fields optional. + +```json +// Request +{ + "display_name": "Work OpenAI", + "base_url": null, + "is_enabled": false +} + +// Response -> same shape as GET items +``` + +### DELETE `/api/providers/{id}` + +Delete a provider config. Returns `204 No Content`. + +### POST `/api/providers/test` + +Test if an API key works. Encrypted key required. + +```json +// Request +{ + "provider_type": "openai", + "encrypted_api_key": "", + "base_url": null // optional override +} + +// Response +{ + "success": true, + "message": "Connected successfully. Found 42 model(s).", + "models_count": 42 +} +``` + +### POST `/api/providers/models` + +List available models from a provider. Encrypted key required. + +```json +// Request +{ + "provider_type": "openai", + "encrypted_api_key": "", + "base_url": null +} + +// Response +[ + { + "id": "gpt-4o", + "display_name": "gpt-4o", + "context_length": 128000, + "owned_by": "openai" + } + // ... +] +``` + +--- + +## 3. Chatting with an External Provider + +Use the **existing** `POST /v1/chat/completions` endpoint — just add provider fields. The backend detects them and proxies to the external API instead of local inference. + +```json +// Request +{ + "messages": [ + { "role": "user", "content": "Hello!" } + ], + "stream": true, + "temperature": 0.7, + "top_p": 0.95, + "max_tokens": 1024, + "presence_penalty": 0.0, + + // --- These trigger external routing --- + "provider_type": "openai", // required (or provider_id) + "external_model": "gpt-4o-mini", // required - model ID at the provider + "encrypted_api_key": "", // required - RSA-encrypted + "provider_id": "a1b2c3d4e5f67890", // optional - saved config ID (for base_url lookup) + "provider_base_url": null // optional - override base URL +} +``` + +### Routing logic + +- If `encrypted_api_key` + (`provider_type` or `provider_id`) are present -> routes to external provider +- Otherwise -> routes to local inference (existing behavior, unchanged) + +### Response format + +Standard OpenAI SSE streaming — same format the frontend already handles. No changes needed to the stream parser. + +``` +data: {"id":"chatcmpl-...","object":"chat.completion.chunk","choices":[{"delta":{"role":"assistant","content":"Hi"},"index":0}]} + +data: {"id":"chatcmpl-...","object":"chat.completion.chunk","choices":[{"delta":{"content":" there!"},"index":0,"finish_reason":"stop"}]} + +data: [DONE] +``` + +### Error format + +If the provider returns an error, it comes as an SSE event: + +``` +data: {"error":{"message":"Invalid API key","type":"provider_error","code":"401","provider":"openai"}} +``` + +### Fields to skip for external providers + +When chatting with an external provider, these local-only fields should **not** be sent (they will be ignored but are unnecessary): + +- `top_k`, `min_p`, `repetition_penalty` +- `image_base64`, `audio_base64` +- `use_adapter` +- `enable_thinking` +- `enable_tools`, `enabled_tools`, `auto_heal_tool_calls`, `max_tool_calls_per_message`, `tool_call_timeout` +- `session_id` + +Only send standard OpenAI fields: `messages`, `stream`, `temperature`, `top_p`, `max_tokens`, `presence_penalty`. + +--- + +## 4. Frontend Storage Model + +The frontend should store in **localStorage**: + +| Key | Value | Notes | +|-----|-------|-------| +| Provider API keys | `{ [provider_id]: "sk-abc123..." }` | Plaintext in localStorage, encrypted before sending | +| Active provider | `provider_id` or `null` | `null` = local inference | +| Active external model | `"gpt-4o-mini"` or `null` | Model ID at the selected provider | + +--- + +## 5. Typical User Flow + +1. User opens **Provider Settings page** +2. Frontend calls `GET /api/providers/registry` -> shows available provider types +3. User picks "OpenAI", enters API key, names it "My OpenAI" +4. Frontend calls `POST /api/providers` to save config (without key) +5. Frontend stores the API key in localStorage keyed by the returned `id` +6. Frontend encrypts key -> calls `POST /api/providers/test` -> shows success/fail +7. User goes to **Chat**, selects the provider + model from the model selector +8. On each message, frontend encrypts the key and adds `provider_type`, `external_model`, `encrypted_api_key` to the existing chat completions request +9. Response streams back in the same SSE format as local inference — no parser changes needed + +--- + +## 6. UI Considerations + +### Provider Settings Page (separate page in app nav) + +- List configured providers with enable/disable toggles +- "Add Provider" form: dropdown (from registry), API key input, display name +- "Test Connection" button per provider +- Edit / Delete actions per provider +- Show provider status (connected / error) + +### Chat Model Selector + +- Add a "Cloud" or "API" section alongside local models +- Group external models by provider (e.g. "OpenAI / gpt-4o") +- When external model selected: set `activeProviderId` + `activeExternalModel` in store +- When local model selected: clear provider state (back to `null`) + +### Chat Page Adaptations + +When an external provider is active: + +- **Hide** local-only UI: context length bar, GGUF settings, LoRA controls, KV cache dtype +- **Hide** local-only features: reasoning toggle, tool calling controls (unless provider supports them) +- **Show** simplified params: temperature, top_p, max_tokens, presence_penalty only +- **Skip** auto-load logic (no local model needed) +- **Show** provider badge/icon next to model name + +--- + +## 7. Error Handling + +| Scenario | How to detect | What to do | +|----------|---------------|------------| +| Public key rotated (server restarted) | 400 with "public key may have changed" | Re-fetch `GET /api/providers/public-key`, re-encrypt, retry | +| Invalid API key | SSE error with `code: "401"` | Show "API key invalid or expired" message | +| Provider down | SSE error with `code: "502"` or `code: "504"` | Show "Provider unavailable, try again later" | +| Rate limited | SSE error with `code: "429"` | Show "Rate limited, please wait" | +| Unknown provider type | 400 from POST endpoints | Should not happen if using registry dropdown |