lm-deluge 0.0.88__py3-none-any.whl → 0.0.90__py3-none-any.whl

lm_deluge/server/app.py

@@ -0,0 +1,388 @@
+"""
+FastAPI application for the LM-Deluge proxy server.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import traceback
+from contextlib import asynccontextmanager
+
+import aiohttp
+from dotenv import load_dotenv
+from fastapi import Depends, FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse
+
+from lm_deluge.models import APIModel, registry
+from lm_deluge.prompt import CachePattern
+from lm_deluge.request_context import RequestContext
+from lm_deluge.tracker import StatusTracker
+
+from .adapters import (
+    anthropic_request_to_conversation,
+    anthropic_request_to_sampling_params,
+    anthropic_tools_to_lm_deluge,
+    api_response_to_anthropic,
+    api_response_to_openai,
+    openai_request_to_conversation,
+    openai_request_to_sampling_params,
+    openai_tools_to_lm_deluge,
+)
+from .auth import verify_anthropic_auth, verify_openai_auth
+from .model_policy import ModelRouter, ProxyModelPolicy
+from .models_anthropic import (
+    AnthropicErrorDetail,
+    AnthropicErrorResponse,
+    AnthropicMessagesRequest,
+    AnthropicMessagesResponse,
+)
+from .models_openai import (
+    OpenAIChatCompletionsRequest,
+    OpenAIChatCompletionsResponse,
+    OpenAIErrorDetail,
+    OpenAIErrorResponse,
+    OpenAIModelInfo,
+    OpenAIModelsResponse,
+)
+
+# Valid cache patterns
+_VALID_CACHE_PATTERNS = {
+    "tools_only",
+    "system_and_tools",
+    "last_user_message",
+    "last_2_user_messages",
+    "last_3_user_messages",
+}
+
+_TRUTHY_VALUES = {"1", "true", "yes", "on"}
+
+
+def get_cache_pattern() -> CachePattern | None:
+    """
+    Get cache pattern from DELUGE_CACHE_PATTERN environment variable.
+
+    Valid values:
+        - none / NONE / unset → no caching
+        - tools_only → cache tools definition
+        - system_and_tools → cache system prompt and tools
+        - last_user_message → cache last user message
+        - last_2_user_messages → cache last 2 user messages
+        - last_3_user_messages → cache last 3 user messages
+    """
+    pattern = os.getenv("DELUGE_CACHE_PATTERN", "").lower().strip()
+    if not pattern or pattern == "none":
+        return None
+    if pattern in _VALID_CACHE_PATTERNS:
+        return pattern  # type: ignore
+    return None
+
+
+def _is_truthy_env(name: str) -> bool:
+    return os.getenv(name, "").strip().lower() in _TRUTHY_VALUES
+
+
+# Global aiohttp session for connection reuse
+_http_session: aiohttp.ClientSession | None = None
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Manage application lifespan - startup and shutdown."""
+    global _http_session
+
+    # Load .env file if present
+    load_dotenv()
+
+    # Create shared aiohttp session
+    connector = aiohttp.TCPConnector(
+        limit=100,
+        limit_per_host=20,
+        keepalive_timeout=30,
+        enable_cleanup_closed=True,
+    )
+    _http_session = aiohttp.ClientSession(connector=connector)
+
+    yield
+
+    # Cleanup
+    if _http_session:
+        await _http_session.close()
+        _http_session = None
+
+
+def _is_model_available(api_model: APIModel) -> bool:
+    """Check if model is available based on configured API keys."""
+    env_var = api_model.api_key_env_var
+    if not env_var:
+        if api_model.api_spec == "bedrock":
+            return bool(os.getenv("AWS_ACCESS_KEY_ID") or os.getenv("AWS_PROFILE"))
+        return False
+    return bool(os.getenv(env_var))
+
+
+def create_app(policy: ProxyModelPolicy | None = None) -> FastAPI:
+    """Create and configure the FastAPI application."""
+    policy = policy or ProxyModelPolicy()
+    router = ModelRouter(policy, registry)
+
+    app = FastAPI(
+        title="LM-Deluge Proxy Server",
+        description="OpenAI and Anthropic compatible API proxy backed by lm-deluge",
+        version="0.1.0",
+        lifespan=lifespan,
+    )
+
+    @app.middleware("http")
+    async def log_requests(request: Request, call_next):
+        if _is_truthy_env("DELUGE_PROXY_LOG_REQUESTS"):
+            body = await request.body()
+            body_text = body.decode("utf-8", errors="replace")
+            if body_text:
+                try:
+                    body_text = json.dumps(json.loads(body_text), indent=2)
+                except Exception:
+                    pass
+            print("DELUGE_PROXY_REQUEST")
+            print(f"{request.method} {request.url}")
+            print("Headers:")
+            print(dict(request.headers))
+            if body_text:
+                print("Body:")
+                print(body_text)
+        return await call_next(request)
+
+    # ========================================================================
+    # Health Check
+    # ========================================================================
+
+    @app.get("/health")
+    async def health_check():
+        """Health check endpoint."""
+        return {"status": "ok"}
+
+    # ========================================================================
+    # OpenAI-Compatible Endpoints
+    # ========================================================================
+
+    @app.get("/v1/models", dependencies=[Depends(verify_openai_auth)])
+    async def list_models(all: bool = False) -> OpenAIModelsResponse:
+        """
+        List available models (OpenAI-compatible).
+
+        By default, only returns models for which the required API key
+        is set in the environment. Use ?all=true to list all registered models.
+        """
+        models = []
+        model_ids = router.list_model_ids(
+            only_available=not all,
+            is_available=lambda model_id: _is_model_available(registry[model_id]),
+        )
+        for model_id in model_ids:
+            models.append(OpenAIModelInfo(id=model_id, owned_by="lm-deluge"))
+
+        return OpenAIModelsResponse(data=models)
+
+    @app.post(
+        "/v1/chat/completions",
+        dependencies=[Depends(verify_openai_auth)],
+        response_model=None,
+    )
+    async def openai_chat_completions(
+        req: OpenAIChatCompletionsRequest,
+    ) -> OpenAIChatCompletionsResponse | JSONResponse:
+        """OpenAI-compatible chat completions endpoint."""
+        # Reject streaming
+        if req.stream:
+            raise HTTPException(
+                status_code=400,
+                detail="Streaming is not supported. Set stream=false.",
+            )
+
+        try:
+            # Get model from registry
+            try:
+                resolved_model = router.resolve(req.model)
+                api_model = APIModel.from_registry(resolved_model)
+            except ValueError as e:
+                raise HTTPException(status_code=400, detail=str(e))
+
+            # Convert request to lm-deluge types
+            conversation = openai_request_to_conversation(req)
+            sampling_params = openai_request_to_sampling_params(req)
+
+            # Convert tools if provided
+            tools = None
+            if req.tools:
+                tools = openai_tools_to_lm_deluge(req.tools)
+
+            # Apply cache pattern only for Anthropic-compatible models
+            cache = None
+            if api_model.api_spec in ("anthropic", "bedrock"):
+                cache = get_cache_pattern()
+
+            # Build RequestContext
+            # We need a minimal StatusTracker for execute_once to work
+            tracker = StatusTracker(
+                max_requests_per_minute=1000,
+                max_tokens_per_minute=1_000_000,
+                max_concurrent_requests=100,
+                use_progress_bar=False,
+            )
+
+            context = RequestContext(
+                task_id=0,
+                model_name=resolved_model,
+                prompt=conversation,
+                sampling_params=sampling_params,
+                tools=tools,
+                cache=cache,
+                status_tracker=tracker,
+                request_timeout=int(os.getenv("DELUGE_PROXY_TIMEOUT", "120")),
+            )
+
+            # Create and execute request
+            request_obj = api_model.make_request(context)
+            response = await request_obj.execute_once()
+
+            # Check for errors
+            if response.is_error:
+                return JSONResponse(
+                    status_code=response.status_code or 500,
+                    content=OpenAIErrorResponse(
+                        error=OpenAIErrorDetail(
+                            message=response.error_message or "Unknown error",
+                            type="api_error",
+                            code=str(response.status_code)
+                            if response.status_code
+                            else None,
+                        )
+                    ).model_dump(),
+                )
+
+            # Convert to OpenAI format
+            return api_response_to_openai(response, resolved_model)
+
+        except HTTPException:
+            raise
+        except Exception as e:
+            traceback.print_exc()
+            return JSONResponse(
+                status_code=500,
+                content=OpenAIErrorResponse(
+                    error=OpenAIErrorDetail(
+                        message=str(e),
+                        type="internal_error",
+                    )
+                ).model_dump(),
+            )
+
+    # ========================================================================
+    # Anthropic-Compatible Endpoints
+    # ========================================================================
+
+    # Support both /v1/messages and /messages for Anthropic SDK compatibility
+    # The Anthropic SDK constructs paths as {base_url}/v1/messages
+    @app.post(
+        "/v1/messages",
+        dependencies=[Depends(verify_anthropic_auth)],
+        response_model=None,
+    )
+    @app.post(
+        "/messages",
+        dependencies=[Depends(verify_anthropic_auth)],
+        response_model=None,
+    )
+    async def anthropic_messages(
+        request: Request,
+        req: AnthropicMessagesRequest,
+    ) -> AnthropicMessagesResponse | JSONResponse:
+        """Anthropic-compatible messages endpoint."""
+        # Reject streaming
+        if req.stream:
+            raise HTTPException(
+                status_code=400,
+                detail="Streaming is not supported. Set stream=false.",
+            )
+
+        try:
+            # Get model from registry
+            try:
+                resolved_model = router.resolve(req.model)
+                api_model = APIModel.from_registry(resolved_model)
+            except ValueError as e:
+                raise HTTPException(status_code=400, detail=str(e))
+
+            # Convert request to lm-deluge types
+            conversation = anthropic_request_to_conversation(req)
+            sampling_params = anthropic_request_to_sampling_params(req)
+
+            # Convert tools if provided
+            tools = None
+            if req.tools:
+                tools = anthropic_tools_to_lm_deluge(req.tools)
+
+            # Apply cache pattern only for Anthropic-compatible models
+            cache = None
+            if api_model.api_spec in ("anthropic", "bedrock"):
+                cache = get_cache_pattern()
+
+            # Build RequestContext
+            tracker = StatusTracker(
+                max_requests_per_minute=1000,
+                max_tokens_per_minute=1_000_000,
+                max_concurrent_requests=100,
+                use_progress_bar=False,
+            )
+
+            extra_headers = None
+            beta_header = request.headers.get("anthropic-beta")
+            if beta_header:
+                extra_headers = {"anthropic-beta": beta_header}
+
+            context = RequestContext(
+                task_id=0,
+                model_name=resolved_model,
+                prompt=conversation,
+                sampling_params=sampling_params,
+                tools=tools,
+                cache=cache,
+                status_tracker=tracker,
+                request_timeout=int(os.getenv("DELUGE_PROXY_TIMEOUT", "120")),
+                extra_headers=extra_headers,
+            )
+
+            # Create and execute request
+            request_obj = api_model.make_request(context)
+            response = await request_obj.execute_once()
+
+            # Check for errors
+            if response.is_error:
+                return JSONResponse(
+                    status_code=response.status_code or 500,
+                    content=AnthropicErrorResponse(
+                        error=AnthropicErrorDetail(
+                            type="api_error",
+                            message=response.error_message or "Unknown error",
+                        )
+                    ).model_dump(),
+                )
+
+            # Convert to Anthropic format
+            return api_response_to_anthropic(response, resolved_model)
+
+        except HTTPException:
+            raise
+        except Exception as e:
+            traceback.print_exc()
+            return JSONResponse(
+                status_code=500,
+                content=AnthropicErrorResponse(
+                    error=AnthropicErrorDetail(
+                        type="internal_error",
+                        message=str(e),
+                    )
+                ).model_dump(),
+            )
+
+    return app

