SPB Git forge

spb/uqo-chat

Public
14commits 1branches 0releases
1.4 MBsize
maindefault branch
17 days agolast push
Python 64.6% TypeScript 33.7% CSS 0.8%
11.8 KB · 280 lines python
Raw Blame History
1"""LLMClient — single entry point for every model call (OpenRouter, OpenAI-compatible).23Handles: streaming SSE parsing (keep-alive comments, fragmented tool_calls, final usage4chunk), native model fallback via `models`, retries on 429/5xx, attribution headers.5"""67from __future__ import annotations89import asyncio10import json11from collections.abc import AsyncIterator12from typing import Any1314import httpx1516from app.core.config import Settings, get_settings17from app.core.logging import get_logger18from app.llm.schemas import StreamEvent1920log = get_logger("llm.openrouter")2122RETRY_STATUS = {429, 500, 502, 503, 504}232425class LLMError(Exception):26    def __init__(self, message: str, status: int | None = None) -> None:27        super().__init__(message)28        self.status = status293031class LLMClient:32    def __init__(self, settings: Settings | None = None) -> None:33        self.settings = settings or get_settings()34        self._client: httpx.AsyncClient | None = None3536    def _headers(self) -> dict[str, str]:37        return {38            "Authorization": f"Bearer {self.settings.OPENROUTER_API_KEY.get_secret_value()}",39            "HTTP-Referer": self.settings.OPENROUTER_APP_URL,40            "X-Title": self.settings.OPENROUTER_APP_NAME,41            "Content-Type": "application/json",42        }4344    def client(self) -> httpx.AsyncClient:45        if self._client is None or self._client.is_closed:46            self._client = httpx.AsyncClient(47                base_url=self.settings.OPENROUTER_BASE_URL,48                timeout=httpx.Timeout(self.settings.LLM_TIMEOUT_SECONDS, connect=15),49                http2=True,50            )51        return self._client5253    async def aclose(self) -> None:54        if self._client and not self._client.is_closed:55            await self._client.aclose()5657    # ------------------------------------------------------------------ body58    def _body(59        self,60        messages: list[dict[str, Any]],61        tools: list[dict[str, Any]] | None,62        models: list[str],63        temperature: float,64        max_tokens: int,65        reasoning: dict[str, Any] | None,66        provider: dict[str, Any] | None,67        user_id_hash: str | None,68        stream: bool,69        response_format: dict[str, Any] | None,70    ) -> dict[str, Any]:71        body: dict[str, Any] = {72            "messages": messages,73            "stream": stream,74            "temperature": temperature,75            "max_tokens": max_tokens,76            "usage": {"include": True},77            "provider": provider or {"data_collection": "deny", "allow_fallbacks": True},78        }79        if len(models) == 1:80            body["model"] = models[0]81        else:82            body["model"] = models[0]83            body["models"] = models84        if tools:85            body["tools"] = tools86            body["tool_choice"] = "auto"87        if reasoning:88            body["reasoning"] = reasoning89        if user_id_hash:90            body["user"] = user_id_hash91        if response_format:92            body["response_format"] = response_format93        return body9495    # --------------------------------------------------------------- stream96    async def stream_chat(97        self,98        messages: list[dict[str, Any]],99        tools: list[dict[str, Any]] | None,100        models: list[str],101        temperature: float = 0.3,102        max_tokens: int | None = None,103        reasoning: dict[str, Any] | None = None,104        provider: dict[str, Any] | None = None,105        user_id_hash: str | None = None,106        response_format: dict[str, Any] | None = None,107    ) -> AsyncIterator[StreamEvent]:108        body = self._body(109            messages, tools, models, temperature,110            max_tokens or self.settings.LLM_MAX_OUTPUT_TOKENS,111            reasoning, provider, user_id_hash, True, response_format,112        )113        attempt = 0114        while True:115            attempt += 1116            try:117                async with self.client().stream(118                    "POST", "/chat/completions", headers=self._headers(), json=body119                ) as resp:120                    if resp.status_code in RETRY_STATUS and attempt < 3:121                        await resp.aread()122                        log.warning("llm_retry", status=resp.status_code, attempt=attempt)123                        await asyncio.sleep(1.5 * attempt)124                        continue125                    if resp.status_code >= 400:126                        text = (await resp.aread()).decode(errors="replace")127                        raise LLMError(f"OpenRouter {resp.status_code}: {text[:400]}",128                                       resp.status_code)129                    async for ev in self._parse_sse(resp):130                        yield ev131                    return132            except (httpx.TransportError, httpx.TimeoutException) as exc:133                if attempt >= 3:134                    yield StreamEvent("error", {"code": "network", "message": str(exc)})135                    return136                log.warning("llm_network_retry", attempt=attempt, error=str(exc))137                await asyncio.sleep(1.5 * attempt)138139    async def _parse_sse(self, resp: httpx.Response) -> AsyncIterator[StreamEvent]:140        """Parse OpenRouter SSE. Yields normalised StreamEvents."""141        tool_buf: dict[int, dict[str, Any]] = {}142        finish_reason: str | None = None143        model_used = ""144        gen_id = ""145        usage: dict[str, Any] | None = None146        async for raw_line in resp.aiter_lines():147            line = raw_line.strip()148            if not line or line.startswith(":"):  # keep-alive ": OPENROUTER PROCESSING"149                continue150            if not line.startswith("data:"):151                continue152            data = line[5:].strip()153            if data == "[DONE]":154                break155            try:156                chunk = json.loads(data)157            except json.JSONDecodeError:158                continue159            if "error" in chunk and chunk["error"]:160                err = chunk["error"]161                yield StreamEvent("error", {"code": str(err.get("code", "upstream")),162                                            "message": str(err.get("message", err))})163                return164            model_used = chunk.get("model") or model_used165            gen_id = chunk.get("id") or gen_id166            if chunk.get("usage"):167                usage = chunk["usage"]168            for choice in chunk.get("choices", []) or []:169                delta = choice.get("delta") or {}170                if delta.get("content"):171                    yield StreamEvent("text_delta", {"delta": delta["content"]})172                if delta.get("reasoning"):173                    yield StreamEvent("reasoning_delta", {"delta": delta["reasoning"]})174                for tc in delta.get("tool_calls") or []:175                    if "index" in tc and tc["index"] is not None:176                        idx = int(tc["index"])177                    elif tc.get("id") and any(b["id"] == tc["id"] for b in tool_buf.values()):178                        idx = next(k for k, b in tool_buf.items() if b["id"] == tc["id"])179                    elif tc.get("id") or not tool_buf:180                        idx = len(tool_buf)  # provider omitted index: new call181                    else:182                        idx = max(tool_buf)  # continuation of the last call183                    buf = tool_buf.get(idx)184                    if buf is None:185                        buf = {"id": tc.get("id") or f"call_{idx}", "name": "", "args": ""}186                        tool_buf[idx] = buf187                        yield StreamEvent("tool_call_start", {"index": idx, "id": buf["id"]})188                    if tc.get("id"):189                        buf["id"] = tc["id"]190                    fn = tc.get("function") or {}191                    if fn.get("name") and not buf["name"]:192                        buf["name"] = fn["name"]193                    if fn.get("arguments"):194                        buf["args"] += fn["arguments"]195                        yield StreamEvent("tool_call_delta", {"index": idx, "delta": fn["arguments"]})196                if choice.get("finish_reason"):197                    finish_reason = choice["finish_reason"]198        for idx in sorted(tool_buf):199            b = tool_buf[idx]200            yield StreamEvent("tool_call_end", {"index": idx, "id": b["id"], "name": b["name"],201                                                "arguments": b["args"]})202        if usage:203            yield StreamEvent("usage", {204                "input_tokens": usage.get("prompt_tokens", 0),205                "output_tokens": usage.get("completion_tokens", 0),206                "cost_usd": float(usage.get("cost", 0.0) or 0.0),207                "model": model_used,208                "generation_id": gen_id,209            })210        if finish_reason is None and tool_buf:211            finish_reason = "tool_calls"212        yield StreamEvent("done", {"finish_reason": finish_reason or "stop", "model": model_used,213                                   "generation_id": gen_id})214215    # ------------------------------------------------------------- complete216    async def complete(217        self,218        messages: list[dict[str, Any]],219        models: list[str],220        temperature: float = 0.0,221        max_tokens: int = 1024,222        response_format: dict[str, Any] | None = None,223        user_id_hash: str | None = None,224    ) -> tuple[str, dict[str, Any]]:225        """Non-streaming helper for small tasks (titles, classification, quiz JSON)."""226        body = self._body(messages, None, models, temperature, max_tokens, None, None,227                          user_id_hash, False, response_format)228        for attempt in range(1, 4):229            try:230                resp = await self.client().post("/chat/completions", headers=self._headers(),231                                                json=body)232            except (httpx.TransportError, httpx.TimeoutException) as exc:233                if attempt == 3:234                    raise LLMError(f"network: {exc}") from exc235                await asyncio.sleep(1.5 * attempt)236                continue237            if resp.status_code in RETRY_STATUS and attempt < 3:238                await asyncio.sleep(1.5 * attempt)239                continue240            if resp.status_code >= 400:241                raise LLMError(f"OpenRouter {resp.status_code}: {resp.text[:400]}", resp.status_code)242            data = resp.json()243            choice = (data.get("choices") or [{}])[0]244            text = (choice.get("message") or {}).get("content") or ""245            usage = data.get("usage") or {}246            return text, {247                "input_tokens": usage.get("prompt_tokens", 0),248                "output_tokens": usage.get("completion_tokens", 0),249                "cost_usd": float(usage.get("cost", 0.0) or 0.0),250                "model": data.get("model", ""),251                "generation_id": data.get("id", ""),252            }253        raise LLMError("OpenRouter: retries exhausted")254255    async def embed(self, texts: list[str]) -> list[list[float]] | None:256        """Optional embeddings via an OpenAI-compatible endpoint (EMBEDDINGS_BASE_URL)."""257        s = self.settings258        if not s.EMBEDDINGS_BASE_URL or not s.MODEL_EMBEDDINGS:259            return None260        headers = {"Content-Type": "application/json"}261        key = s.EMBEDDINGS_API_KEY.get_secret_value() or s.OPENROUTER_API_KEY.get_secret_value()262        if key:263            headers["Authorization"] = f"Bearer {key}"264        async with httpx.AsyncClient(timeout=60) as c:265            r = await c.post(f"{s.EMBEDDINGS_BASE_URL.rstrip('/')}/embeddings", headers=headers,266                             json={"model": s.MODEL_EMBEDDINGS, "input": texts})267            r.raise_for_status()268            data = r.json()["data"]269            return [d["embedding"] for d in sorted(data, key=lambda d: d["index"])]270271272_client: LLMClient | None = None273274275def get_llm() -> LLMClient:276    global _client277    if _client is None:278        _client = LLMClient()279    return _client280