kairo-code 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

kairo/backend/services/byok_service.py

@@ -0,0 +1,204 @@
+"""BYOK (Bring Your Own Key) service for Max plan users.
+
+Handles encrypted storage and proxied inference using user-provided
+OpenAI/Anthropic API keys.
+"""
+
+import logging
+import os
+import secrets
+
+import httpx
+from cryptography.fernet import Fernet
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from backend.models.user_provider_key import UserProviderKey
+
+logger = logging.getLogger(__name__)
+
+# Derive encryption key from JWT secret or dedicated env var.
+# In production, use a dedicated KMS key.
+_ENCRYPTION_KEY: bytes | None = None
+
+
+def _get_fernet() -> Fernet:
+    global _ENCRYPTION_KEY
+    if _ENCRYPTION_KEY is None:
+        key_str = os.environ.get("KAIRO_BYOK_ENCRYPTION_KEY", "")
+        if not key_str:
+            # Derive from JWT secret (acceptable for single-instance deploy)
+            from backend.config import settings
+            import hashlib
+            import base64
+            raw = hashlib.sha256(settings.JWT_SECRET_KEY.encode()).digest()
+            _ENCRYPTION_KEY = base64.urlsafe_b64encode(raw)
+        else:
+            _ENCRYPTION_KEY = key_str.encode()
+    return Fernet(_ENCRYPTION_KEY)
+
+
+def encrypt_key(plaintext_key: str) -> str:
+    """Encrypt an API key for storage."""
+    f = _get_fernet()
+    return f.encrypt(plaintext_key.encode()).decode()
+
+
+def decrypt_key(encrypted_key: str) -> str:
+    """Decrypt a stored API key."""
+    f = _get_fernet()
+    return f.decrypt(encrypted_key.encode()).decode()
+
+
+# Allowed upstream endpoints per provider
+_PROVIDER_ENDPOINTS = {
+    "openai": {
+        "base_url": "https://api.openai.com",
+        "chat_path": "/v1/chat/completions",
+        "models_path": "/v1/models",
+    },
+    "anthropic": {
+        "base_url": "https://api.anthropic.com",
+        "chat_path": "/v1/messages",
+        "models_path": None,
+    },
+}
+
+
+class BYOKService:
+    def __init__(self, db: AsyncSession):
+        self.db = db
+
+    async def store_key(
+        self, user_id: str, provider: str, api_key: str, label: str | None = None,
+    ) -> UserProviderKey:
+        """Store a user's provider API key (encrypted)."""
+        if provider not in _PROVIDER_ENDPOINTS:
+            raise ValueError(f"Unsupported provider: {provider}")
+
+        # Validate the key works by calling the provider
+        await self._validate_key(provider, api_key)
+
+        encrypted = encrypt_key(api_key)
+        suffix = api_key[-4:]
+
+        record = UserProviderKey(
+            user_id=user_id,
+            provider=provider,
+            encrypted_key=encrypted,
+            key_suffix=suffix,
+            label=label,
+        )
+        self.db.add(record)
+        await self.db.commit()
+        await self.db.refresh(record)
+        logger.info("BYOK key stored: user=%s provider=%s id=%s", user_id, provider, record.id)
+        return record
+
+    async def list_keys(self, user_id: str) -> list[UserProviderKey]:
+        """List user's stored provider keys (without decrypting)."""
+        stmt = (
+            select(UserProviderKey)
+            .where(UserProviderKey.user_id == user_id)
+            .order_by(UserProviderKey.created_at.desc())
+        )
+        result = await self.db.execute(stmt)
+        return list(result.scalars().all())
+
+    async def delete_key(self, user_id: str, key_id: str) -> bool:
+        """Delete a stored provider key."""
+        stmt = select(UserProviderKey).where(
+            UserProviderKey.id == key_id, UserProviderKey.user_id == user_id
+        )
+        result = await self.db.execute(stmt)
+        record = result.scalar_one_or_none()
+        if not record:
+            return False
+        await self.db.delete(record)
+        await self.db.commit()
+        logger.info("BYOK key deleted: id=%s user=%s", key_id, user_id)
+        return True
+
+    async def proxy_chat(
+        self, user_id: str, key_id: str, request_body: dict,
+    ) -> dict:
+        """Proxy a chat completion request using the user's stored key."""
+        stmt = select(UserProviderKey).where(
+            UserProviderKey.id == key_id, UserProviderKey.user_id == user_id
+        )
+        result = await self.db.execute(stmt)
+        record = result.scalar_one_or_none()
+        if not record:
+            raise ValueError("Provider key not found")
+
+        provider = record.provider
+        config = _PROVIDER_ENDPOINTS[provider]
+        api_key = decrypt_key(record.encrypted_key)
+
+        headers = {"Content-Type": "application/json"}
+        if provider == "openai":
+            headers["Authorization"] = f"Bearer {api_key}"
+        elif provider == "anthropic":
+            headers["x-api-key"] = api_key
+            headers["anthropic-version"] = "2023-06-01"
+
+        url = f"{config['base_url']}{config['chat_path']}"
+
+        try:
+            async with httpx.AsyncClient(timeout=120.0) as client:
+                resp = await client.post(url, json=request_body, headers=headers)
+                resp.raise_for_status()
+                return resp.json()
+        except httpx.HTTPStatusError as e:
+            # Sanitize error - never leak the API key
+            status = e.response.status_code
+            if status == 401:
+                detail = "Provider returned 401 Unauthorized. Check your API key."
+            elif status == 429:
+                detail = "Provider rate limit exceeded. Try again later."
+            else:
+                detail = f"Provider returned HTTP {status}."
+            logger.error(
+                "BYOK proxy error: user=%s provider=%s status=%d",
+                user_id, provider, status,
+            )
+            raise ValueError(detail)
+        except httpx.RequestError:
+            logger.error("BYOK proxy connection error: user=%s provider=%s", user_id, provider)
+            raise ValueError("Could not connect to provider. Try again later.")
+
+    async def _validate_key(self, provider: str, api_key: str) -> None:
+        """Validate a provider key by making a lightweight API call."""
+        config = _PROVIDER_ENDPOINTS[provider]
+        headers = {}
+
+        if provider == "openai":
+            headers["Authorization"] = f"Bearer {api_key}"
+            url = f"{config['base_url']}{config['models_path']}"
+        elif provider == "anthropic":
+            headers["x-api-key"] = api_key
+            headers["anthropic-version"] = "2023-06-01"
+            # Anthropic has no /models endpoint; send a minimal message
+            url = f"{config['base_url']}{config['chat_path']}"
+        else:
+            return
+
+        try:
+            async with httpx.AsyncClient(timeout=10.0) as client:
+                if provider == "openai":
+                    resp = await client.get(url, headers=headers)
+                else:
+                    # Minimal Anthropic request to validate the key
+                    resp = await client.post(
+                        url, json={
+                            "model": "claude-3-haiku-20240307",
+                            "max_tokens": 1,
+                            "messages": [{"role": "user", "content": "hi"}],
+                        }, headers=headers,
+                    )
+                if resp.status_code == 401:
+                    raise ValueError("Invalid API key. Provider returned 401.")
+                # Any other status (200, 429, etc.) means the key is valid
+        except httpx.RequestError:
+            raise ValueError("Could not reach provider to validate key. Try again.")