lm_deluge/server/auth.py

@@ -0,0 +1,71 @@
+"""
+Optional authentication for the proxy server.
+"""
+
+from __future__ import annotations
+
+import os
+
+from fastapi import Header, HTTPException
+
+
+def get_proxy_api_key() -> str | None:
+    """Get the configured proxy API key from environment."""
+    return os.getenv("DELUGE_PROXY_API_KEY")
+
+
+async def verify_openai_auth(
+    authorization: str | None = Header(default=None),
+) -> None:
+    """
+    Verify OpenAI-style Bearer token authentication.
+    Only enforced if DELUGE_PROXY_API_KEY is set.
+    """
+    expected_key = get_proxy_api_key()
+    if not expected_key:
+        # No auth configured, allow all requests
+        return
+
+    if not authorization:
+        raise HTTPException(
+            status_code=401,
+            detail="Missing Authorization header",
+        )
+
+    if not authorization.startswith("Bearer "):
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid Authorization header format. Expected 'Bearer <token>'",
+        )
+
+    token = authorization.removeprefix("Bearer ").strip()
+    if token != expected_key:
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid API key",
+        )
+
+
+async def verify_anthropic_auth(
+    x_api_key: str | None = Header(default=None, alias="x-api-key"),
+) -> None:
+    """
+    Verify Anthropic-style x-api-key header authentication.
+    Only enforced if DELUGE_PROXY_API_KEY is set.
+    """
+    expected_key = get_proxy_api_key()
+    if not expected_key:
+        # No auth configured, allow all requests
+        return
+
+    if not x_api_key:
+        raise HTTPException(
+            status_code=401,
+            detail="Missing x-api-key header",
+        )
+
+    if x_api_key != expected_key:
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid API key",
+        )

