Coverage for src / lexigram / contracts / ai / llm.py: 5%
154 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-08-19 05:41 +0800
1"""LLM client and prompt protocols."""
3from __future__ import annotations
5from collections.abc import Sequence
6from dataclasses import dataclass, field
7from dataclasses import replace as _dc_replace
8from enum import StrEnum
9from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
11from lexigram.contracts.ai.exceptions import ExtractionError, LLMError
12from lexigram.contracts.ai.multimodal import MessageContent
13from lexigram.contracts.ai.thinking import ThinkingConfig, ThinkingResult
14from lexigram.contracts.infra import AsyncStream
16if TYPE_CHECKING:
17 from lexigram.contracts.ai import ToolDefinition
18 from lexigram.contracts.core import HealthCheckResult
19 from lexigram.contracts.core.result import Result
22@runtime_checkable
23class ChatMessageProtocol(Protocol):
24 """Structural protocol for chat message objects.
26 Any object with ``role`` and ``content`` attributes satisfies this
27 protocol. Implemented concretely by ``lexigram.ai.llm.types.ChatMessage``.
28 """
30 @property
31 def role(self) -> Any:
32 """Message role (e.g. user, assistant, system)."""
34 @property
35 def content(self) -> MessageContent:
36 """Message content."""
39@runtime_checkable
40class CompletionProtocol(Protocol):
41 """Structural protocol for LLM completion objects.
43 Any object with ``content`` and ``model`` attributes satisfies this
44 protocol. Implemented concretely by ``lexigram.ai.llm.types.Completion``.
45 """
47 content: str
48 """The generated completion text."""
50 model: str
51 """The model that produced this completion."""
53 thinking: ThinkingResult | None
54 """Normalised thinking/reasoning output, or ``None`` when not enabled."""
56 usage: dict[str, int] | None
57 """Token usage statistics (prompt, completion, total tokens) or None."""
60@runtime_checkable
61class LLMClientProtocol(Protocol):
62 """Protocol for LLM client implementations.
64 All LLM providers (OpenAI, Anthropic, Google, etc.) should
65 implement this interface.
66 """
68 async def complete(
69 self,
70 messages: Sequence[ChatMessageProtocol],
71 *,
72 model: str | None = None,
73 temperature: float | None = None,
74 max_tokens: int | None = None,
75 tools: Sequence[ToolDefinition] | None = None,
76 stop_sequences: Sequence[str] | None = None,
77 **kwargs: Any,
78 ) -> Result[CompletionProtocol, LLMError]:
79 """Generate a completion from messages.
81 Args:
82 messages: List of chat messages.
83 model: Optional model override.
84 temperature: Optional sampling temperature.
85 max_tokens: Optional max output tokens.
86 tools: Optional tool definitions.
87 stop_sequences: Optional list of stop sequences for early termination.
88 **kwargs: Provider-specific options.
90 Returns:
91 ``Ok(Completion)`` on success, ``Err(LLMError)`` on recoverable failure.
92 """
93 ...
95 def stream_chat(
96 self,
97 messages: list[ChatMessageProtocol],
98 *,
99 model: str | None = None,
100 temperature: float | None = None,
101 max_tokens: int | None = None,
102 tools: list[ToolDefinition] | None = None,
103 stop_sequences: list[str] | None = None,
104 **kwargs: Any,
105 ) -> AsyncStream[StreamChunk, LLMError]:
106 """Start a streaming completion.
108 This method returns an ``AsyncStream`` immediately (not wrapped in
109 ``Result``). The stream is established lazily when iteration begins.
110 Setup failures and mid-stream failures are both surfaced through the
111 stream's typed error channel.
113 Args:
114 messages: List of chat messages.
115 model: Optional model override.
116 temperature: Optional sampling temperature.
117 max_tokens: Optional max output tokens.
118 tools: Optional tool definitions.
119 stop_sequences: Optional list of stop sequences for early termination.
120 **kwargs: Provider-specific options.
122 Returns:
123 ``AsyncStream[StreamChunk, LLMError]`` that yields chunks or
124 surfaces typed errors through terminal operations (``collect()``,
125 ``first()``, ``drain()``).
126 """
127 ...
129 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult:
130 """Perform a lightweight connectivity check.
132 Returns:
133 Structured health check result.
134 """
135 ...
137 async def close(self) -> None:
138 """Close the client and release resources."""
139 ...
142@runtime_checkable
143class EmbeddingClientProtocol(Protocol):
144 """Protocol for embedding client implementations."""
146 async def embed(self, texts: list[str]) -> list[list[float]]:
147 """Generate embeddings for texts.
149 Args:
150 texts: List of strings to embed.
152 Returns:
153 List of embedding vectors.
154 """
155 ...
157 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult:
158 """Perform a lightweight connectivity check.
160 Returns:
161 Structured health check result.
162 """
163 ...
165 async def close(self) -> None:
166 """Close the client and release resources."""
167 ...
170@runtime_checkable
171class StructuredExtractorProtocol(Protocol):
172 """Protocol for extracting typed, validated data from LLM responses.
174 Implementations wrap an ``LLMClientProtocol``, instruct the model to respond
175 with JSON matching a given schema, and validate the parsed response
176 against a target model type.
177 """
179 async def extract(
180 self,
181 prompt: str | list[ChatMessageProtocol],
182 output_model: type[Any],
183 *,
184 max_retries: int = 2,
185 model: str | None = None,
186 **kwargs: Any,
187 ) -> Result[Any, ExtractionError | LLMError]:
188 """Extract structured data from an LLM response.
190 Args:
191 prompt: Text prompt or list of chat messages.
192 output_model: Model class to validate the response against.
193 max_retries: Number of additional attempts on parse/validation failure.
194 model: Optional LLM model name override.
195 **kwargs: Extra options forwarded to the underlying LLM client.
197 Returns:
198 ``Ok(output_model_instance)`` on success.
199 ``Err(ExtractionError)`` on parse/validation failure.
200 ``Err(LLMError)`` when the underlying client returns a recoverable
201 provider/model failure.
202 """
203 ...
206@runtime_checkable
207class PromptTemplateProtocol(Protocol):
208 """Protocol for prompt templates that render to text or chat messages."""
210 @property
211 def name(self) -> str:
212 """Unique name identifying this template."""
213 ...
215 def render(self, **kwargs: Any) -> str | list[dict[str, str]]:
216 """Render the template with the supplied variable values.
218 Args:
219 **kwargs: Variable name → value pairs.
221 Returns:
222 A plain string or a list of ``{"role": ..., "content": ...}`` dicts.
223 """
224 ...
226 def get_variables(self) -> list[str]:
227 """Return the list of declared variable names."""
228 ...
231@runtime_checkable
232class PromptRegistryProtocol(Protocol):
233 """Protocol for a named template registry."""
235 def register(
236 self,
237 name: str,
238 template: PromptTemplateProtocol,
239 *,
240 overwrite: bool = False,
241 ) -> None:
242 """Register a template under *name*."""
243 ...
245 def get(self, name: str) -> PromptTemplateProtocol:
246 """Retrieve a template by name."""
247 ...
249 def list_names(self) -> list[str]:
250 """Return all registered template names."""
251 ...
254@runtime_checkable
255class TokenCounterProtocol(Protocol):
256 """Model-aware token counter.
258 Implementations use the exact tokenizer for the target model (tiktoken for
259 OpenAI, AutoTokenizer for HuggingFace, mistral-common for Mistral). A
260 character-estimate fallback is always available.
262 Note: Methods are synchronous because tokenization is CPU-bound with no I/O.
263 """
265 def count(self, text: str) -> int:
266 """Count tokens in a text string."""
267 ...
269 def count_messages(self, messages: list[ChatMessageProtocol]) -> int:
270 """Count tokens in a list of chat messages, including message overhead."""
271 ...
273 @property
274 def model(self) -> str:
275 """The model this counter is calibrated for."""
276 ...
279@runtime_checkable
280class PromptAssemblerProtocol(Protocol):
281 """Assembles prompt layers in cache-friendly static-to-dynamic order.
283 The static-first ordering maximizes provider-side KV cache reuse. The
284 assembler also injects provider-specific cache annotations (Anthropic's
285 cache_control breakpoints, DeepSeek's 64-token padding, etc.).
286 """
288 def assemble(
289 self,
290 system: str,
291 tools: list[ToolDefinition] | None,
292 reference_docs: list[str] | None,
293 few_shot: list[ChatMessage] | None,
294 history: list[ChatMessage],
295 query: str,
296 provider: str,
297 dynamic_metadata: str | None = None,
298 ) -> list[ChatMessage]:
299 """Assemble a complete prompt message list.
301 Args:
302 system: System instructions (static, cached).
303 tools: Tool/function definitions.
304 reference_docs: Reference documents as text (semi-static).
305 few_shot: Few-shot example messages (static, cached).
306 history: Chat history messages (dynamic).
307 query: Current user query (dynamic).
308 provider: Provider name for cache annotation strategy.
309 dynamic_metadata: Timestamps, user IDs, etc. (dynamic).
311 Returns:
312 Ordered list of ChatMessage instances ready for the LLM client.
313 """
314 ...
317@runtime_checkable
318class PromptRendererProtocol(Protocol):
319 """Renders a prompt template with provided variables.
321 Consumed by: lexigram-ai-platform prompt submodule, lexigram-ai-rag.
322 """
324 def render(self, template: str, variables: dict[str, Any]) -> str:
325 """Render a template string with the supplied variable mapping.
327 Args:
328 template: The raw template string.
329 variables: Mapping of variable names to their values.
331 Returns:
332 The rendered prompt string.
333 """
334 ...
337@runtime_checkable
338class PromptOptimizerProtocol(Protocol):
339 """Optimizes a rendered prompt for a specific model.
341 Consumed by: lexigram-ai-platform prompt submodule.
342 """
344 async def optimize(self, prompt: str, model: str, max_tokens: int) -> str:
345 """Optimize a prompt for the target model and token budget.
347 Args:
348 prompt: The rendered prompt string to optimize.
349 model: Target model identifier.
350 max_tokens: Maximum token budget for the optimized prompt.
352 Returns:
353 The optimized prompt string.
354 """
355 ...
358@runtime_checkable
359class SemanticCacheProtocol(Protocol):
360 """Embedding-based semantic similarity cache for LLM responses.
362 Provides three-tier lookup: exact hash match (Tier 1), vector similarity
363 match (Tier 2), and cache miss (caller invokes LLM).
365 Placement note: Lives in ai/llm.py (not infra/cache/) because it carries
366 AI-domain semantics (model tracking in store(), LLM response caching).
367 """
369 async def lookup(self, query: str) -> str | None:
370 """Look up a query in the cache.
372 Checks Tier 1 (exact hash) then Tier 2 (vector similarity).
374 Args:
375 query: The user query string.
377 Returns:
378 Cached response string, or ``None`` on cache miss.
379 """
380 ...
382 async def store(self, query: str, response: str, model: str) -> None:
383 """Store a query-response pair in both tiers.
385 Args:
386 query: The user query string.
387 response: The LLM response to cache.
388 model: The model that produced the response.
389 """
390 ...
392 async def invalidate(self, query: str) -> bool:
393 """Invalidate a cached entry by query.
395 Args:
396 query: The user query string to invalidate.
398 Returns:
399 ``True`` if the entry was found and removed, ``False`` otherwise.
400 """
401 ...
404@dataclass(frozen=True)
405class StreamEvent:
406 """Protocol for streaming events."""
408 type: str
409 content: str | None = None
410 tool_call: Any | None = None
411 error: str | None = None
414@dataclass(frozen=True)
415class TokenBudget:
416 """Immutable token budget that pipeline stages consult.
418 Every field is calculated once at the start of a request. Stages read
419 remaining capacity via properties — they never mutate the budget.
420 Builder methods return new instances (immutable value semantics).
421 """
423 model_context_limit: int
424 reserved_for_output: int
425 system_prompt_tokens: int
426 tool_definitions_tokens: int
427 rag_context_tokens: int = 0
428 history_tokens: int = 0
430 @property
431 def total_used(self) -> int:
432 """Total tokens consumed by all components."""
433 return (
434 self.system_prompt_tokens
435 + self.tool_definitions_tokens
436 + self.rag_context_tokens
437 + self.history_tokens
438 + self.reserved_for_output
439 )
441 @property
442 def remaining(self) -> int:
443 """Tokens available for additional content."""
444 return max(0, self.model_context_limit - self.total_used)
446 @property
447 def remaining_for_history(self) -> int:
448 """Tokens available specifically for chat history."""
449 used = (
450 self.system_prompt_tokens
451 + self.tool_definitions_tokens
452 + self.rag_context_tokens
453 + self.reserved_for_output
454 )
455 return max(0, self.model_context_limit - used)
457 @property
458 def over_budget(self) -> bool:
459 """True if total consumption exceeds the model context limit."""
460 return self.total_used > self.model_context_limit
462 def with_rag_tokens(self, rag_tokens: int) -> TokenBudget:
463 """Return a new budget with updated RAG token count."""
464 return _dc_replace(self, rag_context_tokens=rag_tokens)
466 def with_history_tokens(self, history_tokens: int) -> TokenBudget:
467 """Return a new budget with updated history token count."""
468 return _dc_replace(self, history_tokens=history_tokens)
471class Role(StrEnum):
472 """Concrete chat message role constants shared across AI packages."""
474 SYSTEM = "system"
475 USER = "user"
476 ASSISTANT = "assistant"
477 TOOL = "tool"
478 FUNCTION = "function"
481@dataclass(frozen=True)
482class ChatMessage:
483 """Concrete chat message data class shared across AI packages.
485 Satisfies ``ChatMessageProtocol``. Suitable for construction in any
486 package that needs to build messages without importing from
487 ``lexigram-ai-llm``.
488 """
490 role: str
491 content: MessageContent
492 name: str | None = None
493 tool_call_id: str | None = None
494 tool_calls: list[ToolCall] | None = None
495 """Native tool calls requested by the LLM (assistant messages only).
497 When present, the message carries the function-call requests of an
498 assistant turn so the provider can re-emit them to the model before
499 the matching ``tool`` role responses.
500 """
501 thinking_blocks: list[dict[str, Any]] | None = None
502 """Raw provider thinking blocks for multi-turn re-injection.
504 Anthropic: list of ``{type, thinking, signature}`` dicts from a prior
505 assistant turn with extended thinking enabled. Must be passed back
506 verbatim to continue a thinking-enabled conversation.
507 """
508 metadata: dict[str, Any] | None = None
509 """Provider-specific metadata for cache control and annotations."""
512@dataclass(frozen=True)
513class Completion:
514 """Minimal completion data class shared across AI packages.
516 Satisfies ``CompletionProtocol``. Used by consumers that build
517 fallback completions without depending on ``lexigram-ai-llm``'s
518 full Pydantic model.
519 """
521 content: str
522 model: str
523 thinking: ThinkingResult | None = None
524 finish_reason: str | None = None
525 metadata: dict[str, Any] = field(default_factory=dict)
528@dataclass(frozen=True)
529class FunctionCall:
530 """Function call request from LLM."""
532 name: str
533 arguments: dict[str, Any] | str = field(default_factory=dict)
536@dataclass(frozen=True)
537class ToolCall:
538 """Tool call request from LLM."""
540 id: str
541 type: str = "function"
542 function: FunctionCall | None = None
545@dataclass(frozen=True)
546class TokenUsage:
547 """Token usage statistics."""
549 prompt_tokens: int
550 completion_tokens: int
551 total_tokens: int
554@dataclass(frozen=True)
555class StreamChunk:
556 """A chunk of streamed completion."""
558 delta: str | None = None
559 model: str | None = None
560 finish_reason: str | None = None
561 role: str | None = None
562 index: int = 0
563 thinking_delta: str | None = None
564 is_thinking: bool = False
567@runtime_checkable
568class CostEstimatorProtocol(Protocol):
569 """Estimates the monetary cost of LLM usage.
571 Implementations map token usage for a model/provider to a cost in
572 USD. Consumed by agents (governance cost tracking) and any package
573 that reports spend. When no estimator is wired, callers must skip
574 cost tracking rather than fabricate estimates.
575 """
577 def estimate_cost(
578 self,
579 model: str,
580 total_tokens: int,
581 provider: str | None = None,
582 prompt_tokens: int = 0,
583 completion_tokens: int = 0,
584 ) -> float:
585 """Estimate cost in USD for the given token usage.
587 When *prompt_tokens* and *completion_tokens* are provided they are
588 priced at the input and output rates respectively — this yields the
589 most accurate estimate. When both are ``0`` the caller has no
590 usage split, so implementations must fall back to *total_tokens*
591 using a documented approximation.
593 Args:
594 model: Model identifier (e.g. ``gpt-4o``).
595 total_tokens: Total tokens consumed (prompt + completion).
596 provider: Provider name (e.g. ``openai``) when pricing
597 differs per provider.
598 prompt_tokens: Input token count from ``Completion.usage``.
599 Defaults to ``0`` (unknown).
600 completion_tokens: Output token count from
601 ``Completion.usage``. Defaults to ``0`` (unknown).
603 Returns:
604 Estimated cost in USD. Return ``0.0`` when pricing is
605 unknown for the model.
606 """
607 ...
610__all__ = [
611 "AsyncStream",
612 "ChatMessage",
613 "ChatMessageProtocol",
614 "Completion",
615 "CompletionProtocol",
616 "CostEstimatorProtocol",
617 "EmbeddingClientProtocol",
618 "ExtractionError",
619 "FunctionCall",
620 "LLMClientProtocol",
621 "LLMError",
622 "PromptAssemblerProtocol",
623 "PromptOptimizerProtocol",
624 "PromptRegistryProtocol",
625 "PromptRendererProtocol",
626 "PromptTemplateProtocol",
627 "Role",
628 "SemanticCacheProtocol",
629 "StreamChunk",
630 "StreamEvent",
631 "StructuredExtractorProtocol",
632 "ThinkingConfig",
633 "ThinkingResult",
634 "TokenBudget",
635 "TokenCounterProtocol",
636 "TokenUsage",
637 "ToolCall",
638]