1"""Reflexion strategy — self-critique and iterative refinement for agent reasoning."""
2
3from __future__ import annotations
4
5import asyncio
6from typing import TYPE_CHECKING, Any, cast
7
8from lexigram.ai.agents.strategies.base import AbstractStrategy
9from lexigram.ai.agents.strategies.token_utils import TokenAccumulator
10from lexigram.ai.agents.types import ReasoningStep, ToolExecutionRecord
11from lexigram.contracts.ai.agents import AgentError, AgentResponse
12from lexigram.contracts.ai.llm import ChatMessage, Role
13from lexigram.logging import (
14 get_logger,
15)
16from lexigram.result import Err, Ok, Result
17
18if TYPE_CHECKING:
19 from lexigram.contracts.ai.agents import ToolProtocol
20 from lexigram.contracts.ai.llm import CompletionProtocol, LLMClientProtocol
21
22logger = get_logger(__name__)
23
24
25_CRITIQUE_SUFFIX = (
26 "\n\n---\nReview your previous response critically. "
27 "Identify any factual errors, logical inconsistencies, missing information, "
28 "or ways to be more helpful. "
29 "If the response is already optimal, write 'NO_CHANGES_NEEDED'. "
30 "Otherwise, describe the specific improvements needed in a numbered list."
31)
32
33_REFINE_PREFIX = (
34 "Based on your self-critique, provide an improved response. "
35 "Apply all identified improvements. Your refined response:\n\n"
36)
37
38
39class ReflexionStrategy(AbstractStrategy):
40 """Reflexion reasoning strategy with self-critique and iterative refinement.
41
42 The Reflexion strategy follows a three-phase loop:
43
44 1. **Generate** — produce an initial response to the user's request.
45 2. **Critique** — ask the same LLM to evaluate its own response,
46 identifying errors or improvements.
47 3. **Refine** — generate an improved response based on the critique.
48
49 The loop repeats up to ``max_iterations`` times, stopping early when
50 the critique declares the response optimal (``NO_CHANGES_NEEDED``).
51
52 Example:
53 >>> from lexigram.ai.agents import Agent
54 >>> from lexigram.ai.agents.strategies import ReflexionStrategy
55 >>>
56 >>> agent = Agent(
57 ... llm=my_llm_client,
58 ... strategy=ReflexionStrategy(max_iterations=3),
59 ... )
60 >>> response = await agent.run("Explain quantum entanglement")
61 """
62
63 def __init__(
64 self,
65 max_iterations: int = 3,
66 temperature_critique: float = 0.3,
67 temperature_refine: float = 0.5,
68 ) -> None:
69 """Initialise the ReflexionStrategy.
70
71 Args:
72 max_iterations: Maximum number of critique–refine cycles (default 3).
73 The strategy stops early when the critique returns
74 ``NO_CHANGES_NEEDED``.
75 temperature_critique: LLM temperature for the self-critique pass.
76 Lower values yield more deterministic critiques.
77 temperature_refine: LLM temperature for the refinement pass.
78 """
79 self.max_iterations = max_iterations
80 self.temperature_critique = temperature_critique
81 self.temperature_refine = temperature_refine
82
83 async def execute(
84 self,
85 message: str,
86 tools: list[ToolProtocol],
87 history: list[dict[str, Any]],
88 llm: LLMClientProtocol,
89 **kwargs: Any,
90 ) -> Result[AgentResponse, AgentError]:
91 """Execute the Reflexion reasoning loop.
92
93 Args:
94 message: The user's input message.
95 tools: Tools available to the agent (informational; reflexion does
96 not call tools directly).
97 history: Conversation history as list of message dicts.
98 llm: LLM client implementing
99 :class:`~lexigram.contracts.ai.LLMClientProtocol`.
100 **kwargs: Additional parameters:
101 - ``system_prompt`` (str): Optional system prompt.
102 - ``timeout`` (float): Per-call timeout in seconds.
103
104 Returns:
105 ``Ok(AgentResponse)`` with the final refined message and the full
106 reasoning trace (initial draft + critique + final response per
107 iteration).
108 ``Err(AgentError)`` if the LLM call fails or is cancelled.
109 """
110 system_prompt: str = kwargs.get("system_prompt", "")
111 timeout: float = kwargs.get("timeout", 60.0)
112
113 steps: list[ReasoningStep] = []
114 tool_calls: list[ToolExecutionRecord] = []
115 usage = TokenAccumulator()
116
117 # Build initial chat messages from history
118 messages: list[ChatMessage] = self._build_messages(
119 message, history, system_prompt
120 )
121
122 # Phase 1: initial generation
123 current_response = await self._call_llm(llm, messages, timeout=timeout)
124 if current_response is None:
125 return Err(
126 AgentError("LLM returned empty response during initial generation")
127 )
128 usage.add(current_response)
129 current_text = (
130 current_response.content
131 if hasattr(current_response, "content")
132 else str(current_response)
133 )
134
135 steps.append(
136 ReasoningStep(
137 step_number=0,
138 thought="[Reflexion] Initial response generated (iteration 0)",
139 action=None,
140 observation=current_text,
141 )
142 )
143
144 logger.debug(
145 "reflexion_initial_response",
146 length=len(current_text),
147 )
148
149 # Phase 2+: critique → refine loop
150 for iteration in range(1, self.max_iterations + 1):
151 # Self-critique pass
152 critique_messages = [
153 *messages,
154 ChatMessage(role=Role.ASSISTANT, content=current_text),
155 ChatMessage(role=Role.USER, content=_CRITIQUE_SUFFIX),
156 ]
157 critique = await self._call_llm(
158 llm,
159 critique_messages,
160 temperature=self.temperature_critique,
161 timeout=timeout,
162 )
163 if critique is None:
164 logger.warning("reflexion_critique_failed", iteration=iteration)
165 break
166 usage.add(critique)
167 critique_text = (
168 critique.content if hasattr(critique, "content") else str(critique)
169 )
170
171 steps.append(
172 ReasoningStep(
173 step_number=(iteration * 2) - 1,
174 thought=f"[Reflexion] Self-critique (iteration {iteration})",
175 action="critique",
176 observation=critique_text,
177 )
178 )
179
180 logger.debug(
181 "reflexion_critique",
182 iteration=iteration,
183 length=len(critique_text),
184 )
185
186 # Early stopping: LLM declared the response optimal
187 if "NO_CHANGES_NEEDED" in critique_text.upper():
188 logger.debug(
189 "reflexion_early_stop",
190 iteration=iteration,
191 reason="NO_CHANGES_NEEDED",
192 )
193 steps.append(
194 ReasoningStep(
195 step_number=iteration * 2,
196 thought=f"[Reflexion] No changes needed — stopping at iteration {iteration}",
197 action=None,
198 observation=current_text,
199 )
200 )
201 break
202
203 # Refinement pass
204 refine_messages = [
205 *critique_messages,
206 ChatMessage(role=Role.ASSISTANT, content=critique_text),
207 ChatMessage(role=Role.USER, content=_REFINE_PREFIX),
208 ]
209 refined = await self._call_llm(
210 llm,
211 refine_messages,
212 temperature=self.temperature_refine,
213 timeout=timeout,
214 )
215 if refined is None:
216 logger.warning("reflexion_refine_failed", iteration=iteration)
217 break
218 usage.add(refined)
219 current_text = (
220 refined.content if hasattr(refined, "content") else str(refined)
221 )
222
223 steps.append(
224 ReasoningStep(
225 step_number=iteration * 2,
226 thought=f"[Reflexion] Refined response (iteration {iteration})",
227 action="refine",
228 observation=current_text,
229 )
230 )
231
232 logger.debug(
233 "reflexion_refined",
234 iteration=iteration,
235 length=len(current_text),
236 )
237
238 return Ok(
239 AgentResponse(
240 message=current_text,
241 steps=steps,
242 tool_calls=tool_calls,
243 total_tokens=usage.total_tokens,
244 prompt_tokens=usage.prompt_tokens,
245 completion_tokens=usage.completion_tokens,
246 )
247 )
248
249 # ------------------------------------------------------------------
250 # Helpers
251 # ------------------------------------------------------------------
252
253 def _build_messages(
254 self,
255 message: str,
256 history: list[dict[str, Any]],
257 system_prompt: str,
258 ) -> list[ChatMessage]:
259 """Convert history + new message into a list of ChatMessage objects."""
260 messages: list[ChatMessage] = []
261
262 if system_prompt:
263 messages.append(ChatMessage(role=Role.SYSTEM, content=system_prompt))
264
265 for entry in history:
266 role_str = entry.get("role", "user")
267 content = entry.get("content", "")
268 try:
269 role = Role(role_str)
270 except ValueError:
271 role = Role.USER
272 messages.append(ChatMessage(role=role, content=content))
273
274 messages.append(ChatMessage(role=Role.USER, content=message))
275 return messages
276
277 async def _call_llm(
278 self,
279 llm: LLMClientProtocol,
280 messages: list[ChatMessage],
281 temperature: float | None = None,
282 timeout: float = 60.0,
283 ) -> CompletionProtocol | None:
284 """Call the LLM and return the completion, or ``None`` on failure."""
285 kwargs: dict[str, Any] = {}
286 if temperature is not None:
287 kwargs["temperature"] = temperature
288
289 try:
290 result = await asyncio.wait_for(
291 llm.complete(cast("list[Any]", messages), **kwargs),
292 timeout=timeout,
293 )
294 except TimeoutError:
295 logger.warning("reflexion_llm_timeout", timeout=timeout)
296 return None
297 except (OSError, ConnectionError, RuntimeError, ValueError) as exc:
298 logger.warning("reflexion_llm_error", error=str(exc))
299 return None
300
301 if not result.is_ok():
302 logger.warning("reflexion_llm_err_result", error=str(result.unwrap_err()))
303 return None
304
305 return result.unwrap()
306
307
308__all__ = ["ReflexionStrategy"]