KoalaTrapLord
af2a71d1a5
* feat: add async DB wrapper and eliminate event-loop blocking * fix: rewrite llm_client.py with native async httpx - removes urlopen and run_in_executor * fix: rewrite routes/coach.py - replace all sync DB/HTTP with async (aiosqlite + httpx) * fix: add async DB shutdown hook to lab_api.py * feat: add aiosqlite>=0.19.0 to requirements
287 lines
10 KiB
Python
287 lines
10 KiB
Python
"""
|
|
Commander AI Lab — Ollama LLM Client
|
|
════════════════════════════════════
|
|
Fully async HTTP client targeting Ollama's OpenAI-compatible
|
|
/v1/chat/completions endpoint. Every network call uses
|
|
httpx.AsyncClient so it NEVER blocks the FastAPI event loop.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
|
|
from .config import (
|
|
LLM_URL,
|
|
LLM_MODEL,
|
|
LLM_TIMEOUT,
|
|
LLM_MAX_RETRIES,
|
|
DEFAULT_TEMPERATURE,
|
|
DEFAULT_MAX_TOKENS,
|
|
)
|
|
|
|
logger = logging.getLogger("coach.llm")
|
|
|
|
|
|
class LLMResponse:
|
|
"""Parsed LLM response container."""
|
|
|
|
def __init__(
|
|
self,
|
|
content: str,
|
|
parsed_json: Optional[dict],
|
|
prompt_tokens: int,
|
|
completion_tokens: int,
|
|
model: str,
|
|
):
|
|
self.content = content
|
|
self.parsed_json = parsed_json
|
|
self.prompt_tokens = prompt_tokens
|
|
self.completion_tokens = completion_tokens
|
|
self.model = model
|
|
|
|
|
|
class LLMClient:
|
|
"""
|
|
Fully async client for Ollama's OpenAI-compatible API.
|
|
|
|
Usage:
|
|
client = LLMClient()
|
|
response = await client.chat(system_prompt, user_prompt)
|
|
if response.parsed_json:
|
|
suggestions = response.parsed_json
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
base_url: str = None,
|
|
model: str = None,
|
|
timeout: int = None,
|
|
):
|
|
self.base_url = (base_url or LLM_URL).rstrip("/")
|
|
self.model = model or LLM_MODEL
|
|
self.timeout = timeout or LLM_TIMEOUT
|
|
self._resolved_model: Optional[str] = None
|
|
|
|
# ── Model resolution ───────────────────────────────────────
|
|
|
|
async def _resolve_model(self) -> str:
|
|
"""Get the actual loaded model name from Ollama /models endpoint."""
|
|
if self._resolved_model:
|
|
return self._resolved_model
|
|
try:
|
|
status = await self.check_connection()
|
|
if status.get("connected") and status.get("active_model"):
|
|
self._resolved_model = status["active_model"]
|
|
logger.info("Resolved Ollama model: %s", self._resolved_model)
|
|
return self._resolved_model
|
|
except Exception:
|
|
pass
|
|
return self.model
|
|
|
|
# ── Request building ──────────────────────────────────────
|
|
|
|
async def _build_request_body(
|
|
self,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
temperature: float,
|
|
max_tokens: int,
|
|
force_json: bool,
|
|
) -> dict:
|
|
model_name = await self._resolve_model()
|
|
body: dict = {
|
|
"model": model_name,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt},
|
|
],
|
|
"temperature": temperature,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
if force_json:
|
|
body["response_format"] = {"type": "json_object"}
|
|
return body
|
|
|
|
# ── Response parsing ──────────────────────────────────────
|
|
|
|
@staticmethod
|
|
def _strip_think_tags(text: str) -> str:
|
|
"""Remove DeepSeek-R1 <think>...</think> reasoning blocks."""
|
|
import re
|
|
return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL).strip()
|
|
|
|
@staticmethod
|
|
def _extract_json_object(text: str) -> Optional[dict]:
|
|
"""Find and parse the first complete JSON object via brace-counting."""
|
|
start = text.find('{')
|
|
if start == -1:
|
|
return None
|
|
depth = 0
|
|
in_string = False
|
|
escape_next = False
|
|
for i in range(start, len(text)):
|
|
ch = text[i]
|
|
if escape_next:
|
|
escape_next = False
|
|
continue
|
|
if ch == '\\':
|
|
if in_string:
|
|
escape_next = True
|
|
continue
|
|
if ch == '"':
|
|
in_string = not in_string
|
|
continue
|
|
if in_string:
|
|
continue
|
|
if ch == '{':
|
|
depth += 1
|
|
elif ch == '}':
|
|
depth -= 1
|
|
if depth == 0:
|
|
candidate = text[start : i + 1]
|
|
try:
|
|
return json.loads(candidate)
|
|
except (json.JSONDecodeError, ValueError):
|
|
next_start = text.find('{', start + 1)
|
|
if next_start != -1 and next_start < i:
|
|
start = next_start
|
|
depth = 1
|
|
return None
|
|
return None
|
|
|
|
def _parse_response(self, raw: dict) -> LLMResponse:
|
|
"""Extract content and usage from an OpenAI-format response."""
|
|
choice = raw.get("choices", [{}])[0]
|
|
content = choice.get("message", {}).get("content", "")
|
|
usage = raw.get("usage", {})
|
|
cleaned = self._strip_think_tags(content)
|
|
parsed: Optional[dict] = None
|
|
# 1. Direct JSON parse
|
|
try:
|
|
parsed = json.loads(cleaned)
|
|
except (json.JSONDecodeError, TypeError):
|
|
pass
|
|
# 2. Markdown fences
|
|
if parsed is None and "```json" in cleaned:
|
|
try:
|
|
s = cleaned.index("```json") + 7
|
|
e = cleaned.index("```", s)
|
|
parsed = json.loads(cleaned[s:e].strip())
|
|
except (json.JSONDecodeError, ValueError, IndexError):
|
|
pass
|
|
if parsed is None and "```" in cleaned:
|
|
try:
|
|
s = cleaned.index("```") + 3
|
|
e = cleaned.index("```", s)
|
|
parsed = json.loads(cleaned[s:e].strip())
|
|
except (json.JSONDecodeError, ValueError, IndexError):
|
|
pass
|
|
# 3. Brace-counting extraction
|
|
if parsed is None:
|
|
parsed = self._extract_json_object(cleaned)
|
|
# 4. Fallback on original content
|
|
if parsed is None and content != cleaned:
|
|
parsed = self._extract_json_object(content)
|
|
if parsed is not None:
|
|
logger.info("Successfully parsed JSON from LLM response")
|
|
else:
|
|
logger.warning(
|
|
"Could not extract JSON from LLM response (%d chars)", len(content)
|
|
)
|
|
return LLMResponse(
|
|
content=content,
|
|
parsed_json=parsed,
|
|
prompt_tokens=usage.get("prompt_tokens", 0),
|
|
completion_tokens=usage.get("completion_tokens", 0),
|
|
model=raw.get("model", self.model),
|
|
)
|
|
|
|
# ── Public API ─────────────────────────────────────────────
|
|
|
|
async def chat(
|
|
self,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
temperature: float = None,
|
|
max_tokens: int = None,
|
|
force_json: bool = True,
|
|
) -> LLMResponse:
|
|
"""
|
|
Async chat call to Ollama. Retries up to LLM_MAX_RETRIES.
|
|
Never blocks the event loop — uses httpx.AsyncClient throughout.
|
|
"""
|
|
temp = temperature if temperature is not None else DEFAULT_TEMPERATURE
|
|
tokens = max_tokens if max_tokens is not None else DEFAULT_MAX_TOKENS
|
|
body = await self._build_request_body(
|
|
system_prompt, user_prompt, temp, tokens, force_json
|
|
)
|
|
last_error: Optional[Exception] = None
|
|
for attempt in range(1, LLM_MAX_RETRIES + 1):
|
|
try:
|
|
logger.info("LLM call attempt %d/%d", attempt, LLM_MAX_RETRIES)
|
|
async with httpx.AsyncClient(
|
|
timeout=httpx.Timeout(float(self.timeout))
|
|
) as client:
|
|
resp = await client.post(
|
|
f"{self.base_url}/chat/completions",
|
|
json=body,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
resp.raise_for_status()
|
|
raw = resp.json()
|
|
result = self._parse_response(raw)
|
|
if force_json and result.parsed_json is None:
|
|
logger.warning(
|
|
"Attempt %d: LLM returned non-JSON, retrying", attempt
|
|
)
|
|
last_error = ValueError("Non-JSON response from LLM")
|
|
continue
|
|
return result
|
|
except httpx.HTTPStatusError as e:
|
|
logger.warning("Attempt %d HTTP error %s", attempt, e.response.status_code)
|
|
last_error = e
|
|
# 400 with json format → retry without response_format constraint
|
|
if e.response.status_code == 400 and force_json:
|
|
logger.warning("Got 400 — retrying without response_format")
|
|
force_json = False
|
|
body = await self._build_request_body(
|
|
system_prompt, user_prompt, temp, tokens, False
|
|
)
|
|
except Exception as e:
|
|
logger.warning("Attempt %d failed: %s", attempt, e)
|
|
last_error = e
|
|
if attempt < LLM_MAX_RETRIES:
|
|
import asyncio
|
|
await asyncio.sleep(2 * attempt)
|
|
raise ConnectionError(
|
|
f"LLM call failed after {LLM_MAX_RETRIES} attempts: {last_error}"
|
|
)
|
|
|
|
async def check_connection(self) -> dict:
|
|
"""Check if Ollama is reachable. Returns status dict."""
|
|
try:
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(5.0)) as client:
|
|
resp = await client.get(
|
|
f"{self.base_url}/models",
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
models = data.get("data", [])
|
|
model_ids = [m.get("id", "unknown") for m in models]
|
|
return {
|
|
"connected": True,
|
|
"models": model_ids,
|
|
"active_model": model_ids[0] if model_ids else self.model,
|
|
}
|
|
except Exception as e:
|
|
return {
|
|
"connected": False,
|
|
"error": str(e),
|
|
"models": [],
|
|
"active_model": None,
|
|
}
|