"""Database client for the private ML Agent DB Space. This module provides a client interface to the separate database Space, keeping user secrets and sessions isolated from the main application. """ import base64 import json import logging import os from typing import Optional import httpx from cryptography.fernet import Fernet logger = logging.getLogger(__name__) # Database Space configuration DB_SPACE_URL = os.environ.get("DB_SPACE_URL", "https://smolagents-ml-agent-db.hf.space") DB_API_KEY = os.environ.get("DB_API_KEY", "dev-key-change-me") # Encryption key for secrets (encryption happens client-side) _cipher: Optional[Fernet] = None def get_cipher() -> Fernet: """Get Fernet cipher for encrypting secrets.""" global _cipher if _cipher is None: key = os.environ.get("ENCRYPTION_KEY") if not key: logger.warning("ENCRYPTION_KEY not set, using ephemeral key (dev only)") key = Fernet.generate_key().decode() _cipher = Fernet(key.encode() if isinstance(key, str) else key) return _cipher def _get_headers() -> dict: """Get headers for DB API requests.""" return {"X-API-Key": DB_API_KEY} async def _request(method: str, path: str, **kwargs) -> dict: """Make a request to the DB Space.""" url = f"{DB_SPACE_URL}{path}" async with httpx.AsyncClient(timeout=30.0) as client: response = await client.request(method, url, headers=_get_headers(), **kwargs) if response.status_code == 404: return None response.raise_for_status() return response.json() def _request_sync(method: str, path: str, **kwargs) -> dict: """Make a synchronous request to the DB Space.""" url = f"{DB_SPACE_URL}{path}" with httpx.Client(timeout=30.0) as client: response = client.request(method, url, headers=_get_headers(), **kwargs) if response.status_code == 404: return None response.raise_for_status() return response.json() # User operations def upsert_user( hf_user_id: str, username: str, name: str = None, picture: str = None ) -> None: """Create or update a user.""" _request_sync( "PUT", f"/users/{hf_user_id}", json={ "hf_user_id": hf_user_id, "username": username, "name": name, "picture": picture, }, ) def get_user(hf_user_id: str) -> dict | None: """Get user by HF user ID.""" return _request_sync("GET", f"/users/{hf_user_id}") # Secrets operations (encrypted client-side) def store_user_secrets(hf_user_id: str, secrets: dict[str, str]) -> None: """Store encrypted user secrets.""" cipher = get_cipher() encrypted = cipher.encrypt(json.dumps(secrets).encode()) encoded = base64.b64encode(encrypted).decode() _request_sync( "PUT", f"/secrets/{hf_user_id}", json={ "encrypted_secrets": encoded, }, ) def get_user_secrets(hf_user_id: str) -> dict[str, str]: """Get decrypted user secrets.""" result = _request_sync("GET", f"/secrets/{hf_user_id}") if not result or not result.get("encrypted_secrets"): return {} cipher = get_cipher() try: encrypted = base64.b64decode(result["encrypted_secrets"]) decrypted = cipher.decrypt(encrypted) return json.loads(decrypted) except Exception as e: logger.error(f"Failed to decrypt secrets for {hf_user_id}: {e}") return {} def delete_user_secrets(hf_user_id: str) -> bool: """Delete user secrets.""" result = _request_sync("DELETE", f"/secrets/{hf_user_id}") return result.get("deleted", False) if result else False # Session operations def save_session( session_id: str, hf_user_id: str, title: str | None = None, messages: list[dict] | None = None, config_snapshot: dict | None = None, ) -> None: """Save or update a session.""" _request_sync( "PUT", f"/sessions/{session_id}", json={ "session_id": session_id, "hf_user_id": hf_user_id, "title": title, "messages": messages, "config_snapshot": config_snapshot, }, ) def get_session(session_id: str) -> dict | None: """Get a session by ID.""" return _request_sync("GET", f"/sessions/{session_id}") def list_user_sessions(hf_user_id: str, include_archived: bool = False) -> list[dict]: """List all sessions for a user.""" params = {"include_archived": str(include_archived).lower()} result = _request_sync("GET", f"/users/{hf_user_id}/sessions", params=params) return result if result else [] def archive_session(session_id: str) -> bool: """Archive a session.""" result = _request_sync("POST", f"/sessions/{session_id}/archive") return result.get("archived", False) if result else False def delete_session(session_id: str) -> bool: """Delete a session.""" result = _request_sync("DELETE", f"/sessions/{session_id}") return result.get("deleted", False) if result else False # Async versions for use in async contexts async def upsert_user_async( hf_user_id: str, username: str, name: str = None, picture: str = None ) -> None: """Create or update a user (async).""" await _request( "PUT", f"/users/{hf_user_id}", json={ "hf_user_id": hf_user_id, "username": username, "name": name, "picture": picture, }, ) async def save_session_async( session_id: str, hf_user_id: str, title: str | None = None, messages: list[dict] | None = None, config_snapshot: dict | None = None, ) -> None: """Save or update a session (async).""" await _request( "PUT", f"/sessions/{session_id}", json={ "session_id": session_id, "hf_user_id": hf_user_id, "title": title, "messages": messages, "config_snapshot": config_snapshot, }, )