AI Paper Reading: Agentic Reasoning Implementation Guide 2026

"""
Core loop implementation of the ReAct pattern.

The LLM reasons in natural language, calls tools, and observes results
in a cycle that repeats up to max_steps.
"""
from dataclasses import dataclass, field
from typing import Callable, Optional
from enum import Enum
import json
import re


class StepType(Enum):
    THOUGHT = "thought"
    ACTION = "action"
    OBSERVATION = "observation"
    FINAL_ANSWER = "final_answer"


@dataclass
class AgentStep:
    step_type: StepType
    content: str
    tool_name: Optional[str] = None
    tool_input: Optional[dict] = None
    token_count: int = 0


@dataclass
class AgentTrace:
    """Complete record of agent execution."""
    question: str
    steps: list[AgentStep] = field(default_factory=list)
    final_answer: Optional[str] = None
    total_tokens: int = 0
    total_tool_calls: int = 0

    def add_step(self, step: AgentStep):
        self.steps.append(step)
        self.total_tokens += step.token_count
        if step.step_type == StepType.ACTION:
            self.total_tool_calls += 1


class ReActAgent:
    """ReAct pattern agent.

    Takes an LLM and a set of tools, and performs iterative
    reasoning-action-observation loops for a given question.
    """

    SYSTEM_PROMPT = """You are a helpful assistant that solves problems step by step.
For each step, you MUST output exactly one of:
- Thought: <your reasoning about what to do next>
- Action: <tool_name>({"param": "value"})
- Final Answer: <your final response to the user>

Available tools:
{tool_descriptions}

Rules:
- Always think before acting.
- After observing a tool result, think about what it means before the next action.
- When you have enough information, provide Final Answer.
"""

    def __init__(
        self,
        llm: Callable,   # (messages: list[dict]) -> str
        tools: dict[str, Callable],
        tool_descriptions: dict[str, str],
        max_steps: int = 10,
        max_tokens_per_step: int = 1024,
    ):
        self.llm = llm
        self.tools = tools
        self.tool_descriptions = tool_descriptions
        self.max_steps = max_steps
        self.max_tokens_per_step = max_tokens_per_step

    def run(self, question: str) -> AgentTrace:
        trace = AgentTrace(question=question)

        # Insert tool descriptions into the system prompt
        tool_desc_text = "\n".join(
            f"- {name}: {desc}"
            for name, desc in self.tool_descriptions.items()
        )
        system_msg = self.SYSTEM_PROMPT.format(tool_descriptions=tool_desc_text)

        messages = [
            {"role": "system", "content": system_msg},
            {"role": "user", "content": question},
        ]

        for step_num in range(self.max_steps):
            # Ask the LLM to generate the next step
            response = self.llm(messages)
            parsed = self._parse_response(response)

            if parsed.step_type == StepType.FINAL_ANSWER:
                trace.final_answer = parsed.content
                trace.add_step(parsed)
                break

            trace.add_step(parsed)
            messages.append({"role": "assistant", "content": response})

            if parsed.step_type == StepType.ACTION and parsed.tool_name:
                # Execute the tool
                observation = self._execute_tool(
                    parsed.tool_name, parsed.tool_input or {}
                )
                obs_step = AgentStep(
                    step_type=StepType.OBSERVATION,
                    content=observation,
                )
                trace.add_step(obs_step)
                messages.append({
                    "role": "user",
                    "content": f"Observation: {observation}",
                })

        return trace

    def _parse_response(self, response: str) -> AgentStep:
        """Parse LLM output to distinguish between Thought/Action/Final Answer."""
        response = response.strip()

        # Check for Final Answer
        if response.lower().startswith("final answer:"):
            return AgentStep(
                step_type=StepType.FINAL_ANSWER,
                content=response[len("final answer:"):].strip(),
            )

        # Parse Action: Action: tool_name({"key": "value"})
        action_match = re.match(
            r'Action:\s*(\w+)\((\{.*\})\)', response, re.DOTALL
        )
        if action_match:
            tool_name = action_match.group(1)
            try:
                tool_input = json.loads(action_match.group(2))
            except json.JSONDecodeError:
                tool_input = {}
            return AgentStep(
                step_type=StepType.ACTION,
                content=response,
                tool_name=tool_name,
                tool_input=tool_input,
            )

        # Everything else is treated as Thought
        return AgentStep(
            step_type=StepType.THOUGHT,
            content=response,
        )

    def _execute_tool(self, tool_name: str, tool_input: dict) -> str:
        """Execute a tool and return the result as a string."""
        if tool_name not in self.tools:
            return f"Error: Unknown tool '{tool_name}'. Available: {list(self.tools.keys())}"
        try:
            result = self.tools[tool_name](**tool_input)
            return str(result)
        except Exception as e:
            return f"Error executing {tool_name}: {type(e).__name__}: {str(e)}"

