backends/openai_compat.py

from __future__ import annotations

import base64
import math
from typing import Any

import httpx

from app.config import Settings
from backends.base import InputType


class EmbedDimensionError(ValueError):
    pass


class OpenAICompatError(RuntimeError):
    pass


def _normalize_base(url: str) -> str:
    return url.rstrip("/")


class OpenAICompatClient:
    def __init__(
        self,
        *,
        base_url: str,
        api_key: str,
        timeout_s: float,
        client: httpx.Client | None = None,
    ) -> None:
        self.base_url = _normalize_base(base_url)
        self.api_key = api_key
        self._owns = client is None
        self._client = client or httpx.Client(
            base_url=self.base_url,
            timeout=timeout_s,
            headers={"Authorization": f"Bearer {api_key}"},
        )

    def close(self) -> None:
        if self._owns:
            self._client.close()

    def health(self) -> bool:
        try:
            response = self._client.get("/models")
            return response.status_code < 500
        except httpx.HTTPError:
            return False

    def chat_completions(self, body: dict[str, Any]) -> dict[str, Any]:
        response = self._client.post("/chat/completions", json=body)
        try:
            response.raise_for_status()
        except httpx.HTTPStatusError as exc:
            raise OpenAICompatError(
                f"chat/completions {exc.response.status_code}: {exc.response.text[:500]}"
            ) from exc
        return response.json()

    def embeddings(self, body: dict[str, Any]) -> dict[str, Any]:
        response = self._client.post("/embeddings", json=body)
        if response.status_code == 404:
            # vLLM pooling runner
            response = self._client.post("/pooling", json={**body, "task": "embed"})
        try:
            response.raise_for_status()
        except httpx.HTTPStatusError as exc:
            raise OpenAICompatError(
                f"embeddings {exc.response.status_code}: {exc.response.text[:500]}"
            ) from exc
        return response.json()


class OpenAICompatLLM:
    def __init__(
        self,
        settings: Settings,
        *,
        name: str,
        accepts_images: bool,
        extra_body: dict[str, Any] | None = None,
        client: httpx.Client | None = None,
    ) -> None:
        self.name = name
        self.accepts_images = accepts_images
        self.model = settings.llm_model
        self.max_tokens = settings.llm_max_tokens
        self.extra_body = extra_body or {}
        self._http = OpenAICompatClient(
            base_url=settings.llm_base_url,
            api_key=settings.llm_api_key,
            timeout_s=settings.llm_timeout_s,
            client=client,
        )

    def health(self) -> bool:
        return self._http.health()

    def complete_json(
        self,
        *,
        system: str,
        user: str,
        image_jpeg: bytes | None = None,
    ) -> str:
        if self.accepts_images and image_jpeg:
            b64 = base64.b64encode(image_jpeg).decode("ascii")
            user_content: Any = [
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{b64}"},
                },
                {"type": "text", "text": user},
            ]
        else:
            user_content = user
        body: dict[str, Any] = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user_content},
            ],
            "temperature": 0,
            "max_tokens": self.max_tokens,
            "response_format": {"type": "json_object"},
        }
        body.update(self.extra_body)
        payload = self._http.chat_completions(body)
        try:
            return str(payload["choices"][0]["message"]["content"] or "")
        except (KeyError, IndexError, TypeError) as exc:
            raise OpenAICompatError(f"unexpected chat response: {payload!r}") from exc


def apply_embed_prefix(text: str, input_type: InputType, *, enabled: bool) -> str:
    if not enabled:
        return text
    prefix = "query: " if input_type == "query" else "passage: "
    stripped = text.lstrip()
    if stripped.startswith("query:") or stripped.startswith("passage:"):
        return text
    return prefix + text


def l2_normalize(vec: list[float]) -> list[float]:
    norm = math.sqrt(sum(x * x for x in vec)) or 1.0
    return [x / norm for x in vec]


def parse_embedding_payload(payload: dict[str, Any]) -> list[list[float]]:
    if "data" in payload:
        rows = sorted(payload["data"], key=lambda row: row.get("index", 0))
        return [list(map(float, row["embedding"])) for row in rows]
    if "embeddings" in payload:
        embeddings = payload["embeddings"]
        if isinstance(embeddings, dict) and "float" in embeddings:
            embeddings = embeddings["float"]
        return [list(map(float, row)) for row in embeddings]
    raise OpenAICompatError(f"unexpected embed response keys: {list(payload)}")


class OpenAICompatEmbed:
    def __init__(
        self,
        settings: Settings,
        *,
        name: str,
        client: httpx.Client | None = None,
    ) -> None:
        self.name = name
        self.dim = settings.embed_dim
        self.model = settings.embed_model
        self.prefix = settings.embed_prefix
        self._http = OpenAICompatClient(
            base_url=settings.embed_base_url or settings.llm_base_url,
            api_key=settings.embed_api_key or settings.llm_api_key,
            timeout_s=settings.embed_timeout_s,
            client=client,
        )

    def health(self) -> bool:
        return self._http.health()

    def embed(self, texts: list[str], *, input_type: InputType) -> list[list[float]]:
        if not texts:
            return []
        prefixed = [
            apply_embed_prefix(text, input_type, enabled=self.prefix) for text in texts
        ]
        body = {
            "model": self.model,
            "input": prefixed,
            "encoding_format": "float",
            "input_type": input_type,
        }
        payload = self._http.embeddings(body)
        vectors = [l2_normalize(vec) for vec in parse_embedding_payload(payload)]
        for vec in vectors:
            if len(vec) != self.dim:
                raise EmbedDimensionError(
                    f"embed dim {len(vec)} != configured {self.dim}. "
                    "Never mix Gemma-3840 and Nemotron-2048 in one index."
                )
        return vectors