8.2 The Agentic Loop: ReAct, Plan-Execute & Reflection¶

while not done:
    action = model.think(context)
    observation = environment.execute(action)
    context += (action, observation)

REACT_SYSTEM_PROMPT = """\
You are a helpful AI assistant that solves tasks by alternating between
Thought, Action, and Observation steps.

Format:
    Thought: <your reasoning about what to do next>
    Action: <tool_name>(<arguments>)
    Observation: <filled in by the system>
    ... (repeat as needed)
    Action: finish(<your final answer>)

Available tools:
{tool_descriptions}

Rules:
- Always think before acting.
- Actions must be exactly one tool call per step.
- When you have a final answer, call finish() — do not keep looping.
- If a tool returns an error, reason about the error and try a different approach.
"""

"""
minimal_react.py — A self-contained ReAct agent.

Requires:
    pip install openai>=1.0

Usage:
    OPENAI_API_KEY=sk-... python minimal_react.py
"""

import os
import re
import math
import json
import datetime
from typing import Callable

from openai import OpenAI

# ---------------------------------------------------------------------------
# 1. Tool registry
# ---------------------------------------------------------------------------

# Each tool is a plain Python function.  We store its name and a short
# description that goes into the system prompt.
TOOLS: dict[str, tuple[Callable, str]] = {}

def register_tool(description: str):
    """Decorator that registers a function as an agent tool."""
    def decorator(fn):
        TOOLS[fn.__name__] = (fn, description)
        return fn
    return decorator


@register_tool("Search the web for a query. Returns a short summary string.")
def search(query: str) -> str:
    # In production, call a real search API (e.g. Serper, Brave, Tavily).
    # Here we fake it for the purpose of the demo.
    fake_db = {
        "Tokyo population": "Tokyo metro ~37 million (2024).",
        "New York population": "NYC metro ~20 million (2024).",
        "speed of light": "299,792,458 metres per second.",
        "Python creator": "Guido van Rossum created Python.",
    }
    for key, value in fake_db.items():
        if key.lower() in query.lower():
            return value
    return f"No results found for: {query}"


@register_tool("Evaluate a Python arithmetic expression and return the result.")
def calculator(expression: str) -> str:
    # Restrict to safe arithmetic — no exec() on untrusted input.
    allowed = set("0123456789+-*/()., ")
    if not all(c in allowed for c in expression):
        return "Error: only arithmetic expressions allowed."
    try:
        result = eval(expression, {"__builtins__": {}}, {"math": math})
        return str(result)
    except Exception as e:
        return f"Error: {e}"


@register_tool("Return today's date.")
def get_date() -> str:
    return datetime.date.today().isoformat()


# ---------------------------------------------------------------------------
# 2. ReAct agent class
# ---------------------------------------------------------------------------