lm_deluge/server/model_policy.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import random
+from typing import Callable, Literal
+
+import yaml
+from pydantic import BaseModel, Field, model_validator
+
+from lm_deluge.models import APIModel, registry
+
+RouteStrategy = Literal["round_robin", "random", "weighted"]
+PolicyMode = Literal["allow_user_pick", "force_default", "alias_only"]
+
+
+def _dedupe_keep_order(values: list[str]) -> list[str]:
+    seen: set[str] = set()
+    output = []
+    for value in values:
+        if value in seen:
+            continue
+        seen.add(value)
+        output.append(value)
+    return output
+
+
+class RouteConfig(BaseModel):
+    models: list[str]
+    strategy: RouteStrategy = "round_robin"
+    weights: list[float] | None = None
+
+    @model_validator(mode="after")
+    def _validate_route(self) -> "RouteConfig":
+        if not self.models:
+            raise ValueError("route models must not be empty")
+        if self.strategy == "weighted":
+            if not self.weights:
+                raise ValueError("weighted strategy requires weights")
+            if len(self.weights) != len(self.models):
+                raise ValueError("weights must match models length")
+            if any(weight < 0 for weight in self.weights):
+                raise ValueError("weights must be non-negative")
+            if sum(self.weights) <= 0:
+                raise ValueError("weights must sum to a positive value")
+        return self
+
+
+class ProxyModelPolicy(BaseModel):
+    mode: PolicyMode = "allow_user_pick"
+    allowed_models: list[str] | None = None
+    default_model: str | None = None
+    routes: dict[str, RouteConfig] = Field(default_factory=dict)
+    expose_aliases: bool = False
+
+    def validate_against_registry(self, model_registry: dict[str, APIModel]) -> None:
+        registry_keys = set(model_registry.keys())
+        allowed_models = (
+            _dedupe_keep_order(self.allowed_models)
+            if self.allowed_models is not None
+            else None
+        )
+
+        if allowed_models is not None:
+            unknown = [model for model in allowed_models if model not in registry_keys]
+            if unknown:
+                raise ValueError(f"Unknown allowed models: {', '.join(unknown)}")
+
+        for alias, route in self.routes.items():
+            if alias in registry_keys:
+                raise ValueError(
+                    f"Route alias '{alias}' conflicts with a registry model id"
+                )
+            for model_id in route.models:
+                if model_id not in registry_keys:
+                    raise ValueError(
+                        f"Route '{alias}' references unknown model '{model_id}'"
+                    )
+                if allowed_models is not None and model_id not in allowed_models:
+                    raise ValueError(
+                        f"Route '{alias}' uses model '{model_id}' not in allowlist"
+                    )
+
+        if self.mode == "force_default" and not self.default_model:
+            raise ValueError("force_default mode requires default_model")
+
+        if self.default_model:
+            if self.default_model not in self.routes:
+                if self.default_model not in registry_keys:
+                    raise ValueError(f"Default model '{self.default_model}' is unknown")
+                if (
+                    allowed_models is not None
+                    and self.default_model not in allowed_models
+                ):
+                    raise ValueError(
+                        f"Default model '{self.default_model}' not in allowlist"
+                    )
+
+        if self.mode == "alias_only" and not self.routes:
+            raise ValueError("alias_only mode requires at least one route alias")
+
+    def allowed_raw_models(self, model_registry: dict[str, APIModel]) -> list[str]:
+        if self.allowed_models is None:
+            return list(model_registry.keys())
+        return _dedupe_keep_order(self.allowed_models)
+
+
+class ModelRouter:
+    def __init__(
+        self,
+        policy: ProxyModelPolicy,
+        model_registry: dict[str, APIModel] | None = None,
+        *,
+        rng: random.Random | None = None,
+    ) -> None:
+        self.policy = policy
+        self.model_registry = model_registry or registry
+        self.policy.validate_against_registry(self.model_registry)
+        self._rng = rng or random.Random()
+        self._round_robin_index: dict[str, int] = {}
+
+    def resolve(self, requested_model: str) -> str:
+        target = requested_model
+        if self.policy.mode == "force_default":
+            if not self.policy.default_model:
+                raise ValueError("No default model configured")
+            target = self.policy.default_model
+
+        if target in self.policy.routes:
+            return self._select_from_route(target)
+
+        if self.policy.mode == "alias_only":
+            raise ValueError(f"Model '{requested_model}' is not an exposed alias")
+
+        if target not in self.model_registry:
+            raise ValueError(f"Model '{target}' not found in registry")
+        if (
+            self.policy.allowed_models is not None
+            and target not in self.policy.allowed_models
+        ):
+            raise ValueError(f"Model '{target}' is not allowed by proxy policy")
+        return target
+
+    def list_model_ids(
+        self,
+        *,
+        only_available: bool,
+        is_available: Callable[[str], bool],
+    ) -> list[str]:
+        models: list[str] = []
+
+        if self.policy.mode != "alias_only":
+            raw_models = self.policy.allowed_raw_models(self.model_registry)
+            if only_available:
+                raw_models = [model for model in raw_models if is_available(model)]
+            models.extend(raw_models)
+
+        if self.policy.mode == "alias_only" or self.policy.expose_aliases:
+            aliases = list(self.policy.routes.keys())
+            if only_available:
+                aliases = [
+                    alias
+                    for alias in aliases
+                    if any(
+                        is_available(model)
+                        for model in self.policy.routes[alias].models
+                    )
+                ]
+            models.extend(aliases)
+
+        return models
+
+    def _select_from_route(self, alias: str) -> str:
+        route = self.policy.routes[alias]
+        models = route.models
+        if len(models) == 1:
+            return models[0]
+
+        if route.strategy == "round_robin":
+            index = self._round_robin_index.get(alias, 0)
+            selected = models[index % len(models)]
+            self._round_robin_index[alias] = index + 1
+            return selected
+
+        if route.strategy == "weighted":
+            assert route.weights is not None
+            return self._rng.choices(models, weights=route.weights, k=1)[0]
+
+        return self._rng.choice(models)
+
+
+def load_policy_data(path: str | None) -> dict:
+    if not path:
+        return {}
+    with open(path, "r", encoding="utf-8") as handle:
+        data = yaml.safe_load(handle) or {}
+    if not isinstance(data, dict):
+        return {}
+    if "model_policy" in data:
+        return data["model_policy"] or {}
+    proxy_block = data.get("proxy")
+    if isinstance(proxy_block, dict) and "model_policy" in proxy_block:
+        return proxy_block["model_policy"] or {}
+    return data
+
+
+def build_policy(
+    *,
+    path: str | None = None,
+    overrides: dict | None = None,
+) -> ProxyModelPolicy:
+    data = load_policy_data(path)
+    if overrides:
+        data.update(overrides)
+    policy = ProxyModelPolicy(**data)
+    policy.validate_against_registry(registry)
+    return policy