"""
Agent tool definitions for production environments.

Each tool has built-in input validation, timeout, and cost limits,
and returns results in a structured format.
"""
from dataclasses import dataclass
from typing import Any, Optional
import httpx
import time


@dataclass
class ToolResult:
    success: bool
    data: Any
    error: Optional[str] = None
    execution_time_ms: float = 0.0
    cost_usd: float = 0.0


class WebSearchTool:
    """Web search tool.

    Used when the agent needs to look up the latest information.
    Has built-in rate limiting and cost controls.
    """

    def __init__(
        self,
        api_key: str,
        max_results: int = 5,
        timeout_seconds: float = 10.0,
        max_calls_per_minute: int = 10,
    ):
        self.api_key = api_key
        self.max_results = max_results
        self.timeout_seconds = timeout_seconds
        self.max_calls_per_minute = max_calls_per_minute
        self._call_timestamps: list[float] = []

    def _check_rate_limit(self) -> bool:
        now = time.time()
        self._call_timestamps = [
            ts for ts in self._call_timestamps if now - ts < 60
        ]
        return len(self._call_timestamps) < self.max_calls_per_minute

    def __call__(self, query: str) -> ToolResult:
        if not query or len(query) > 500:
            return ToolResult(
                success=False,
                data=None,
                error="Query must be 1-500 characters",
            )

        if not self._check_rate_limit():
            return ToolResult(
                success=False,
                data=None,
                error=f"Rate limit exceeded: max {self.max_calls_per_minute}/min",
            )

        start = time.monotonic()
        try:
            # Actual search API call (e.g., Tavily, Serper, etc.)
            with httpx.Client(timeout=self.timeout_seconds) as client:
                response = client.get(
                    "https://api.search-provider.com/search",
                    params={"q": query, "max_results": self.max_results},
                    headers={"Authorization": f"Bearer {self.api_key}"},
                )
                response.raise_for_status()
                elapsed = (time.monotonic() - start) * 1000
                self._call_timestamps.append(time.time())

                return ToolResult(
                    success=True,
                    data=response.json(),
                    execution_time_ms=elapsed,
                    cost_usd=0.001,  # Estimated cost per call
                )
        except httpx.TimeoutException:
            return ToolResult(
                success=False,
                data=None,
                error=f"Search timed out after {self.timeout_seconds}s",
                execution_time_ms=(time.monotonic() - start) * 1000,
            )
        except httpx.HTTPStatusError as e:
            return ToolResult(
                success=False,
                data=None,
                error=f"HTTP {e.response.status_code}: {e.response.text[:200]}",
                execution_time_ms=(time.monotonic() - start) * 1000,
            )


class CodeExecutionTool:
    """Code execution tool.

    The agent executes Python code for calculations or data processing.
    For security, only allowed modules can be imported, and execution time
    and memory are limited.
    """

    ALLOWED_MODULES = {"math", "statistics", "json", "re", "datetime", "collections"}

    def __init__(self, timeout_seconds: float = 5.0):
        self.timeout_seconds = timeout_seconds

    def __call__(self, code: str) -> ToolResult:
        if not code or len(code) > 5000:
            return ToolResult(
                success=False,
                data=None,
                error="Code must be 1-5000 characters",
            )

        # Import check: only allowed modules can be used
        import_lines = [
            line.strip() for line in code.splitlines()
            if line.strip().startswith("import ") or line.strip().startswith("from ")
        ]
        for line in import_lines:
            module = line.split()[1].split(".")[0]
            if module not in self.ALLOWED_MODULES:
                return ToolResult(
                    success=False,
                    data=None,
                    error=f"Module '{module}' not allowed. Allowed: {self.ALLOWED_MODULES}",
                )

        start = time.monotonic()
        try:
            # Execute in a restricted environment
            local_vars: dict = {}
            exec(code, {"__builtins__": {}}, local_vars)  # noqa: S102

            elapsed = (time.monotonic() - start) * 1000
            # Return 'result' variable if it exists
            result = local_vars.get("result", str(local_vars))

            return ToolResult(
                success=True,
                data=result,
                execution_time_ms=elapsed,
            )
        except Exception as e:
            return ToolResult(
                success=False,
                data=None,
                error=f"{type(e).__name__}: {str(e)}",
                execution_time_ms=(time.monotonic() - start) * 1000,
            )

