Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/reasoning/iterative.py: 87%
86 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1from __future__ import annotations
3from datetime import UTC, datetime
4from typing import Any
6from lexigram.ai.rag.reasoning.base import (
7 AbstractReasoner,
8 ReasoningResult,
9 ReasoningStep,
10 ReasoningStrategy,
11)
12from lexigram.contracts import (
13 ChatMessage,
14 LLMClientProtocol,
15)
16from lexigram.contracts.data.vector.protocols import VectorCollectionProtocol
19class IterativeRefinementReasoner(AbstractReasoner):
20 """Iteratively refine answers through multiple passes."""
22 def __init__(
23 self,
24 llm_client: LLMClientProtocol,
25 vector_store: VectorCollectionProtocol,
26 max_iterations: int = 3,
27 top_k: int = 5,
28 temperature: float = 0.5,
29 ):
30 """Initialize iterative refinement reasoner."""
31 self.llm_client = llm_client
32 self.vector_store = vector_store
33 self.max_iterations = max_iterations
34 self.top_k = top_k
35 self.temperature = temperature
37 async def reason(
38 self,
39 query: str,
40 initial_context: list[Any] | None = None,
41 **kwargs,
42 ) -> ReasoningResult:
43 """Perform iterative refinement reasoning."""
44 steps: list[ReasoningStep] = []
45 current_answer = ""
47 # Initial retrieval
48 retrieved_docs = await self.vector_store.search(query=query, limit=self.top_k) # type: ignore[call-arg,arg-type]
50 # Extract context
51 context_texts: list[str] = []
52 for doc in retrieved_docs:
53 if hasattr(doc, "content"):
54 text = doc.content
55 if text is not None:
56 context_texts.append(text)
57 elif isinstance(doc, dict) and "content" in doc:
58 context_texts.append(doc["content"])
59 elif isinstance(doc, str):
60 context_texts.append(doc)
62 # Add initial context if provided
63 if initial_context:
64 for doc in initial_context:
65 if hasattr(doc, "content"):
66 text = doc.content
67 if text is not None:
68 context_texts.append(text)
69 elif isinstance(doc, dict) and "content" in doc:
70 context_texts.append(doc["content"])
71 elif isinstance(doc, str):
72 context_texts.append(doc)
74 for iteration in range(1, self.max_iterations + 1):
75 if iteration == 1:
76 # Initial answer
77 answer = await self._generate_initial_answer(query, context_texts)
78 reasoning = "Initial answer generation"
79 else:
80 # Refine previous answer
81 answer, critique = await self._refine_answer(
82 query,
83 current_answer,
84 context_texts,
85 )
86 reasoning = f"Refinement iteration {iteration}: {critique}"
88 step = ReasoningStep(
89 step_number=iteration,
90 question=f"Iteration {iteration}",
91 context=retrieved_docs if iteration == 1 else [],
92 reasoning=reasoning,
93 answer=answer,
94 confidence=min(0.5 + (iteration * 0.15), 0.95),
95 metadata={"iteration": iteration},
96 )
97 steps.append(step)
98 current_answer = answer
100 return ReasoningResult(
101 query=query,
102 final_answer=current_answer,
103 steps=steps,
104 strategy=ReasoningStrategy.ITERATIVE_REFINEMENT,
105 total_hops=len(steps),
106 overall_confidence=steps[-1].confidence if steps else 0.0,
107 metadata={
108 "max_iterations": self.max_iterations,
109 "timestamp": datetime.now(UTC).isoformat(),
110 },
111 )
113 async def _generate_initial_answer(
114 self,
115 query: str,
116 context: list[str],
117 ) -> str:
118 """Generate initial answer."""
119 context_str = "\n\n".join(f"[{i + 1}] {ctx}" for i, ctx in enumerate(context))
121 prompt = f"""Context:
122{context_str}
124Question: {query}
126Provide a comprehensive answer based on the context:"""
128 result = await self.llm_client.complete(
129 messages=[
130 ChatMessage(role="system", content="You are a helpful assistant."),
131 ChatMessage(role="user", content=prompt),
132 ],
133 temperature=self.temperature,
134 max_tokens=500,
135 )
136 if result.is_err():
137 raise result.unwrap_err()
138 response = result.unwrap()
140 # Extract text
141 if hasattr(response, "content"):
142 return response.content
143 if hasattr(response, "choices") and response.choices:
144 return response.choices[0].message.content
145 if isinstance(response, dict) and "content" in response:
146 return response["content"]
147 return str(response)
149 async def _refine_answer(
150 self,
151 query: str,
152 previous_answer: str,
153 context: list[str],
154 ) -> tuple[str, str]:
155 """Refine previous answer."""
156 context_str = "\n\n".join(f"[{i + 1}] {ctx}" for i, ctx in enumerate(context))
158 prompt = f"""Question: {query}
160Previous Answer:
161{previous_answer}
163Context:
164{context_str}
166Please refine the previous answer by:
1671. Correcting any inaccuracies
1682. Adding missing important information from the context
1693. Improving clarity and organization
171Provide your critique and refined answer:"""
173 result = await self.llm_client.complete(
174 messages=[
175 ChatMessage(
176 role="system",
177 content="You are a helpful assistant that improves and refines answers.",
178 ),
179 ChatMessage(role="user", content=prompt),
180 ],
181 temperature=self.temperature,
182 max_tokens=600,
183 )
184 if result.is_err():
185 raise result.unwrap_err()
186 response = result.unwrap()
188 # Extract text
189 if hasattr(response, "content"):
190 text = response.content
191 elif hasattr(response, "choices") and response.choices:
192 text = response.choices[0].message.content
193 elif isinstance(response, dict) and "content" in response:
194 text = response["content"]
195 else:
196 text = str(response)
198 # Try to separate critique and answer
199 if "Refined Answer:" in text:
200 parts = text.split("Refined Answer:", 1)
201 critique = parts[0].strip()
202 answer = parts[1].strip()
203 elif "Answer:" in text:
204 parts = text.split("Answer:", 1)
205 critique = parts[0].strip()
206 answer = parts[1].strip()
207 else:
208 critique = "General refinement"
209 answer = text
211 return answer, critique