class ReActAgent:
    """
    A minimal ReAct agent that uses raw text generation (no function-calling
    API) so the Thought/Action/Observation format is fully visible.
    """

    MAX_STEPS = 10       # hard safety limit
    STOP_TOKEN = "finish"

    def __init__(self, model: str = "gpt-4o-mini"):
        self.client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
        self.model = model
        self._build_system_prompt()

    def _build_system_prompt(self) -> None:
        tool_descriptions = "\n".join(
            f"  - {name}({self._sig(fn)}): {desc}"
            for name, (fn, desc) in TOOLS.items()
        )
        self.system_prompt = f"""\
You are a helpful AI assistant that solves tasks step-by-step.
Use the following format EXACTLY:

Thought: <your reasoning>
Action: <tool_name>(<comma-separated arguments>)
...
(repeat Thought/Action until you have the answer)
Action: finish(<your final answer as a string>)

Available tools:
{tool_descriptions}

Important:
- Each Action line must contain exactly one tool call.
- Do NOT fabricate Observation lines — they will be injected for you.
- Stop immediately once you call finish().
"""

    @staticmethod
    def _sig(fn: Callable) -> str:
        """Return a compact signature string for a tool function."""
        import inspect
        sig = inspect.signature(fn)
        params = ", ".join(
            f"{p}: {v.annotation.__name__}"
            if v.annotation is not inspect.Parameter.empty
            else p
            for p, v in sig.parameters.items()
        )
        return params

    @staticmethod
    def _parse_action(text: str) -> tuple[str, list[str]] | None:
        """
        Extract the tool name and arguments from a line like:
            Action: search("Tokyo population")
        Returns (tool_name, [arg1, arg2, ...]) or None if no action found.
        """
        # Match lines starting with "Action:" then capture tool_name(args)
        match = re.search(r"Action:\s*(\w+)\(([^)]*)\)", text)
        if not match:
            return None
        tool_name = match.group(1)
        # Split args by comma and strip whitespace / surrounding quotes
        raw_args = match.group(2)
        args = [a.strip().strip('"\'') for a in raw_args.split(",") if a.strip()]
        return tool_name, args

    @staticmethod
    def _execute_action(tool_name: str, args: list[str]) -> str:
        """Call the registered tool and return its string output."""
        if tool_name == ReActAgent.STOP_TOKEN:
            # finish() is a pseudo-tool handled by the loop.
            return "__FINISH__"
        if tool_name not in TOOLS:
            return f"Error: unknown tool '{tool_name}'. Available: {list(TOOLS.keys())}"
        fn, _ = TOOLS[tool_name]
        try:
            result = fn(*args)
            return str(result)
        except TypeError as e:
            return f"Error calling {tool_name}: {e}"

    def run(self, task: str, verbose: bool = True) -> str:
        """
        Run the ReAct loop for a given task.
        Returns the final answer string.
        """
        # Conversation history: system + user task, then alternating
        # assistant (thought+action) and user (observation) messages.
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user",   "content": task},
        ]

        for step in range(self.MAX_STEPS):
            # ---- Model generation ----
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                # Stop generation right after "Action: finish(...)"
                # We use a stop sequence to prevent the model generating
                # a fake Observation line.
                stop=["Observation:"],
                temperature=0.0,   # deterministic for agentic tasks
                max_tokens=512,
            )
            assistant_text = response.choices[0].message.content.strip()

            if verbose:
                print(f"\n[Step {step+1}]\n{assistant_text}")

            # ---- Parse action ----
            parsed = self._parse_action(assistant_text)
            if parsed is None:
                # Model produced text but no parseable Action — inject an
                # error observation to get it back on track.
                observation = "Error: no Action found. Please follow the format."
            else:
                tool_name, args = parsed
                observation = self._execute_action(tool_name, args)

            # ---- Check for finish ----
            if observation == "__FINISH__":
                # Extract the argument to finish() as the final answer.
                match = re.search(r"finish\(([^)]*)\)", assistant_text)
                final_answer = match.group(1).strip('"\'') if match else assistant_text
                if verbose:
                    print(f"\n[Final Answer] {final_answer}")
                return final_answer

            # ---- Append to context ----
            # Assistant turn includes the thought + action.
            messages.append({"role": "assistant", "content": assistant_text})
            # User turn carries the environment observation.
            messages.append({
                "role": "user",
                "content": f"Observation: {observation}"
            })

            if verbose:
                print(f"Observation: {observation}")

        # Exceeded MAX_STEPS
        return "Error: agent exceeded maximum steps without finishing."


# ---------------------------------------------------------------------------
# 3. Demo
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    agent = ReActAgent(model="gpt-4o-mini")
    answer = agent.run(
        "How many times larger is Tokyo's metro population than New York's? "
        "Give me the ratio to 2 decimal places."
    )
    print(f"\nAnswer: {answer}")

def tool_schemas() -> list[dict]:
    """Convert the TOOLS registry into OpenAI-style JSON tool schemas."""
    import inspect
    schemas = []
    for name, (fn, desc) in TOOLS.items():
        params = {p: {"type": "string"}          # all our demo tools take strings
                  for p in inspect.signature(fn).parameters}
        schemas.append({
            "type": "function",
            "function": {
                "name": name,
                "description": desc,
                "parameters": {
                    "type": "object",
                    "properties": params,
                    "required": list(params),
                },
            },
        })
    return schemas


