Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/wrappers.py: 39%
100 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""LLM client wrappers — caching and provenance enrichment.
3Both wrappers preserve the :class:`~lexigram.contracts.ai.llm.LLMClientProtocol`
4shape (``complete`` / ``stream_chat`` / ``health_check`` / ``close``) so they
5compose freely with each other and with :class:`LLMAuditBridge.wrap`.
7Recommended layering (outer → inner)::
9 LLMCacheWrapper( # short-circuits cache hits before audit fires
10 LLMAuditBridge.wrap( # logs every real LLM call
11 CompletionEnricher( # fills provider / model_revision / prompt_hash
12 provider_client
13 )
14 )
15 )
17The cache wrapper sits outermost so cache hits don't trigger audit log
18entries (no real LLM call happened). The enricher sits innermost so the
19audit + cache layers see fully-populated :class:`Completion` provenance.
20"""
22from __future__ import annotations
24from collections.abc import Sequence
25from typing import TYPE_CHECKING, Any, cast
27from lexigram.ai.llm.caching.types import build_llm_cache_key
28from lexigram.contracts.ai.exceptions import LLMError
29from lexigram.contracts.ai.llm import (
30 ChatMessageProtocol,
31 CompletionProtocol,
32 LLMClientProtocol,
33 StreamChunk,
34)
35from lexigram.contracts.infra import AsyncStream
36from lexigram.logging import get_logger
37from lexigram.result import Result
38from lexigram.security.hashing import ambient as hashing
40if TYPE_CHECKING:
41 from lexigram.ai.llm.protocols import LLMCacheProtocol
42 from lexigram.ai.llm.types import Completion
44logger = get_logger(__name__)
47def _hash_messages_for_cache_key(messages: Any) -> str:
48 """Build a stable prompt string from ``messages`` for cache-key hashing.
50 Accepts any iterable that yields chat-message-like objects. Items may
51 expose ``.role`` and ``.content`` attributes, or be plain dicts with
52 those keys, or any other shape — they're stringified deterministically
53 for hashing only.
54 """
55 parts: list[str] = []
56 try:
57 iterator = iter(messages)
58 except TypeError:
59 return str(messages)
61 for msg in iterator:
62 role = getattr(msg, "role", None)
63 content = getattr(msg, "content", None)
64 if role is None and isinstance(msg, dict):
65 role = msg.get("role")
66 content = msg.get("content")
67 parts.append(f"{role}:{content}")
68 return "|".join(parts)
71def _compute_prompt_hash(messages: Any) -> bytes:
72 """Compute a SHA-256 digest over the canonical message representation."""
73 prompt = _hash_messages_for_cache_key(messages)
74 return bytes.fromhex(hashing.hash_hex(prompt))
77class _CacheWrappedClient:
78 """Caches ``complete()`` results keyed by ``build_llm_cache_key``.
80 Cache hits short-circuit before the wrapped client runs, so neither the
81 underlying provider nor downstream wrappers (audit, etc.) see the call.
82 """
84 def __init__(
85 self,
86 client: LLMClientProtocol,
87 cache: LLMCacheProtocol,
88 provider: str,
89 model: str,
90 model_revision: str | None = None,
91 ttl_seconds: float | None = None,
92 ) -> None:
93 self._client = client
94 self._cache = cache
95 self._provider = provider
96 self._model = model
97 self._model_revision = model_revision
98 self._ttl_seconds = ttl_seconds
100 def _key(self, messages: Any, kwargs: dict[str, Any]) -> str:
101 prompt = _hash_messages_for_cache_key(messages)
102 # Include kwargs in the prompt portion so different temperature /
103 # max_tokens / tool sets produce distinct entries.
104 kw_repr = ",".join(f"{k}={kwargs[k]}" for k in sorted(kwargs))
105 prompt_with_kwargs = f"{prompt}::{kw_repr}" if kw_repr else prompt
106 return build_llm_cache_key(
107 provider=self._provider,
108 model=self._model,
109 prompt=prompt_with_kwargs,
110 model_revision=self._model_revision,
111 prompt_hash=_compute_prompt_hash(messages),
112 )
114 async def complete(
115 self,
116 messages: Sequence[ChatMessageProtocol],
117 **kwargs: Any,
118 ) -> Result[CompletionProtocol, LLMError]:
119 key = self._key(messages, kwargs)
120 cached = await self._cache.get(key)
121 if cached is not None:
122 logger.debug("llm_cache_hit", key_prefix=key[:8])
123 return cast("Result[CompletionProtocol, LLMError]", cached)
125 result = await self._client.complete(messages, **kwargs)
126 if result.is_ok():
127 try:
128 await self._cache.set(key, result, ttl=self._ttl_seconds)
129 except Exception:
130 logger.exception("llm_cache_set_failed", key_prefix=key[:8])
131 return result
133 def stream_chat(
134 self,
135 messages: list[ChatMessageProtocol],
136 **kwargs: Any,
137 ) -> AsyncStream[StreamChunk, LLMError]:
138 # Streaming is not cached — pass-through.
139 return self._client.stream_chat(messages, **kwargs)
141 async def health_check(self, timeout: float = 5.0) -> Any:
142 return await self._client.health_check(timeout=timeout)
144 async def close(self) -> None:
145 return await self._client.close()
148class LLMCacheWrapper:
149 """Static factory namespace for the cache wrapper.
151 Use :meth:`wrap` to attach a cache to a client.
152 """
154 @staticmethod
155 def wrap(
156 client: LLMClientProtocol,
157 cache: LLMCacheProtocol,
158 *,
159 provider: str,
160 model: str,
161 model_revision: str | None = None,
162 ttl_seconds: float | None = None,
163 ) -> LLMClientProtocol:
164 """Wrap *client* with a content-keyed cache.
166 Args:
167 client: The underlying LLM client.
168 cache: Backing cache (must implement
169 :class:`~lexigram.ai.llm.protocols.LLMCacheProtocol`).
170 provider: Provider identifier (e.g. ``"anthropic"``).
171 model: Model name (e.g. ``"claude-sonnet-4-6"``).
172 model_revision: Optional pinned revision; participates in the key.
173 ttl_seconds: Optional per-entry TTL; falls back to cache default.
175 Returns:
176 A wrapped client that short-circuits cache hits.
177 """
178 return _CacheWrappedClient(
179 client=client,
180 cache=cache,
181 provider=provider,
182 model=model,
183 model_revision=model_revision,
184 ttl_seconds=ttl_seconds,
185 )
188class _EnrichedClient:
189 """Fills missing provenance fields on returned :class:`Completion` objects.
191 The provider clients populate ``model`` and ``content`` reliably; revision
192 is only set when the response carries it. This wrapper computes the
193 ``prompt_hash`` from inputs and back-fills ``provider`` and
194 ``model_revision`` from configuration when the underlying client did not
195 set them. Already-populated fields are never overwritten.
196 """
198 def __init__(
199 self,
200 client: LLMClientProtocol,
201 provider: str,
202 model: str,
203 model_revision: str | None = None,
204 ) -> None:
205 self._client = client
206 self._provider = provider
207 self._model = model
208 self._model_revision = model_revision
210 def _enrich(self, completion: Completion, messages: Any) -> Completion:
211 if not completion.provider:
212 completion.provider = self._provider
213 if not completion.model:
214 completion.model = self._model
215 if not completion.model_revision and self._model_revision:
216 completion.model_revision = self._model_revision
217 if completion.prompt_hash is None:
218 completion.prompt_hash = _compute_prompt_hash(messages)
219 return completion
221 async def complete(
222 self,
223 messages: Sequence[ChatMessageProtocol],
224 **kwargs: Any,
225 ) -> Result[CompletionProtocol, LLMError]:
226 result = await self._client.complete(messages, **kwargs)
227 if result.is_ok():
228 try:
229 completion = result.unwrap()
230 self._enrich(cast("Completion", completion), messages)
231 except Exception:
232 logger.exception("completion_enrichment_failed")
233 return result
235 def stream_chat(
236 self,
237 messages: list[ChatMessageProtocol],
238 **kwargs: Any,
239 ) -> AsyncStream[StreamChunk, LLMError]:
240 return self._client.stream_chat(messages, **kwargs)
242 async def health_check(self, timeout: float = 5.0) -> Any:
243 return await self._client.health_check(timeout=timeout)
245 async def close(self) -> None:
246 return await self._client.close()
249class CompletionEnricher:
250 """Static factory namespace for the provenance-enrichment wrapper.
252 Use :meth:`wrap` to attach enrichment to a client.
253 """
255 @staticmethod
256 def wrap(
257 client: LLMClientProtocol,
258 *,
259 provider: str,
260 model: str,
261 model_revision: str | None = None,
262 ) -> LLMClientProtocol:
263 """Wrap *client* so returned ``Completion`` objects carry provenance.
265 Args:
266 client: The underlying LLM client.
267 provider: Provider name used to fill ``Completion.provider`` when
268 the underlying client leaves it empty.
269 model: Model name used to fill ``Completion.model`` when missing.
270 model_revision: Pinned revision used to fill
271 ``Completion.model_revision`` when the underlying client did
272 not populate it from the response.
274 Returns:
275 A wrapped client that enriches successful completions in place.
276 """
277 return _EnrichedClient(
278 client=client,
279 provider=provider,
280 model=model,
281 model_revision=model_revision,
282 )
285__all__ = ["CompletionEnricher", "LLMCacheWrapper"]