Coverage for src / lexigram / contracts / ai / rag.py: 0%
52 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-08-15 18:57 +0800
1"""RAG pipeline, retrieval, and reranking protocols."""
3from __future__ import annotations
5from dataclasses import dataclass
6from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
8from lexigram.contracts.ai.exceptions import RAGError
10if TYPE_CHECKING:
11 from lexigram.contracts.ai.vector import DocumentProtocol, SearchResultProtocol
12 from lexigram.contracts.core.result import Result
15# RAG Errors
16class RetrievalError(RAGError):
17 """Error raised during document retrieval."""
19 _code = "LEX_ERR_RAG_002"
22class SynthesisError(RAGError):
23 """Error raised during response synthesis."""
25 _code = "LEX_ERR_RAG_003"
28class ChunkingError(RAGError):
29 """Error raised during document chunking."""
31 _code = "LEX_ERR_RAG_004"
34@runtime_checkable
35class ChunkProtocol(Protocol):
36 """Structural protocol for document chunks."""
38 @property
39 def text(self) -> str:
40 """Chunk text content."""
41 ...
43 @property
44 def metadata(self) -> dict[str, Any]:
45 """Chunk metadata."""
46 ...
48 @property
49 def score(self) -> float | None:
50 """Optional relevance score."""
51 ...
54@dataclass(frozen=True)
55class RAGContext:
56 """Inputs to a RAG pipeline."""
58 query: str
59 config: dict[str, Any] | None = None
60 filters: dict[str, Any] | None = None
61 session_id: str | None = None
64@dataclass(frozen=True)
65class RAGResponse:
66 """Outputs from a RAG pipeline."""
68 answer: str
69 sources: list[SearchResultProtocol]
70 citations: list[Any] | None = None
71 confidence: float | None = None
74@runtime_checkable
75class DocumentLoaderProtocol(Protocol):
76 """Protocol for loading documents."""
78 async def load(self, source: str, **kwargs: Any) -> list[DocumentProtocol]:
79 """Load documents from a source."""
80 ...
83@runtime_checkable
84class SynthesizerProtocol(Protocol):
85 """Protocol for synthesizing a final answer from context."""
87 async def synthesize(
88 self,
89 query: str,
90 context: list[SearchResultProtocol],
91 **kwargs: Any,
92 ) -> Result[RAGResponse, RAGError]:
93 """Synthesize retrieved documents into an answer."""
94 ...
97@runtime_checkable
98class RAGPipelineProtocol(Protocol):
99 """Protocol for RAG pipeline execution.
101 Orchestrates retrieval, synthesis, and quality stages for a
102 given query context.
103 """
105 async def execute(self, context: RAGContext) -> Result[RAGResponse, RAGError]:
106 """Execute the full RAG pipeline for the given context.
108 Args:
109 context: Pipeline context containing query and config.
111 Returns:
112 Updated context with retrieved and synthesised response.
113 """
114 ...
117@runtime_checkable
118class RetrievalStrategyProtocol(Protocol):
119 """Protocol for pluggable RAG retrieval and ranking strategies.
121 Implementations take a query and a set of candidate documents and
122 return an ordered subset.
123 """
125 async def retrieve(
126 self,
127 query: str,
128 candidates: list[SearchResultProtocol],
129 *,
130 top_k: int = 5,
131 **kwargs: Any,
132 ) -> list[SearchResultProtocol]:
133 """Rank and return the top-k most relevant candidates.
135 Args:
136 query: Query string.
137 candidates: Retrieved candidate documents.
138 top_k: Maximum number of results to return.
139 **kwargs: Strategy-specific options.
141 Returns:
142 Ordered list of the most relevant documents.
143 """
144 ...
147@runtime_checkable
148class RerankingStrategyProtocol(Protocol):
149 """Protocol for cross-encoder or LLM-based reranking strategies.
151 Applied after initial retrieval to reorder documents by relevance.
152 """
154 async def rerank(
155 self,
156 query: str,
157 documents: list[SearchResultProtocol],
158 *,
159 top_k: int | None = None,
160 ) -> list[SearchResultProtocol]:
161 """Reorder documents by relevance to query.
163 Args:
164 query: Query string.
165 documents: Documents to rerank.
166 top_k: If set, return only the top-k results.
168 Returns:
169 Reranked (and optionally truncated) list of documents.
170 """
171 ...
174@runtime_checkable
175class RAGEvaluatorProtocol(Protocol):
176 """Protocol for evaluating RAG pipeline quality.
178 Implementations run metrics (faithfulness, relevance, etc.) on a
179 completed RAG interaction and return a structured report.
180 """
182 async def evaluate(
183 self,
184 query: str,
185 retrieved_docs: list[Any],
186 generated_answer: str,
187 **kwargs: Any,
188 ) -> Any:
189 """Evaluate a completed RAG interaction.
191 Args:
192 query: The original user query.
193 retrieved_docs: Documents retrieved by the pipeline.
194 generated_answer: The synthesized answer produced by the pipeline.
195 **kwargs: Additional evaluation parameters.
197 Returns:
198 An evaluation report (``RAGEvaluationReport`` or similar).
199 """
200 ...
203@runtime_checkable
204class PromptCompressorProtocol(Protocol):
205 """Compresses text to fit within a token budget.
207 Implementations range from learned compression (LLMLingua-2) to
208 heuristic truncation. The protocol guarantees that the returned text
209 fits within target_token_count.
211 Placement note: Lives in ai/rag.py because its primary consumer is the
212 RAG context compression stage. Also consumed by lexigram-ai-memory.
213 """
215 async def compress(
216 self,
217 text: str,
218 target_token_count: int,
219 force_tokens: list[str] | None = None,
220 ) -> str:
221 """Compress text to fit within target_token_count.
223 Args:
224 text: The raw text to compress.
225 target_token_count: Maximum tokens in the result.
226 force_tokens: Tokens that must never be removed.
228 Returns:
229 Compressed text fitting within the budget.
230 """
231 ...
234__all__ = [
235 "ChunkProtocol",
236 "ChunkingError",
237 "DocumentLoaderProtocol",
238 "PromptCompressorProtocol",
239 "RAGContext",
240 "RAGError",
241 "RAGEvaluatorProtocol",
242 "RAGPipelineProtocol",
243 "RAGResponse",
244 "RerankingStrategyProtocol",
245 "RetrievalError",
246 "RetrievalStrategyProtocol",
247 "SynthesisError",
248 "SynthesizerProtocol",
249]