def run_tool_calling(client, model: str, task: str, max_steps: int = 10) -> str:
    """The same ReAct loop, but with structured tool calls instead of text parsing."""
    messages: list[dict] = [{"role": "user", "content": task}]

    for _ in range(max_steps):
        resp = client.chat.completions.create(
            model=model,
            messages=messages,
            tools=tool_schemas(),
            tool_choice="auto",
            temperature=0.0,
        )
        msg = resp.choices[0].message
        # Echo the assistant turn back verbatim — it carries the tool_call ids
        # (and, for reasoning models, the reasoning block the next turn needs).
        messages.append(msg.model_dump(exclude_none=True))

        if not msg.tool_calls:
            return msg.content or ""      # no tool call => this is the final answer

        for call in msg.tool_calls:
            args = json.loads(call.function.arguments)   # structured: no regex, no
            fn, _ = TOOLS[call.function.name]            # quoting or paren bugs
            try:
                result = str(fn(**args))
            except Exception as e:        # tool errors are observations, not crashes
                result = f"Error: {e}"
            messages.append({"role": "tool",
                             "tool_call_id": call.id,
                             "content": result})

    return "Error: exceeded max steps."

# pip install langgraph langchain-openai
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import InMemorySaver

@tool
def calculator(expression: str) -> str:
    """Evaluate a simple arithmetic expression."""   # docstring == tool description
    allowed = set("0123456789+-*/()., ")
    if not all(c in allowed for c in expression):
        return "Error: only arithmetic expressions allowed."
    return str(eval(expression, {"__builtins__": {}}, {}))

agent = create_react_agent(
    ChatOpenAI(model="gpt-4o-mini", temperature=0),
    tools=[calculator],
    checkpointer=InMemorySaver(),   # durable state: a run can be paused and resumed
)

result = agent.invoke(
    {"messages": [{"role": "user", "content": "What is 37 / 20?"}]},
    config={"configurable": {"thread_id": "demo-1"},
            "recursion_limit": 25},   # LangGraph's equivalent of MAX_STEPS
)
print(result["messages"][-1].content)

# Serve your own checkpoint; the loop above then needs only a different base_url.
vllm serve ./stack-100m-sft --port 8000

client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")

PLAN_PROMPT = """\
Break the following task into an ordered list of subtasks.
Each subtask should be atomic enough to be completed in a single
focused ReAct session. Output valid JSON only.

Task: {task}

Output format:
{{
  "goal": "<restate the goal>",
  "steps": [
    {{"id": 1, "description": "<step>", "depends_on": []}},
    {{"id": 2, "description": "<step>", "depends_on": [1]}},
    ...
  ]
}}
"""

import json

def make_plan(client, model: str, task: str) -> dict:
    """Ask the LLM to decompose a task into steps."""
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": PLAN_PROMPT.format(task=task)}],
        response_format={"type": "json_object"},  # JSON mode
        temperature=0.0,
    )
    return json.loads(response.choices[0].message.content)


def execute_plan(agent: ReActAgent, plan: dict) -> dict[int, str]:
    """
    Execute each step in dependency order.
    Returns a dict mapping step id → result string.
    """
    results: dict[int, str] = {}

    for step in plan["steps"]:
        step_id = step["id"]
        desc = step["description"]
        deps = step.get("depends_on", [])

        # Build a context-enriched task that incorporates prior results
        prior_context = "\n".join(
            f"Step {d} result: {results[d]}"
            for d in deps
            if d in results
        )
        full_task = f"{desc}\n\nContext from prior steps:\n{prior_context}" if prior_context else desc

        print(f"\n=== Executing step {step_id}: {desc} ===")
        results[step_id] = agent.run(full_task, verbose=True)

    return results

REFLECT_PROMPT = """\
The following agent trajectory failed to complete the task correctly.
Please analyze what went wrong and provide specific, actionable advice
for how the agent should approach this task differently next time.

Task: {task}

Trajectory:
{trajectory}

Error or failure reason: {failure_reason}

Reflection (be specific and concrete):
"""

def reflect(client, model: str, task: str, trajectory: str, failure: str) -> str:
    """Produce a verbal reflection on a failed trajectory."""
    response = client.chat.completions.create(
        model=model,
        messages=[{
            "role": "user",
            "content": REFLECT_PROMPT.format(
                task=task,
                trajectory=trajectory,
                failure_reason=failure,
            )
        }],
        temperature=0.0,
        max_tokens=256,
    )
    return response.choices[0].message.content.strip()

