Livrare LOT 1 - Didi
This commit is contained in:
commit
5380c3fc63
990 changed files with 133308 additions and 0 deletions
260
ai_platform/modules/didi_brain/shared/llm_client.py
Normal file
260
ai_platform/modules/didi_brain/shared/llm_client.py
Normal file
|
|
@ -0,0 +1,260 @@
|
|||
"""LLM client — async wrapper over the OpenAI-compat router on :14011.
|
||||
|
||||
Design:
|
||||
- One LlmClient instance per long-running process (reusable httpx.AsyncClient).
|
||||
- Roles (REASONING / FAST / VISION) decide model + backend automatically.
|
||||
- Hard timeouts and retries are enforced; no LLM call hangs forever.
|
||||
- Helpers for the two patterns we use most: structured JSON and single-token NLI.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from tenacity import (
|
||||
AsyncRetrying,
|
||||
retry_if_exception_type,
|
||||
stop_after_attempt,
|
||||
wait_exponential,
|
||||
)
|
||||
|
||||
from shared.config import LlmRole, settings
|
||||
from shared.logging import get_logger
|
||||
|
||||
log = get_logger(__name__)
|
||||
|
||||
|
||||
class LlmError(Exception):
|
||||
"""Wraps any failure from the LLM stack with enough context to debug."""
|
||||
|
||||
def __init__(self, message: str, *, status: int | None = None, body: str | None = None):
|
||||
super().__init__(message)
|
||||
self.status = status
|
||||
self.body = body
|
||||
|
||||
|
||||
class RoleNotAvailable(LlmError):
|
||||
"""The requested role has no enabled model behind it."""
|
||||
|
||||
|
||||
class LlmClient:
|
||||
"""Async OpenAI-compat client targeting the unified router."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str | None = None,
|
||||
timeout: float = 120.0,
|
||||
max_attempts: int = 3,
|
||||
):
|
||||
self._base_url = (base_url or settings.llm_router_url).rstrip("/")
|
||||
self._timeout = timeout
|
||||
self._max_attempts = max_attempts
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if settings.llm_router_api_key:
|
||||
headers["Authorization"] = f"Bearer {settings.llm_router_api_key}"
|
||||
self._http = httpx.AsyncClient(
|
||||
base_url=self._base_url,
|
||||
headers=headers,
|
||||
timeout=httpx.Timeout(timeout, connect=10.0),
|
||||
)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._http.aclose()
|
||||
|
||||
async def __aenter__(self) -> LlmClient:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: Any) -> None:
|
||||
await self.aclose()
|
||||
|
||||
# ------------------------------------------------------------------ core
|
||||
async def chat(
|
||||
self,
|
||||
*,
|
||||
role: LlmRole = LlmRole.REASONING,
|
||||
messages: list[dict[str, str]],
|
||||
temperature: float = 0.1,
|
||||
max_tokens: int = 1024,
|
||||
stop: list[str] | None = None,
|
||||
extra: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Raw chat completion. Returns the parsed response dict.
|
||||
|
||||
The response shape is OpenAI-compatible:
|
||||
{"choices": [{"message": {"content": "..."}}], "usage": {...}, "backend": "..."}
|
||||
"""
|
||||
model_info = settings.model_for(role)
|
||||
if not model_info:
|
||||
raise RoleNotAvailable(f"Role {role.value} has no enabled model")
|
||||
model_id, backend = model_info
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": model_id,
|
||||
"backend": backend,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
}
|
||||
if stop:
|
||||
payload["stop"] = stop
|
||||
if extra:
|
||||
payload.update(extra)
|
||||
|
||||
async for attempt in AsyncRetrying(
|
||||
stop=stop_after_attempt(self._max_attempts),
|
||||
wait=wait_exponential(multiplier=1, min=1, max=10),
|
||||
retry=retry_if_exception_type((httpx.HTTPError, httpx.TimeoutException)),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
resp = await self._http.post("/v1/chat/completions", json=payload)
|
||||
if resp.status_code >= 400:
|
||||
raise LlmError(
|
||||
f"HTTP {resp.status_code} from LLM router",
|
||||
status=resp.status_code,
|
||||
body=resp.text[:1000],
|
||||
)
|
||||
data = resp.json()
|
||||
if "choices" not in data:
|
||||
raise LlmError(
|
||||
"LLM response missing 'choices'",
|
||||
body=json.dumps(data)[:1000],
|
||||
)
|
||||
return data
|
||||
raise LlmError("retry loop exited without result") # unreachable
|
||||
|
||||
# ------------------------------------------------------------------ helpers
|
||||
async def chat_text(
|
||||
self,
|
||||
*,
|
||||
role: LlmRole = LlmRole.REASONING,
|
||||
system: str | None = None,
|
||||
user: str,
|
||||
temperature: float = 0.1,
|
||||
max_tokens: int = 1024,
|
||||
stop: list[str] | None = None,
|
||||
) -> tuple[str, dict[str, Any]]:
|
||||
"""Convenience: send (system?, user) → return (text, usage_dict)."""
|
||||
msgs: list[dict[str, str]] = []
|
||||
if system:
|
||||
msgs.append({"role": "system", "content": system})
|
||||
msgs.append({"role": "user", "content": user})
|
||||
data = await self.chat(
|
||||
role=role,
|
||||
messages=msgs,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
stop=stop,
|
||||
)
|
||||
text = data["choices"][0]["message"]["content"]
|
||||
usage = data.get("usage", {}) | {"backend": data.get("backend", "")}
|
||||
return text, usage
|
||||
|
||||
async def chat_json(
|
||||
self,
|
||||
*,
|
||||
role: LlmRole = LlmRole.REASONING,
|
||||
system: str | None = None,
|
||||
user: str,
|
||||
max_tokens: int = 2048,
|
||||
temperature: float = 0.0,
|
||||
) -> tuple[dict[str, Any] | list[Any], dict[str, Any]]:
|
||||
"""Send a request expected to return JSON. Strips fences if present.
|
||||
|
||||
Raises LlmError if the response is not parseable as JSON.
|
||||
"""
|
||||
text, usage = await self.chat_text(
|
||||
role=role,
|
||||
system=system,
|
||||
user=user,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
cleaned = _extract_json(text)
|
||||
try:
|
||||
return json.loads(cleaned), usage
|
||||
except json.JSONDecodeError as e:
|
||||
raise LlmError(
|
||||
f"LLM did not return valid JSON: {e}",
|
||||
body=text[:1000],
|
||||
) from e
|
||||
|
||||
async def chat_label(
|
||||
self,
|
||||
*,
|
||||
role: LlmRole = LlmRole.REASONING,
|
||||
system: str,
|
||||
user: str,
|
||||
allowed: list[str],
|
||||
max_tokens: int = 16,
|
||||
) -> tuple[str, dict[str, Any]]:
|
||||
"""Single-label classification. Returns the matched label uppercased.
|
||||
|
||||
Useful for NLI (SUPPORT/CONTRADICT/NEUTRAL), credibility tiers, etc.
|
||||
Raises LlmError if no allowed label is found in the response.
|
||||
"""
|
||||
text, usage = await self.chat_text(
|
||||
role=role,
|
||||
system=system,
|
||||
user=user,
|
||||
temperature=0.0,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
upper = text.upper()
|
||||
for label in allowed:
|
||||
if label.upper() in upper:
|
||||
return label.upper(), usage
|
||||
raise LlmError(
|
||||
f"LLM response did not contain any allowed label {allowed}",
|
||||
body=text[:500],
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ health
|
||||
async def list_models(self) -> list[dict[str, Any]]:
|
||||
resp = await self._http.get("/v1/models")
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
return data.get("data", [])
|
||||
|
||||
async def list_backends(self) -> list[str]:
|
||||
try:
|
||||
resp = await self._http.get("/v1/backends")
|
||||
resp.raise_for_status()
|
||||
return resp.json().get("backends", [])
|
||||
except httpx.HTTPError:
|
||||
return []
|
||||
|
||||
|
||||
# --------------------------------------------------------------------- utilities
|
||||
_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.DOTALL)
|
||||
|
||||
|
||||
def _extract_json(text: str) -> str:
|
||||
"""Strip markdown code fences and isolate the JSON object/array if needed."""
|
||||
text = text.strip()
|
||||
m = _FENCE_RE.search(text)
|
||||
if m:
|
||||
return m.group(1).strip()
|
||||
# Find first { or [ and last matching close
|
||||
starts = [text.find("{"), text.find("[")]
|
||||
starts = [s for s in starts if s >= 0]
|
||||
if not starts:
|
||||
return text
|
||||
start = min(starts)
|
||||
return text[start:].strip()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def llm_client() -> AsyncIterator[LlmClient]:
|
||||
"""`async with llm_client() as llm:` for short-lived scripts."""
|
||||
client = LlmClient()
|
||||
try:
|
||||
yield client
|
||||
finally:
|
||||
await client.aclose()
|
||||
Loading…
Add table
Add a link
Reference in a new issue