diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index e1bdd0ca0c..b57c89f1b4 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -30,7 +30,7 @@ class ExternalProviderClient: 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)) + self._client = httpx.AsyncClient(timeout = httpx.Timeout(timeout, connect = 10.0)) def _auth_headers(self) -> dict[str, str]: """Build authentication headers using the provider's registry config.""" @@ -101,12 +101,12 @@ class ExternalProviderClient: async with self._client.stream( "POST", url, - json=body, - headers=self._auth_headers(), + 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") + error_text = error_body.decode("utf-8", errors = "replace") logger.error( "External provider returned %d: %s", response.status_code, @@ -127,7 +127,7 @@ class ExternalProviderClient: if line.strip(): yield line except GeneratorExit: - await response.aclose() # set PoolByteStream._closed=True FIRST + await response.aclose() # set PoolByteStream._closed=True FIRST await lines_gen.aclose() # now safe — aclose() is a no-op raise finally: @@ -137,17 +137,23 @@ class ExternalProviderClient: 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 + 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 + 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 + 502, + f"Error communicating with {self.provider_type}: {exc}", + self.provider_type, ) async def _stream_anthropic( @@ -175,8 +181,13 @@ class ExternalProviderClient: for msg in messages: if msg.get("role") == "system": content = msg.get("content", "") - system = content if isinstance(content, str) else \ - "\n".join(p["text"] for p in content if p.get("type") == "text") + system = ( + content + if isinstance(content, str) + else "\n".join( + p["text"] for p in content if p.get("type") == "text" + ) + ) continue content = msg.get("content") @@ -191,24 +202,31 @@ class ExternalProviderClient: if url.startswith("data:"): # data:image/png;base64, → split header and data header, _, b64data = url.partition(",") - media_type = header.split(";")[0].replace("data:", "") or "image/jpeg" - anthropic_parts.append({ - "type": "image", - "source": { - "type": "base64", - "media_type": media_type, - "data": b64data, - }, - }) + media_type = ( + header.split(";")[0].replace("data:", "") + or "image/jpeg" + ) + anthropic_parts.append( + { + "type": "image", + "source": { + "type": "base64", + "media_type": media_type, + "data": b64data, + }, + } + ) else: # Remote URL — Anthropic supports url source type natively - anthropic_parts.append({ - "type": "image", - "source": { - "type": "url", - "url": url, - }, - }) + anthropic_parts.append( + { + "type": "image", + "source": { + "type": "url", + "url": url, + }, + } + ) filtered.append({"role": msg["role"], "content": anthropic_parts}) else: filtered.append(msg) @@ -239,16 +257,20 @@ class ExternalProviderClient: async with self._client.stream( "POST", url, - json=body, - headers=self._auth_headers(), + 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") + error_text = error_body.decode("utf-8", errors = "replace") logger.error( - "Anthropic returned %d: %s", response.status_code, error_text[:500] + "Anthropic returned %d: %s", + response.status_code, + error_text[:500], + ) + yield _error_sse_line( + response.status_code, error_text, self.provider_type ) - yield _error_sse_line(response.status_code, error_text, self.provider_type) return lines_gen = response.aiter_lines().__aiter__() @@ -263,7 +285,7 @@ class ExternalProviderClient: if not line.startswith("data:"): continue - data_str = line[len("data:"):].strip() + data_str = line[len("data:") :].strip() if not data_str: continue @@ -280,11 +302,13 @@ class ExternalProviderClient: chunk = { "id": completion_id, "object": "chat.completion.chunk", - "choices": [{ - "index": 0, - "delta": {"content": delta.get("text", "")}, - "finish_reason": None, - }], + "choices": [ + { + "index": 0, + "delta": {"content": delta.get("text", "")}, + "finish_reason": None, + } + ], } yield f"data: {_json.dumps(chunk)}" @@ -294,20 +318,26 @@ class ExternalProviderClient: chunk = { "id": completion_id, "object": "chat.completion.chunk", - "choices": [{ - "index": 0, - "delta": {}, - "finish_reason": _finish_reason_map.get(stop_reason, "stop"), - }], + "choices": [ + { + "index": 0, + "delta": {}, + "finish_reason": _finish_reason_map.get( + stop_reason, "stop" + ), + } + ], } yield f"data: {_json.dumps(chunk)}" elif event_type == "message_stop": yield "data: [DONE]" - await response.aclose() # set PoolByteStream._closed=True FIRST + await ( + response.aclose() + ) # set PoolByteStream._closed=True FIRST break except GeneratorExit: - await response.aclose() # set PoolByteStream._closed=True FIRST + await response.aclose() # set PoolByteStream._closed=True FIRST await lines_gen.aclose() # now safe — aclose() is a no-op raise finally: @@ -316,13 +346,25 @@ class ExternalProviderClient: 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) + 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) + 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) + yield _error_sse_line( + 502, + f"Error communicating with {self.provider_type}: {exc}", + self.provider_type, + ) async def chat_completion( self, @@ -347,8 +389,8 @@ class ExternalProviderClient: response = await self._client.post( f"{self.base_url}/chat/completions", - json=body, - headers=self._auth_headers(), + json = body, + headers = self._auth_headers(), ) response.raise_for_status() return response.json() @@ -363,7 +405,7 @@ class ExternalProviderClient: try: response = await self._client.get( f"{self.base_url}/models", - headers=self._auth_headers(), + headers = self._auth_headers(), ) response.raise_for_status() data = response.json() diff --git a/studio/backend/core/inference/key_exchange.py b/studio/backend/core/inference/key_exchange.py index c92bd9d9d9..982b54501a 100644 --- a/studio/backend/core/inference/key_exchange.py +++ b/studio/backend/core/inference/key_exchange.py @@ -29,13 +29,17 @@ 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_exponent = 65537, + key_size = 2048, + ) + _public_key_pem = ( + _private_key.public_key() + .public_bytes( + serialization.Encoding.PEM, + serialization.PublicFormat.SubjectPublicKeyInfo, + ) + .decode("utf-8") ) - _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") @@ -63,9 +67,9 @@ def decrypt_api_key(encrypted_b64: str) -> str: plaintext = _private_key.decrypt( ciphertext, padding.OAEP( - mgf=padding.MGF1(algorithm=hashes.SHA256()), - algorithm=hashes.SHA256(), - label=None, + mgf = padding.MGF1(algorithm = hashes.SHA256()), + algorithm = hashes.SHA256(), + label = None, ), ) return plaintext.decode("utf-8") diff --git a/studio/backend/main.py b/studio/backend/main.py index 590039b151..d86907f763 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -116,6 +116,7 @@ async def lifespan(app: FastAPI): # 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(): diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py index 572951cffe..d9eb6703ae 100644 --- a/studio/backend/models/providers.py +++ b/studio/backend/models/providers.py @@ -16,15 +16,23 @@ from pydantic import BaseModel, Field 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" + 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" ) - 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 ────────────────────────────────────────── @@ -33,32 +41,36 @@ class ProviderRegistryEntry(BaseModel): 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')") + 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.", + 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") + 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") + 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 ───────────────────────────────────────────────── @@ -67,19 +79,23 @@ class ProviderResponse(BaseModel): 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") + 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") + 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)" + None, description = "Custom base URL (overrides registry default)" ) @@ -89,18 +105,20 @@ class ProviderModelsRequest(BaseModel): 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") + 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)" + 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") + 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)" + None, description = "Number of models found (if test succeeded)" ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9e7453b7b0..966aebd6ff 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -869,10 +869,12 @@ def _build_external_messages( if part.type == "text": parts.append({"type": "text", "text": part.text}) elif part.type == "image_url": - parts.append({ - "type": "image_url", - "image_url": {"url": part.image_url.url}, - }) + parts.append( + { + "type": "image_url", + "image_url": {"url": part.image_url.url}, + } + ) result.append({"role": msg.role, "content": parts}) else: # Non-vision provider — strip images, keep text only @@ -944,6 +946,7 @@ async def _proxy_to_external_provider( # Build messages preserving multimodal content for vision-capable providers from core.inference.providers import get_provider_info as _get_provider_info + _pinfo = _get_provider_info(provider_type) or {} _supports_vision = _pinfo.get("supports_vision", False) chat_messages = _build_external_messages(payload.messages, _supports_vision) diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index 0cf891fabc..ede58d04e4 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -55,7 +55,7 @@ async def get_public_key( # ── Provider registry (static) ─────────────────────────────────── -@router.get("/registry", response_model=list[ProviderRegistryEntry]) +@router.get("/registry", response_model = list[ProviderRegistryEntry]) async def list_registry( current_subject: str = Depends(get_current_subject), ): @@ -66,7 +66,7 @@ async def list_registry( # ── Provider config CRUD ────────────────────────────────────────── -@router.get("/", response_model=list[ProviderResponse]) +@router.get("/", response_model = list[ProviderResponse]) async def list_provider_configs( current_subject: str = Depends(get_current_subject), ): @@ -74,19 +74,19 @@ async def list_provider_configs( 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"], + 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) +@router.post("/", response_model = ProviderResponse, status_code = 201) async def create_provider_config( payload: ProviderCreate, current_subject: str = Depends(get_current_subject), @@ -95,8 +95,8 @@ async def create_provider_config( info = get_provider_info(payload.provider_type) if info is None: raise HTTPException( - status_code=400, - detail=f"Unknown provider type: {payload.provider_type}. " + status_code = 400, + detail = f"Unknown provider type: {payload.provider_type}. " f"Use GET /api/providers/registry to see available types.", ) @@ -104,25 +104,25 @@ async def create_provider_config( 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, + 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"], + 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) +@router.put("/{provider_id}", response_model = ProviderResponse) async def update_provider_config( provider_id: str, payload: ProviderUpdate, @@ -131,30 +131,30 @@ async def update_provider_config( """Update a saved provider configuration.""" existing = providers_db.get_provider(provider_id) if not existing: - raise HTTPException(status_code=404, detail="Provider not found") + 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, + 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") + 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"], + 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) +@router.delete("/{provider_id}", status_code = 204) async def delete_provider_config( provider_id: str, current_subject: str = Depends(get_current_subject), @@ -162,13 +162,13 @@ async def delete_provider_config( """Delete a saved provider configuration.""" deleted = providers_db.delete_provider(provider_id) if not deleted: - raise HTTPException(status_code=404, detail="Provider not found") + raise HTTPException(status_code = 404, detail = "Provider not found") # ── Test connectivity ───────────────────────────────────────────── -@router.post("/test", response_model=ProviderTestResult) +@router.post("/test", response_model = ProviderTestResult) async def test_provider( payload: ProviderTestRequest, current_subject: str = Depends(get_current_subject), @@ -182,8 +182,8 @@ async def test_provider( info = get_provider_info(payload.provider_type) if info is None: raise HTTPException( - status_code=400, - detail=f"Unknown provider type: {payload.provider_type}", + status_code = 400, + detail = f"Unknown provider type: {payload.provider_type}", ) try: @@ -191,31 +191,31 @@ async def test_provider( 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.", + 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, + 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), + 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, + success = False, + message = f"Connection failed: {exc}", + models_count = None, ) finally: await client.close() @@ -224,7 +224,7 @@ async def test_provider( # ── List models from provider ───────────────────────────────────── -@router.post("/models", response_model=list[ProviderModelInfo]) +@router.post("/models", response_model = list[ProviderModelInfo]) async def list_provider_models( payload: ProviderModelsRequest, current_subject: str = Depends(get_current_subject), @@ -237,8 +237,8 @@ async def list_provider_models( info = get_provider_info(payload.provider_type) if info is None: raise HTTPException( - status_code=400, - detail=f"Unknown provider type: {payload.provider_type}", + status_code = 400, + detail = f"Unknown provider type: {payload.provider_type}", ) try: @@ -246,34 +246,34 @@ async def list_provider_models( 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.", + 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, + 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"), + 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}", + 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 index 079c769f50..ca47fcbd80 100644 --- a/studio/backend/storage/providers_db.py +++ b/studio/backend/storage/providers_db.py @@ -135,9 +135,7 @@ 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() + row = conn.execute("SELECT * FROM llm_providers WHERE id = ?", (id,)).fetchone() return dict(row) if row else None finally: conn.close() diff --git a/studio/backend/tests/test_providers_api.py b/studio/backend/tests/test_providers_api.py index b309304abc..4eaf9d8b4f 100644 --- a/studio/backend/tests/test_providers_api.py +++ b/studio/backend/tests/test_providers_api.py @@ -41,18 +41,17 @@ 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-2506"), - "gemini": ("GEMINI_API_KEY", "gemini-3-flash-preview"), - "cohere": ("COHERE_API_KEY", "command-a-03-2025"), + "openai": ("OPENAI_API_KEY", "gpt-4o-mini"), + "mistral": ("MISTRAL_API_KEY", "mistral-small-2506"), + "gemini": ("GEMINI_API_KEY", "gemini-3-flash-preview"), + "cohere": ("COHERE_API_KEY", "command-a-03-2025"), "openrouter": ("OPENROUTER_API_KEY", "openai/gpt-4o-mini"), - "anthropic": ("ANTHROPIC_API_KEY", "claude-haiku-4-5"), - "deepseek": ("DEEPSEEK_API_KEY", "deepseek-chat"), + "anthropic": ("ANTHROPIC_API_KEY", "claude-haiku-4-5"), + "deepseek": ("DEEPSEEK_API_KEY", "deepseek-chat"), } PROVIDER_KEYS: dict[str, str] = { - ptype: os.getenv(env_var, "") - for ptype, (env_var, _) in _PROVIDER_CONFIGS.items() + ptype: os.getenv(env_var, "") for ptype, (env_var, _) in _PROVIDER_CONFIGS.items() } EXPECTED_PROVIDER_TYPES = set(_PROVIDER_CONFIGS.keys()) @@ -79,7 +78,7 @@ def _parse_sse_stream(response: requests.Response) -> tuple[str, bool]: raw_line = raw_line.decode("utf-8") if not raw_line.startswith("data:"): continue - data = raw_line[len("data:"):].strip() + data = raw_line[len("data:") :].strip() if data == "[DONE]": saw_done = True break @@ -101,7 +100,7 @@ def _parse_sse_stream(response: requests.Response) -> tuple[str, bool]: # ── Session-scoped fixtures ──────────────────────────────────────── -@pytest.fixture(scope="session") +@pytest.fixture(scope = "session") def auth_headers() -> dict[str, str]: """ Log in once per session and return auth headers. @@ -125,8 +124,8 @@ def auth_headers() -> dict[str, str]: resp = requests.post( _url("/api/auth/login"), - json={"username": USERNAME, "password": PASSWORD}, - timeout=10, + json = {"username": USERNAME, "password": PASSWORD}, + timeout = 10, ) assert resp.status_code == 200, f"Login failed ({resp.status_code}): {resp.text}" body = resp.json() @@ -139,13 +138,13 @@ def auth_headers() -> dict[str, str]: new_password = os.getenv("STUDIO_TEST_NEW_PASSWORD") or f"{PASSWORD}-test" change_resp = requests.post( _url("/api/auth/change-password"), - headers={"Authorization": f"Bearer {token}"}, - json={"current_password": PASSWORD, "new_password": new_password}, - timeout=10, - ) - assert change_resp.status_code == 200, ( - f"Auto password-change failed ({change_resp.status_code}): {change_resp.text}" + headers = {"Authorization": f"Bearer {token}"}, + json = {"current_password": PASSWORD, "new_password": new_password}, + timeout = 10, ) + assert ( + change_resp.status_code == 200 + ), f"Auto password-change failed ({change_resp.status_code}): {change_resp.text}" token = change_resp.json()["access_token"] print( f"\n NOTE: Bootstrap password changed automatically.\n" @@ -156,13 +155,13 @@ def auth_headers() -> dict[str, str]: return {"Authorization": f"Bearer {token}"} -@pytest.fixture(scope="session") +@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, + headers = auth_headers, + timeout = 10, ) assert resp.status_code == 200, f"Public key fetch failed: {resp.text}" pem = resp.json().get("public_key", "") @@ -170,7 +169,7 @@ def public_key_pem(auth_headers: dict[str, str]) -> str: return pem -@pytest.fixture(scope="session") +@pytest.fixture(scope = "session") def vision_image_data_url() -> str: """ Download the sloth image once per session and return it as a base64 data URI. @@ -179,14 +178,14 @@ def vision_image_data_url() -> str: the image inline — Gemini's OpenAI-compatible layer does not fetch external HTTP URLs, so raw image_url links silently produce empty replies for Gemini. """ - resp = requests.get(_VISION_IMAGE_URL, timeout=30) + resp = requests.get(_VISION_IMAGE_URL, timeout = 30) resp.raise_for_status() content_type = resp.headers.get("Content-Type", "image/jpeg").split(";")[0].strip() b64 = base64.b64encode(resp.content).decode("utf-8") return f"data:{content_type};base64,{b64}" -@pytest.fixture(scope="session") +@pytest.fixture(scope = "session") def encrypt_key(public_key_pem: str): """ Return a callable encrypt_key(plaintext: str) -> str (base64 RSA-OAEP ciphertext). @@ -200,9 +199,9 @@ def encrypt_key(public_key_pem: str): ciphertext = rsa_pub.encrypt( plaintext.encode("utf-8"), padding.OAEP( - mgf=padding.MGF1(algorithm=hashes.SHA256()), - algorithm=hashes.SHA256(), - label=None, + mgf = padding.MGF1(algorithm = hashes.SHA256()), + algorithm = hashes.SHA256(), + label = None, ), ) return base64.b64encode(ciphertext).decode("utf-8") @@ -219,21 +218,27 @@ class TestAuth: assert PASSWORD, "STUDIO_TEST_PASSWORD not set" resp = requests.post( _url("/api/auth/login"), - json={"username": USERNAME, "password": PASSWORD}, - timeout=10, + json = {"username": USERNAME, "password": PASSWORD}, + timeout = 10, ) - assert resp.status_code == 200, f"Login failed ({resp.status_code}): {resp.text}" + 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')}") + 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): + 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) @@ -250,12 +255,14 @@ class TestRegistry: """GET /api/providers/registry returns all 7 supported providers.""" resp = requests.get( _url("/api/providers/registry"), - headers=auth_headers, - timeout=10, + 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}" + 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: @@ -265,8 +272,8 @@ class TestRegistry: """All 7 provider_type values are present in the registry.""" resp = requests.get( _url("/api/providers/registry"), - headers=auth_headers, - timeout=10, + headers = auth_headers, + timeout = 10, ) assert resp.status_code == 200 returned_types = {p["provider_type"] for p in resp.json()} @@ -275,10 +282,17 @@ class TestRegistry: 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) + 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"): + 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 @@ -299,11 +313,13 @@ class TestProviderCRUD: """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, + 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}" + 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" @@ -314,13 +330,15 @@ class TestProviderCRUD: 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 ( + 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}" - ) + 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]): @@ -329,11 +347,13 @@ class TestProviderCRUD: 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, + 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.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}'") @@ -342,13 +362,17 @@ class TestProviderCRUD: assert TestProviderCRUD._created_id, "No created_id" resp = requests.delete( _url(f"/api/providers/{TestProviderCRUD._created_id}"), - headers=auth_headers, - timeout=10, + headers = auth_headers, + timeout = 10, ) - assert resp.status_code == 204, f"Delete failed ({resp.status_code}): {resp.text}" + 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) + 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") @@ -363,10 +387,10 @@ _INFERENCE_PARAMS = [ ptype, model, PROVIDER_KEYS.get(ptype, ""), - id=ptype, - marks=pytest.mark.skipif( + id = ptype, + marks = pytest.mark.skipif( not PROVIDER_KEYS.get(ptype, ""), - reason=f"no {env_var} set", + reason = f"no {env_var} set", ), ) for ptype, (env_var, model) in _PROVIDER_CONFIGS.items() @@ -392,15 +416,17 @@ class TestProviderInference: 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, + 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}" + 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')}" - ) + 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) @@ -416,11 +442,13 @@ class TestProviderInference: 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, + 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}" + 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}" @@ -449,14 +477,14 @@ class TestProviderInference: } with requests.post( _url("/v1/chat/completions"), - headers={**auth_headers, "Content-Type": "application/json"}, - json=payload, - stream=True, - timeout=60, + 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]}" - ) + 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}" @@ -467,17 +495,19 @@ class TestProviderInference: # ── TestVisionInference ───────────────────────────────────────────── # Sloth photo — used to test vision routing across providers -_VISION_IMAGE_URL = "https://www.travelexcellence.com/images/where-to-see-sloths-in-costa-rica.jpg" +_VISION_IMAGE_URL = ( + "https://www.travelexcellence.com/images/where-to-see-sloths-in-costa-rica.jpg" +) _VISION_PARAMS = [ pytest.param( ptype, model, PROVIDER_KEYS.get(ptype, ""), - id=ptype, - marks=pytest.mark.skipif( + id = ptype, + marks = pytest.mark.skipif( not PROVIDER_KEYS.get(ptype, ""), - reason=f"no key for {ptype}", + reason = f"no key for {ptype}", ), ) for ptype, (_, model) in _PROVIDER_CONFIGS.items() @@ -504,13 +534,21 @@ class TestVisionInference: """Image URL + text message → non-empty streamed reply.""" encrypted = encrypt_key(api_key) payload = { - "messages": [{ - "role": "user", - "content": [ - {"type": "text", "text": "Which animal is in this image? Reply in one word."}, - {"type": "image_url", "image_url": {"url": vision_image_data_url}}, - ], - }], + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Which animal is in this image? Reply in one word.", + }, + { + "type": "image_url", + "image_url": {"url": vision_image_data_url}, + }, + ], + } + ], "stream": True, "max_tokens": 215, "provider_type": provider_type, @@ -519,14 +557,14 @@ class TestVisionInference: } with requests.post( _url("/v1/chat/completions"), - headers={**auth_headers, "Content-Type": "application/json"}, - json=payload, - stream=True, - timeout=60, + headers = {**auth_headers, "Content-Type": "application/json"}, + json = payload, + stream = True, + timeout = 60, ) as resp: - assert resp.status_code == 200, ( - f"Vision request failed ({resp.status_code}): {resp.text[:300]}" - ) + assert ( + resp.status_code == 200 + ), f"Vision request failed ({resp.status_code}): {resp.text[:300]}" reply, saw_done = _parse_sse_stream(resp) assert reply.strip(), f"Empty reply from {provider_type}/{model}" @@ -548,17 +586,21 @@ class TestLocalInferenceUnaffected: """ resp = requests.post( _url("/v1/chat/completions"), - headers={**auth_headers, "Content-Type": "application/json"}, - json={ + headers = {**auth_headers, "Content-Type": "application/json"}, + json = { "messages": [{"role": "user", "content": "Hello"}], "stream": False, }, - timeout=15, + timeout = 15, ) allowed = {200, 400, 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)" + 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")