def critic_gate(client, model: str, context: str, proposed_action: str,
                threshold: float = 0.6) -> tuple[bool, str]:
    """
    Ask a critic whether the proposed action is correct/safe.
    Returns (should_proceed, explanation).
    """
    prompt = f"""\
Context so far:
{context}

Proposed action:
{proposed_action}

On a scale of 0 to 1, rate whether this action is correct, safe,
and aligned with the task goal. Then explain briefly.
Format: SCORE: <float>\\nEXPLANATION: <text>
"""
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.0, max_tokens=128,
    )
    text = response.choices[0].message.content
    match = re.search(r"SCORE:\s*([\d.]+)", text)
    score = float(match.group(1)) if match else 0.0
    return (score >= threshold), text

from dataclasses import dataclass, field
from copy import deepcopy

@dataclass
class AgentNode:
    """A node in the search tree."""
    messages: list         # conversation history up to this point
    score: float = 1.0     # cumulative log-prob or value estimate
    depth: int = 0
    actions: list = field(default_factory=list)  # actions taken so far

def beam_react_agent(
    client,
    model: str,
    task: str,
    beam_width: int = 3,
    max_depth: int = 8,
    value_fn=None,   # optional callable(node) -> float
) -> str:
    """
    A beam-search version of the ReAct loop.
    Returns the answer from the highest-scoring terminal node.
    """
    system = "You solve tasks step-by-step using Thought/Action format."
    root = AgentNode(
        messages=[
            {"role": "system", "content": system},
            {"role": "user",   "content": task},
        ]
    )
    beam = [root]

    for depth in range(max_depth):
        candidates = []
        for node in beam:
            # Sample multiple continuations from this node
            for _ in range(beam_width):
                resp = client.chat.completions.create(
                    model=model,
                    messages=node.messages,
                    stop=["Observation:"],
                    temperature=0.7,  # need diversity for beam search
                    max_tokens=256,
                )
                text = resp.choices[0].message.content.strip()

                # Check for finish
                if "finish(" in text:
                    match = re.search(r"finish\(([^)]*)\)", text)
                    answer = match.group(1).strip('"\'') if match else text
                    # Score this leaf
                    leaf_score = value_fn(node) if value_fn else node.score
                    candidates.append((leaf_score, answer, None))
                    continue

                # Execute action and get observation.  Both helpers are
                # @staticmethod, so we can reuse them without an agent instance.
                parsed = ReActAgent._parse_action(text)
                if parsed:
                    tool_name, args = parsed
                    obs = ReActAgent._execute_action(tool_name, args)
                else:
                    obs = "Error: no valid action found."

                # Compute node score (simple: add log-prob proxy from token count)
                new_msgs = deepcopy(node.messages)
                new_msgs.append({"role": "assistant", "content": text})
                new_msgs.append({"role": "user", "content": f"Observation: {obs}"})
                child_score = node.score * (0.95 ** 1)  # decay per step
                new_node = AgentNode(
                    messages=new_msgs,
                    score=child_score,
                    depth=depth + 1,
                    actions=node.actions + [text],
                )
                candidates.append((child_score, None, new_node))

        # Separate finished and unfinished candidates
        finished = [(s, a) for s, a, n in candidates if n is None]
        if finished:
            best_score, best_answer = max(finished, key=lambda x: x[0])
            return best_answer

        # Keep top beam_width nodes
        next_nodes = [(s, n) for s, _, n in candidates if n is not None]
        next_nodes.sort(key=lambda x: x[0], reverse=True)
        beam = [n for _, n in next_nodes[:beam_width]]

    # If we exhaust depth, return best partial result
    return "Exhausted search depth without finding answer."

MAX_CONTEXT_FRACTION = 0.85  # stop at 85% of context window

def context_full(messages: list, model_context_limit: int = 128_000) -> bool:
    """Rough check: count characters as a proxy for tokens (4 chars ≈ 1 token)."""
    total_chars = sum(len(m["content"]) for m in messages)
    estimated_tokens = total_chars / 4
    return estimated_tokens > MAX_CONTEXT_FRACTION * model_context_limit

from collections import Counter

def detect_loop(action_history: list[str], window: int = 4, threshold: int = 3) -> bool:
    """
    Returns True if the same action appears >= threshold times
    in the last 'window' steps.
    """
    recent = action_history[-window:]
    counts = Counter(recent)
    return any(c >= threshold for c in counts.values())