Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/clients/base.py: 29%
161 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Base class for Lexigram LLM clients.
3Provides common retry, exponential backoff, circuit breaker, and token limit
4checks for all LLM providers. Subclasses implement provider-specific logic.
5"""
7from __future__ import annotations
9from abc import ABC, abstractmethod
10import asyncio
11from collections.abc import AsyncIterator
12import types
13from typing import Any, Protocol, Self, cast
15from lexigram.ai.llm.config import ClientConfig
16from lexigram.ai.llm.exceptions import (
17 LLMAuthenticationError,
18 LLMError,
19 LLMRateLimitError,
20 LLMTimeoutError,
21)
22from lexigram.ai.llm.health import ProviderCircuitBreaker
23from lexigram.ai.llm.thinking import normalize_thinking_text
24from lexigram.ai.llm.types import (
25 ChatMessage,
26 Completion,
27 StreamChunk,
28 ToolCall,
29)
30from lexigram.contracts.ai.thinking import ThinkingResult
31from lexigram.contracts.core.health import HealthCheckResult
32from lexigram.contracts.infra import AsyncStream
33from lexigram.logging import (
34 get_logger,
35)
36from lexigram.result import Err, Ok, Result
38logger = get_logger(__name__)
41class _ClientCircuitBreaker(Protocol):
42 async def is_available(self) -> bool: ...
43 async def record_success(self) -> None: ...
44 async def record_failure(self, error: Exception) -> None: ...
47class AbstractLLMClient(ABC):
48 """Base class for all LLM client implementations.
50 Provides common retry, backoff, circuit breaker, and pre-flight token
51 checking. Subclasses implement only provider-specific API calls and
52 response parsing via `_do_complete` and `_do_stream_chat`.
53 """
55 def __init__(
56 self,
57 config: ClientConfig,
58 max_retries: int = 3,
59 circuit_breaker: _ClientCircuitBreaker | None = None,
60 ):
61 """Initialize the base client with config.
63 Args:
64 config: LLM configuration object.
65 max_retries: Maximum number of retries for transient errors.
66 circuit_breaker: Optional circuit breaker for resilience.
67 If not provided, creates a default ProviderCircuitBreaker.
68 """
69 self.config = config
70 self.max_retries = max_retries
71 self._circuit_breaker = circuit_breaker or ProviderCircuitBreaker(
72 provider_name=config.provider.value
73 )
74 self._closed = False
76 async def complete(
77 self,
78 messages: list[ChatMessage],
79 **kwargs: Any,
80 ) -> Result[Completion, LLMError]:
81 """Template method: retry + circuit breaker + pre-flight + subclass."""
83 if not await self._circuit_breaker.is_available():
84 return Err(LLMError("Circuit breaker is open due to recent failures."))
86 if await self._exceeds_token_limit(messages, **kwargs):
87 return Err(LLMError("Request exceeds model maximum context window."))
89 for attempt in range(self.max_retries + 1):
90 try:
91 result = await self._do_complete(messages, **kwargs)
92 if result.is_ok():
93 await self._circuit_breaker.record_success()
94 return Ok(self._normalize_thinking(result.unwrap()))
96 err = result.unwrap_err()
97 if self._is_retryable_error(err):
98 await self._circuit_breaker.record_failure(err)
99 if attempt < self.max_retries:
100 await self._backoff(attempt)
101 continue
103 return result # Non-retryable or max retries reached
105 except Exception as e:
106 await self._circuit_breaker.record_failure(e)
107 if not self._is_retryable_exception(e) or attempt == self.max_retries:
108 raise
109 await self._backoff(attempt)
111 return Err(LLMError("Max retries exceeded"))
113 def stream_chat(
114 self,
115 messages: list[ChatMessage],
116 **kwargs: Any,
117 ) -> AsyncStream[StreamChunk, LLMError]:
118 """Template method for streaming chat with retry + circuit breaker.
120 Returns an ``AsyncStream`` immediately. All setup, retry, and
121 circuit-breaker logic is executed lazily inside the stream's async
122 generator. Setup failures and mid-stream failures are both surfaced
123 through ``AsyncStream``'s typed error channel.
124 """
126 async def _establish_stream() -> AsyncIterator[StreamChunk]:
127 """Async generator that handles setup, retry, and streaming."""
128 # Check circuit breaker
129 if not await self._circuit_breaker.is_available():
130 raise LLMError("Circuit breaker is open due to recent failures.")
132 # Check token limit
133 if await self._exceeds_token_limit(messages, **kwargs):
134 raise LLMError("Request exceeds model maximum context window.")
136 # Retry loop for setup
137 for attempt in range(self.max_retries + 1):
138 try:
139 result = await self._do_stream_chat(messages, **kwargs)
140 if result.is_ok():
141 await self._circuit_breaker.record_success()
142 stream_iter = result.unwrap()
143 # Stream the chunks with timeout
144 async for chunk in self._timeout_stream(
145 stream_iter, self.config.timeout / 2
146 ):
147 yield chunk
148 return
150 err = result.unwrap_err()
151 if self._is_retryable_error(err):
152 await self._circuit_breaker.record_failure(err)
153 if attempt < self.max_retries:
154 await self._backoff(attempt)
155 continue
157 raise err
159 except Exception as e:
160 await self._circuit_breaker.record_failure(e)
161 if (
162 not self._is_retryable_exception(e)
163 or attempt == self.max_retries
164 ):
165 raise
166 await self._backoff(attempt)
168 raise LLMError("Max retries exceeded")
170 return AsyncStream(
171 _establish_stream(),
172 error_adapter=self._coerce_stream_error,
173 )
175 async def chat(
176 self,
177 messages: list[ChatMessage],
178 tools: list[ToolCall] | None = None,
179 **kwargs: Any,
180 ) -> Result[Completion, LLMError]:
181 """Template method: completion with tools."""
183 if not await self._circuit_breaker.is_available():
184 return Err(LLMError("Circuit breaker is open due to recent failures."))
186 if await self._exceeds_token_limit(messages, **kwargs):
187 return Err(LLMError("Request exceeds model maximum context window."))
189 for attempt in range(self.max_retries + 1):
190 try:
191 result = await self._do_chat(messages, tools, **kwargs)
192 if result.is_ok():
193 await self._circuit_breaker.record_success()
194 return Ok(self._normalize_thinking(result.unwrap()))
196 err = result.unwrap_err()
197 if self._is_retryable_error(err):
198 await self._circuit_breaker.record_failure(err)
199 if attempt < self.max_retries:
200 await self._backoff(attempt)
201 continue
203 return result
205 except Exception as e:
206 await self._circuit_breaker.record_failure(e)
207 if not self._is_retryable_exception(e) or attempt == self.max_retries:
208 raise
209 await self._backoff(attempt)
211 return Err(LLMError("Max retries exceeded"))
213 async def _timeout_stream(
214 self, stream: AsyncIterator[StreamChunk], chunk_timeout: float
215 ) -> AsyncIterator[StreamChunk]:
216 """Wrap an async iterator with a per-chunk timeout."""
217 try:
218 while True:
219 chunk = await asyncio.wait_for(anext(stream), timeout=chunk_timeout)
220 yield chunk
221 except StopAsyncIteration:
222 return
223 except TimeoutError:
224 logger.warning("Stream chunk timed out after %s seconds", chunk_timeout)
225 raise LLMError("Streaming timed out between chunks")
227 def _coerce_stream_error(self, exc: Exception) -> LLMError:
228 """Convert iterator failures into typed stream errors."""
229 if isinstance(exc, LLMError):
230 return exc
231 return LLMError(str(exc))
233 def _is_retryable_error(self, err: LLMError) -> bool:
234 """Determine if a returned error should trigger a retry.
236 Rate limits and request timeouts are NOT retried: the router's
237 cascade advances to the next entry instead of hammering a throttled
238 or overloaded model in place (a timeout retry costs another full
239 timeout window).
240 """
241 if isinstance(err, (LLMRateLimitError, LLMTimeoutError)):
242 return False
243 if "connection" in str(err).lower() or "timeout" in str(err).lower():
244 return True
245 return "503" in str(err) or "502" in str(err) or "504" in str(err)
247 def _is_retryable_exception(self, e: Exception) -> bool:
248 """Determine if a raised exception should trigger a retry.
250 Rate limits and request timeouts fail fast (cascade advancement
251 handles them); authentication and validation errors are permanent.
252 """
253 if isinstance(e, (LLMRateLimitError, LLMTimeoutError)):
254 return False
255 if isinstance(e, LLMAuthenticationError):
256 return False
257 return not ("invalid" in str(e).lower() or "not found" in str(e).lower())
259 async def _backoff(self, attempt: int) -> None:
260 """Apply exponential backoff before the next retry."""
261 delay = (2**attempt) + (0.1 * attempt) # basic jitter
262 logger.info("Applying backoff of %.2fs before retry %d", delay, attempt + 1)
263 await asyncio.sleep(delay)
265 async def _exceeds_token_limit(
266 self, messages: list[ChatMessage], **kwargs: Any
267 ) -> bool:
268 """Pre-flight check: return True if messages exceed the configured context window.
270 Uses :class:`~lexigram.ai.llm.pricing.tokens.TiktokenCounter` for accurate
271 model-specific counting (falls back to character estimation when tiktoken is
272 unavailable). The check is skipped if ``context_window`` is not set in
273 ``config.extra``.
275 Args:
276 messages: Chat messages to count.
277 **kwargs: Ignored; accepted for call-site compatibility.
279 Returns:
280 ``True`` when the token count exceeds the context window limit,
281 ``False`` otherwise (including when no limit is configured).
282 """
283 context_window: int | None = self.config.extra.get("context_window")
284 if context_window is None:
285 return False
287 from lexigram.ai.llm.pricing.tokens import TiktokenCounter
289 counter = TiktokenCounter(model=self.config.model)
290 token_count = counter.count_messages(messages)
291 return token_count > context_window
293 async def close(self) -> None:
294 """Close the client and cleanup resources."""
295 self._closed = True
297 def _apply_thinking(self, params: dict[str, Any]) -> None:
298 """Inject provider-specific thinking parameters into an API payload dict.
300 Called from ``_do_complete``, ``_do_stream_chat``, and ``_do_chat`` in
301 each provider immediately after the base payload is assembled and before
302 the API call is made. The default implementation is a no-op so providers
303 without thinking support require no change.
305 Override in providers that support extended thinking/reasoning to avoid
306 repeating the injection logic across all three call-path methods.
308 Args:
309 params: Mutable API payload dict that will be sent to the provider.
310 """
312 def _normalize_thinking(self, completion: Completion) -> Completion:
313 """Strip inline thinking tags from completion content and populate ThinkingResult.
315 This method is a universal post-processing step applied in the
316 ``complete()`` and ``chat()`` template methods to ALL clients. It handles
317 models that embed reasoning tokens directly in their text output using
318 known tag formats (Qwen3 ``<think>``, Gemma-4 ``<|channel>thought``, etc.).
320 Guard: if ``completion.thinking`` is already set (Anthropic, Gemini, Bedrock,
321 OpenRouter all extract thinking natively from structured API fields), this
322 method returns the completion unchanged — normalization is a no-op for them.
324 Args:
325 completion: Raw completion from ``_do_complete`` or ``_do_chat``.
327 Returns:
328 A new Completion with clean ``content`` and populated ``thinking``,
329 or the original unchanged if no inline thinking was detected.
330 """
331 if completion.thinking is not None:
332 return completion
334 clean_content, thinking_text = normalize_thinking_text(completion.content)
335 if thinking_text is None:
336 return completion
338 return cast(
339 "Completion",
340 completion.model_copy(
341 update={
342 "content": clean_content,
343 "thinking": ThinkingResult(content=thinking_text),
344 }
345 ),
346 )
348 @abstractmethod
349 async def _do_complete(
350 self,
351 messages: list[ChatMessage],
352 **kwargs: Any,
353 ) -> Result[Completion, LLMError]:
354 """Provider-specific complete implementation."""
356 @abstractmethod
357 async def _do_stream_chat(
358 self,
359 messages: list[ChatMessage],
360 **kwargs: Any,
361 ) -> Result[AsyncIterator[StreamChunk], LLMError]:
362 """Provider-specific streaming implementation."""
364 @abstractmethod
365 async def _do_chat(
366 self,
367 messages: list[ChatMessage],
368 tools: list[ToolCall] | None = None,
369 **kwargs: Any,
370 ) -> Result[Completion, LLMError]:
371 """Provider-specific chat (with tools) implementation."""
373 @abstractmethod
374 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult:
375 """Perform health check against the provider."""
377 async def __aenter__(self) -> Self:
378 return self
380 async def __aexit__(
381 self,
382 exc_type: type[BaseException] | None,
383 exc_val: BaseException | None,
384 exc_tb: types.TracebackType | None,
385 ) -> None:
386 await self.close()