| 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: |
| |
| 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 |
|
|