"""
Agent working memory management.

Maintains the full conversation history, but when passing to the LLM,
summarizes/selects based on importance to fit within the context window.
"""
from dataclasses import dataclass, field
from typing import Optional
import hashlib


@dataclass
class MemoryEntry:
    role: str
    content: str
    step_number: int
    importance: float = 0.5  # 0.0 ~ 1.0
    token_count: int = 0
    content_hash: str = ""

    def __post_init__(self):
        if not self.content_hash:
            self.content_hash = hashlib.md5(
                self.content.encode()
            ).hexdigest()[:8]


class SlidingWindowMemory:
    """Sliding window + importance-based memory management.

    Always keeps the most recent K messages,
    and selects older messages based on importance scores.
    """

    def __init__(
        self,
        max_tokens: int = 8192,
        recent_window: int = 6,     # Always keep the most recent N
        system_prompt_tokens: int = 500,
    ):
        self.max_tokens = max_tokens
        self.recent_window = recent_window
        self.system_prompt_tokens = system_prompt_tokens
        self.entries: list[MemoryEntry] = []

    def add(self, entry: MemoryEntry):
        # Prevent duplicates
        if any(e.content_hash == entry.content_hash for e in self.entries):
            return
        self.entries.append(entry)

    def get_context(self, system_message: str) -> list[dict]:
        """Construct the message list to pass to the LLM.

        1. System prompt is always included
        2. Most recent recent_window messages are always included
        3. The rest are included within budget based on importance
        """
        budget = self.max_tokens - self.system_prompt_tokens
        messages = [{"role": "system", "content": system_message}]

        if not self.entries:
            return messages

        # Secure recent messages first
        recent = self.entries[-self.recent_window:]
        older = self.entries[:-self.recent_window] if len(self.entries) > self.recent_window else []

        recent_tokens = sum(e.token_count for e in recent)

        # Add important older messages within budget
        remaining_budget = budget - recent_tokens
        selected_older = sorted(older, key=lambda e: e.importance, reverse=True)

        included_older = []
        for entry in selected_older:
            if remaining_budget <= 0:
                break
            if entry.token_count <= remaining_budget:
                included_older.append(entry)
                remaining_budget -= entry.token_count

        # Sort chronologically and compose messages
        included_older.sort(key=lambda e: e.step_number)
        all_entries = included_older + recent

        for entry in all_entries:
            messages.append({"role": entry.role, "content": entry.content})

        return messages

    def mark_important(self, step_number: int, importance: float = 1.0):
        """Increase the importance of a specific step.

        Used to mark tool execution results, key findings, etc.
        """
        for entry in self.entries:
            if entry.step_number == step_number:
                entry.importance = importance
                break

"""
Guardrails for controlling the cost and resource usage of agent execution.
"""
from dataclasses import dataclass
from typing import Optional
import time


@dataclass
class AgentBudget:
    max_llm_calls: int = 15
    max_tool_calls: int = 10
    max_total_tokens: int = 50_000
    max_cost_usd: float = 0.50
    max_wall_time_seconds: float = 120.0


@dataclass
class AgentUsage:
    llm_calls: int = 0
    tool_calls: int = 0
    total_tokens: int = 0
    total_cost_usd: float = 0.0
    start_time: float = 0.0

    def elapsed_seconds(self) -> float:
        return time.time() - self.start_time if self.start_time else 0.0


