backends.py

"""Backend abstraction for open-source LLM runtimes.

Every supported backend exposes an OpenAI-compatible HTTP surface. The
abstraction here normalises base URLs, auth headers, feature flags, and
model-listing behaviour so the gateway can target any of them uniformly.

Supported BACKEND_KIND values:

- ``vllm``      vLLM OpenAI server (default recommendation)
- ``ollama``    Ollama (expose /v1)
- ``llamacpp``  llama.cpp ``llama-server`` with --api
- ``tgi``       Hugging Face Text Generation Inference (OpenAI-compatible mode)
- ``sglang``    SGLang server in OpenAI-compatible mode
- ``localai``   LocalAI
- ``lmstudio``  LM Studio local server
- ``openai``    Generic OpenAI-compatible endpoint (any other project)

The gateway never assumes a single backend owns the feature set. Anything not
natively supported (e.g. embeddings on a runtime that lacks them) returns a
structured 501 so callers get a clean error instead of mystery failures.
"""

from __future__ import annotations

from dataclasses import dataclass, field

from .config import Settings


@dataclass(frozen=True)
class BackendCapabilities:
    chat: bool = True
    completions: bool = True
    models: bool = True
    embeddings: bool = True
    streaming: bool = True
    tools: bool = True


@dataclass(frozen=True)
class BackendProfile:
    kind: str
    display_name: str
    base_url: str
    health_path: str = "/health"
    models_path: str = "/models"
    chat_path: str = "/chat/completions"
    completions_path: str = "/completions"
    embeddings_path: str = "/embeddings"
    capabilities: BackendCapabilities = field(default_factory=BackendCapabilities)
    extra_headers: dict[str, str] = field(default_factory=dict)

    @property
    def health_url(self) -> str:
        root = self.base_url
        if root.endswith("/v1"):
            root = root[:-3]
        return f"{root.rstrip('/')}{self.health_path}"


_PROFILES: dict[str, BackendCapabilities] = {
    "vllm": BackendCapabilities(),
    "openai": BackendCapabilities(),
    "ollama": BackendCapabilities(embeddings=True),
    "llamacpp": BackendCapabilities(embeddings=True, tools=False),
    "tgi": BackendCapabilities(embeddings=False),
    "sglang": BackendCapabilities(),
    "localai": BackendCapabilities(),
    "lmstudio": BackendCapabilities(),
}

_DISPLAY: dict[str, str] = {
    "vllm": "vLLM",
    "openai": "OpenAI-compatible",
    "ollama": "Ollama",
    "llamacpp": "llama.cpp",
    "tgi": "Text Generation Inference",
    "sglang": "SGLang",
    "localai": "LocalAI",
    "lmstudio": "LM Studio",
}

_HEALTH_PATHS: dict[str, str] = {
    "ollama": "/api/tags",
    "llamacpp": "/health",
    "tgi": "/health",
    "lmstudio": "/v1/models",
}


def resolve_backend(settings: Settings) -> BackendProfile:
    kind = (settings.backend_kind or "openai").lower()
    if kind not in _PROFILES:
        kind = "openai"

    extra_headers: dict[str, str] = {}
    if settings.backend_api_key:
        extra_headers["Authorization"] = f"Bearer {settings.backend_api_key}"

    return BackendProfile(
        kind=kind,
        display_name=_DISPLAY.get(kind, kind),
        base_url=settings.backend_base_url.rstrip("/"),
        health_path=_HEALTH_PATHS.get(kind, "/health"),
        capabilities=_PROFILES[kind],
        extra_headers=extra_headers,
    )