| 2 | schemas, full cost accounting in `llm_jobs`. Deterministic processing always runs before any call here.""" |
2 | schemas, full cost accounting in `llm_jobs`. Deterministic processing always runs before any call here.""" |
| 3 | from __future__ import annotations |
3 | from __future__ import annotations |
| 4 | |
4 | |
| |
5 | +import asyncio |
| 5 | import json |
6 | import json |
| 6 | import logging |
7 | import logging |
| 7 | import re |
8 | import re |
| 67 | return httpx.AsyncClient(base_url=self.base_url, headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}, |
68 | return httpx.AsyncClient(base_url=self.base_url, headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}, |
| 68 | timeout=httpx.Timeout(self.timeout_s, connect=30)) |
69 | timeout=httpx.Timeout(self.timeout_s, connect=30)) |
| 69 | |
70 | |
| |
71 | + async def _post(self, client: httpx.AsyncClient, path: str, body: dict[str, Any]) -> httpx.Response: |
| |
72 | + """Local inference servers answer 503/429 while they load or swap a model: wait and retry instead of failing the job.""" |
| |
73 | + delay = 15.0 |
| |
74 | + for attempt in range(6): |
| |
75 | + r = await client.post(path, json=body) |
| |
76 | + if r.status_code in (503, 429, 502, 504) and attempt < 5: |
| |
77 | + retry_after = r.headers.get("retry-after") |
| |
78 | + wait = min(120.0, float(retry_after)) if retry_after and retry_after.replace(".", "", 1).isdigit() else delay |
| |
79 | + log.info("llm server busy, retrying", extra={"status": r.status_code, "wait": wait, "path": path}) |
| |
80 | + await asyncio.sleep(wait) |
| |
81 | + delay = min(120.0, delay * 1.7) |
| |
82 | + continue |
| |
83 | + return r |
| |
84 | + return r |
| |
85 | + |
| 70 | async def complete(self, *, model: str, system: str, user: str, max_tokens: int, temperature: float, json_mode: bool) -> tuple[str, dict[str, Any]]: |
86 | async def complete(self, *, model: str, system: str, user: str, max_tokens: int, temperature: float, json_mode: bool) -> tuple[str, dict[str, Any]]: |
| 71 | body: dict[str, Any] = {"model": model, "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}], |
87 | body: dict[str, Any] = {"model": model, "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}], |
| 72 | "max_tokens": max_tokens, "temperature": temperature, "stream": False} |
88 | "max_tokens": max_tokens, "temperature": temperature, "stream": False} |
| 73 | if json_mode: |
89 | if json_mode: |
| 74 | body["response_format"] = {"type": "json_object"} |
90 | body["response_format"] = {"type": "json_object"} |
| 75 | async with self._client() as client: |
91 | async with self._client() as client: |
| 76 | − r = await client.post("/chat/completions", json=body) |
92 | + r = await self._post(client, "/chat/completions", body) |
| 77 | if r.status_code == 400 and json_mode: |
93 | if r.status_code == 400 and json_mode: |
| 78 | body.pop("response_format") |
94 | body.pop("response_format") |
| 79 | − r = await client.post("/chat/completions", json=body) |
95 | + r = await self._post(client, "/chat/completions", body) |
| 80 | r.raise_for_status() |
96 | r.raise_for_status() |
| 81 | data = r.json() |
97 | data = r.json() |
| 82 | choice = (data.get("choices") or [{}])[0] |
98 | choice = (data.get("choices") or [{}])[0] |
| 85 | |
101 | |
| 86 | async def embed(self, *, model: str, texts: list[str]) -> list[list[float]]: |
102 | async def embed(self, *, model: str, texts: list[str]) -> list[list[float]]: |
| 87 | async with self._client() as client: |
103 | async with self._client() as client: |
| 88 | − r = await client.post("/embeddings", json={"model": model, "input": texts}) |
104 | + r = await self._post(client, "/embeddings", {"model": model, "input": texts}) |
| 89 | r.raise_for_status() |
105 | r.raise_for_status() |
| 90 | data = r.json() |
106 | data = r.json() |
| 91 | items = sorted(data.get("data") or [], key=lambda d: d.get("index", 0)) |
107 | items = sorted(data.get("data") or [], key=lambda d: d.get("index", 0)) |
| 92 | |
108 | |