class BudgetGuard:
    """Agent execution budget monitor.

    Calls check() before each step to verify whether the budget
    has been exceeded. If exceeded, the agent should terminate
    early with the results gathered so far.
    """

    def __init__(self, budget: AgentBudget):
        self.budget = budget
        self.usage = AgentUsage()

    def start(self):
        self.usage.start_time = time.time()

    def record_llm_call(self, tokens: int, cost_usd: float):
        self.usage.llm_calls += 1
        self.usage.total_tokens += tokens
        self.usage.total_cost_usd += cost_usd

    def record_tool_call(self, cost_usd: float = 0.0):
        self.usage.tool_calls += 1
        self.usage.total_cost_usd += cost_usd

    def check(self) -> Optional[str]:
        """Returns the reason if budget is exceeded. Returns None if within budget."""
        if self.usage.llm_calls >= self.budget.max_llm_calls:
            return f"LLM call limit reached: {self.usage.llm_calls}/{self.budget.max_llm_calls}"

        if self.usage.tool_calls >= self.budget.max_tool_calls:
            return f"Tool call limit reached: {self.usage.tool_calls}/{self.budget.max_tool_calls}"

        if self.usage.total_tokens >= self.budget.max_total_tokens:
            return f"Token limit reached: {self.usage.total_tokens}/{self.budget.max_total_tokens}"

        if self.usage.total_cost_usd >= self.budget.max_cost_usd:
            return f"Cost limit reached: ${self.usage.total_cost_usd:.3f}/${self.budget.max_cost_usd:.3f}"

        elapsed = self.usage.elapsed_seconds()
        if elapsed >= self.budget.max_wall_time_seconds:
            return f"Time limit reached: {elapsed:.1f}s/{self.budget.max_wall_time_seconds}s"

        return None

    def summary(self) -> dict:
        return {
            "llm_calls": f"{self.usage.llm_calls}/{self.budget.max_llm_calls}",
            "tool_calls": f"{self.usage.tool_calls}/{self.budget.max_tool_calls}",
            "tokens": f"{self.usage.total_tokens}/{self.budget.max_total_tokens}",
            "cost_usd": f"${self.usage.total_cost_usd:.4f}/${self.budget.max_cost_usd:.4f}",
            "elapsed_s": f"{self.usage.elapsed_seconds():.1f}/{self.budget.max_wall_time_seconds}",
        }

"""
Agent evaluation framework.

Comprehensively measures efficiency, tool usage appropriateness,
and reasoning quality in addition to accuracy.
"""
from dataclasses import dataclass


@dataclass
class AgentEvalMetrics:
    # Accuracy
    final_answer_correct: bool
    partial_credit: float          # 0.0 ~ 1.0 (partial score)

    # Efficiency
    total_steps: int
    total_tool_calls: int
    total_tokens: int
    total_cost_usd: float
    wall_time_seconds: float

    # Tool usage quality
    unnecessary_tool_calls: int    # Number of unnecessary tool calls
    failed_tool_calls: int         # Number of failed tool calls
    tool_call_accuracy: float      # Rate of calling the right tool with the right input

    # Reasoning quality
    reasoning_coherence: float     # Logical consistency of reasoning (0.0 ~ 1.0)
    hallucination_count: int       # Number of unsupported claims

    @property
    def efficiency_score(self) -> float:
        """Efficiency score: how few resources were used to reach the correct answer."""
        if not self.final_answer_correct:
            return 0.0
        # Lower is more efficient -> convert via inverse
        step_penalty = min(self.total_steps / 10, 1.0)
        cost_penalty = min(self.total_cost_usd / 0.10, 1.0)
        return max(0.0, 1.0 - (step_penalty + cost_penalty) / 2)

    @property
    def overall_score(self) -> float:
        """Overall score."""
        weights = {
            "accuracy": 0.4,
            "efficiency": 0.2,
            "tool_quality": 0.2,
            "reasoning": 0.2,
        }
        accuracy = 1.0 if self.final_answer_correct else self.partial_credit
        return (
            weights["accuracy"] * accuracy
            + weights["efficiency"] * self.efficiency_score
            + weights["tool_quality"] * self.tool_call_accuracy
            + weights["reasoning"] * self.reasoning_coherence
        )