Coverage for src/asketmc_bot/llm_client.py: 67%
153 statements
« prev ^ index » next coverage.py v7.5.1, created at 2026-04-08 17:14 +0000
1# src/asketmc_bot/llm_client.py
2#!/usr/bin/env python3.10
3"""LLM client for OpenRouter (remote) and Ollama (local) with a circuit breaker.
5- Manages a shared aiohttp session.
6- Retries remote calls with backoff and JSON hardening.
7- Falls back to a local model when the remote path fails or is blocked.
8- Exposes breaker state for fast, synchronous checks from other threads.
9"""
11from __future__ import annotations
13import asyncio
14import json
15import logging
16import time
17from dataclasses import dataclass
18from typing import Any, List, Optional, Tuple
20from aiohttp import ClientSession, TCPConnector, ClientTimeout
23# ── Config DTO ────────────────────────────────────────────────────────────────
24@dataclass(frozen=True)
25class LLMConfig:
26 """Configuration required by the LLM client."""
28 api_url: str
29 or_model: str
30 or_max_tokens: int
31 openrouter_api_key: str
32 ollama_url: str
33 local_model: str
34 http_conn_limit: int
35 or_retries: int
36 http_timeout_total: int
37 breaker_base_block_sec: int
38 breaker_max_block_sec: int
41# ── Async HTTP session holder ─────────────────────────────────────────────────
42class AsyncSessionHolder:
43 """Create and reuse a single aiohttp ClientSession with a connection limit."""
45 def __init__(self, *, limit: int, timeout_total: int) -> None:
46 self._lock = asyncio.Lock()
47 self._session: Optional[ClientSession] = None
48 self._limit = int(limit)
49 self._timeout_total = int(timeout_total)
51 async def get(self) -> ClientSession:
52 """Return an open session, creating it if needed."""
53 async with self._lock:
54 if self._session is None or self._session.closed:
55 self._session = ClientSession(
56 connector=TCPConnector(limit=self._limit),
57 timeout=ClientTimeout(total=self._timeout_total),
58 )
59 return self._session
61 async def close(self) -> None:
62 """Close the session if open."""
63 async with self._lock:
64 if self._session and not self._session.closed: 64 ↛ 63line 64 didn't jump to line 63
65 await self._session.close()
66 self._session = None
69# ── Circuit Breaker ───────────────────────────────────────────────────────────
70class CircuitBreaker:
71 """Closed/Open/Half-open breaker with exponential backoff."""
73 def __init__(self, base_block: int, max_block: int) -> None:
74 self._state = "closed"
75 self._open_until = 0.0
76 self._block = max(1, int(base_block))
77 self._max_block = max(self._block, int(max_block))
78 self._lock = asyncio.Lock()
80 async def allow(self) -> bool:
81 """Return True if a remote call is allowed now."""
82 async with self._lock:
83 now = time.time()
84 if self._state == "open":
85 if now >= self._open_until:
86 self._state = "half_open"
87 return True
88 return False
89 return True
91 async def on_success(self, *, base_block: int) -> None:
92 """Reset breaker after a successful remote call."""
93 async with self._lock:
94 self._state = "closed"
95 self._block = max(1, int(base_block))
97 async def on_failure(self) -> None:
98 """Open breaker and schedule the next allow after the current block."""
99 async with self._lock:
100 self._state = "open"
101 self._open_until = time.time() + self._block
102 self._block = min(self._block * 2, self._max_block)
104 async def state(self) -> str:
105 """Return breaker state: closed | half_open | open."""
106 async with self._lock:
107 return self._state
110# ── LLM Client ────────────────────────────────────────────────────────────────
111class LLMClient:
112 """High-level client that wraps remote (OpenRouter) and local (Ollama) calls."""
114 def __init__(self, config: LLMConfig, logger: Optional[logging.Logger] = None) -> None:
115 self._cfg = config
116 self._log = logger or logging.getLogger("llm")
117 self._session_holder = AsyncSessionHolder(
118 limit=config.http_conn_limit, timeout_total=config.http_timeout_total
119 )
120 self._breaker = CircuitBreaker(
121 base_block=config.breaker_base_block_sec, max_block=config.breaker_max_block_sec
122 )
123 self._loop: Optional[asyncio.AbstractEventLoop] = None
125 # ---- Loop attachment for cross-thread sync helpers -----------------------
126 def attach_loop(self, loop: asyncio.AbstractEventLoop) -> None:
127 """Attach the running loop to enable sync helpers from other threads."""
128 self._loop = loop
130 # ---- Public helpers -------------------------------------------------------
131 async def is_remote_blocked(self) -> bool:
132 """Async check if remote path is currently blocked by the breaker."""
133 return (await self._breaker.state()) == "open"
135 def is_remote_blocked_sync(self) -> bool:
136 """Thread-safe check usable from non-async contexts."""
137 if self._loop is None: 137 ↛ 138line 137 didn't jump to line 138, because the condition on line 137 was never true
138 return False
139 fut = asyncio.run_coroutine_threadsafe(self._breaker.state(), self._loop)
140 try:
141 return fut.result(timeout=0.5) == "open"
142 except Exception:
143 return False
145 async def close(self) -> None:
146 """Close underlying resources."""
147 await self._session_holder.close()
149 # ---- Core API -------------------------------------------------------------
150 async def query_model(
151 self,
152 messages: Optional[List[dict]] = None,
153 sys_prompt: Optional[str] = None,
154 ctx_txt: Optional[str] = None,
155 q: Optional[str] = None,
156 timeout_sec: int = 240,
157 ) -> Tuple[str, bool]:
158 """Query remote model with local fallback. Returns (text, used_fallback)."""
159 if messages is None: 159 ↛ 172line 159 didn't jump to line 172, because the condition on line 159 was always true
160 if not (sys_prompt and ctx_txt and q):
161 raise ValueError(
162 "query_model requires either `messages` or (`sys_prompt`, `ctx_txt`, `q`)."
163 )
164 messages = [
165 {"role": "system", "content": sys_prompt.strip()},
166 {
167 "role": "user",
168 "content": f"CONTEXT:\n{ctx_txt.strip()}\n\nQUESTION: {q.strip()}\nANSWER:",
169 },
170 ]
172 used_fallback = False
174 if await self._breaker.allow():
175 text, category = await self._call_openrouter(messages)
176 if text is not None:
177 await self._breaker.on_success(base_block=self._cfg.breaker_base_block_sec)
178 return text, used_fallback
179 await self._breaker.on_failure()
180 used_fallback = True
181 else:
182 used_fallback = True
184 # Fallback to local model
185 if sys_prompt is None or ctx_txt is None or q is None: 185 ↛ 186line 185 didn't jump to line 186, because the condition on line 185 was never true
186 sys_prompt = next((m.get("content", "") for m in messages if m.get("role") == "system"), "")
187 user = next((m.get("content", "") for m in messages if m.get("role") == "user"), "")
188 prompt_text = f"{sys_prompt or ''}\n\n{user or ''}"
189 else:
190 prompt_text = (
191 f"{sys_prompt.strip()}\n\nCONTEXT:\n{ctx_txt.strip()}\n\nQUESTION: {q.strip()}\nANSWER:"
192 )
193 text = await self.call_local_llm(prompt_text, timeout_sec=timeout_sec)
194 return text, used_fallback
196 async def call_local_llm(self, prompt_text: str, timeout_sec: Optional[int] = None) -> str:
197 """Call local Ollama model with a raw prompt."""
198 session = await self._session_holder.get()
199 try:
200 async with session.post( 200 ↛ exit, 200 ↛ 2052 missed branches: 1) line 200 didn't return from function 'call_local_llm', because the return on line 209 wasn't executed or the return on line 210 wasn't executed, 2) line 200 didn't jump to line 205, because
201 self._cfg.ollama_url,
202 json={"model": self._cfg.local_model, "prompt": prompt_text, "stream": False},
203 timeout=ClientTimeout(total=timeout_sec or self._cfg.http_timeout_total),
204 ) as resp:
205 raw = await resp.text()
206 try:
207 data = json.loads(raw)
208 except Exception:
209 return "⚠️ Local LLM returned non-JSON response."
210 return (data or {}).get("response", "❌ No response.")
211 except asyncio.TimeoutError: 211 ↛ 213line 211 didn't jump to line 213
212 return "⚠️ Local LLM did not respond (timeout)."
213 except Exception as exc:
214 return f"⚠️ Local LLM error: {exc}"
216 # ---- Internals ------------------------------------------------------------
217 async def _call_openrouter(self, messages: List[dict]) -> Tuple[Optional[str], Optional[str]]:
218 """Call OpenRouter with retries; return (text, error_category).
220 error_category ∈ {None, 'auth', 'transient', 'other'}.
221 """
222 session = await self._session_holder.get()
223 last_exc: Optional[BaseException] = None
225 for attempt in range(1, self._cfg.or_retries + 1):
226 try:
227 async with session.post(
228 self._cfg.api_url,
229 json={
230 "model": self._cfg.or_model,
231 "messages": messages,
232 "max_tokens": self._cfg.or_max_tokens,
233 },
234 headers={"Authorization": f"Bearer {self._cfg.openrouter_api_key}"},
235 ) as resp:
236 if resp.status == 401:
237 self._log.error("OpenRouter unauthorized (401)")
238 return None, "auth"
239 if resp.status in {429, 500, 502, 503, 504}:
240 raise RuntimeError(f"HTTP {resp.status}")
242 raw = await resp.text()
243 try:
244 data: Any = json.loads(raw)
245 except Exception:
246 self._log.warning("OpenRouter non-JSON response: %s", raw[:500])
247 return None, "other"
249 msg = (
250 (((data or {}).get("choices") or [{}])[0].get("message") or {}).get("content")
251 )
252 if not msg:
253 self._log.warning("OpenRouter empty content: %s", str(data)[:500])
254 return None, "other"
255 return msg, None
256 except Exception as exc:
257 last_exc = exc
258 wait = min(2**attempt, 10) + 0.1 * attempt
259 self._log.warning(
260 "OpenRouter attempt %s/%s failed: %s; backoff %.1fs",
261 attempt,
262 self._cfg.or_retries,
263 exc,
264 wait,
265 )
266 if attempt < self._cfg.or_retries:
267 await asyncio.sleep(wait)
269 if last_exc is not None:
270 self._log.error("OpenRouter retries exhausted: %s", last_exc)
271 return None, "transient"