Agent Architectures: Control Loops, State & Planning › Code Implementation

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional


@dataclass
class Observation:
    """Input received from the environment at each iteration."""

    content: str
    source: str  # 'user', 'tool', or 'system'
    step: int


@dataclass
class Action:
    """The agent's decision about what to do next."""

    tool_name: Optional[str]  # None signals a final answer (termination)
    tool_input: Dict[str, Any]
    reasoning: str


@dataclass
class AgentStep:
    """A complete record of one agent loop iteration."""

    observation: Observation
    thought: str
    action: Action
    result: Optional[str] = None


@dataclass
class AgentState:
    """The full state of the agent across all iterations."""

    steps: List[AgentStep] = field(default_factory=list)
    working_memory: Dict[str, Any] = field(default_factory=dict)
    current_plan: List[str] = field(default_factory=list)
    iteration_count: int = 0

import operator


def calculator(expression: str) -> str:
    """
    Safe calculator that evaluates simple arithmetic expressions.
    Supports addition, subtraction, multiplication, and division.
    Parses operations explicitly to avoid any code execution concerns.
    """
    # Remove all whitespace
    expression = expression.replace(" ", "")

    # Try to match a simple two-operand expression
    pattern = re.compile(r"^(-?\d+(?:\.\d+)?)([\+\-\*/])(-?\d+(?:\.\d+)?)$")
    match = pattern.match(expression)
    if not match:
        return "Error: Expression must be in the form 'a op b' (e.g. '100 / 4')"

    left = float(match.group(1))
    op_char = match.group(2)
    right = float(match.group(3))

    ops = {
        "+": operator.add,
        "-": operator.sub,
        "*": operator.mul,
        "/": operator.truediv,
    }

    if op_char not in ops:
        return f"Error: Unsupported operator '{op_char}'"

    if op_char == "/" and right == 0:
        return "Error: Division by zero"

    result = ops[op_char](left, right)
    # Return integer if result is whole number
    if result == int(result):
        return str(int(result))
    return str(round(result, 6))


# Tool registry: maps tool names to their functions and schemas
TOOLS = {
    "calculator": {
        "function": calculator,
        "description": "Calculate simple two-operand arithmetic. Input: expression like '150 - 45'.",
        "parameters": {
            "type": "object",
            "properties": {
                "expression": {
                    "type": "string",
                    "description": "Arithmetic expression to evaluate, e.g. '150 - 45'",
                }
            },
            "required": ["expression"],
        },
    }
}

import re


class SimpleAgent:
    """
    Agent architecture showing loop, state, planning, and termination.
    Uses rule-based reasoning to simulate LLM decision making.
    """

    def __init__(self, max_iterations: int = 5, tools: Dict = None):
        self.max_iterations = max_iterations
        self.tools = tools or {}
        self.state = AgentState()

    def plan(self, observation: Observation) -> List[str]:
        """
        Generate a task plan based on the observation content.
        A real agent would call an LLM here.
        """
        content = observation.content.lower()
        if "calculate" in content or any(
            op in content for op in ["+", "-", "*", "/"]
        ):
            return ["extract_expression", "calculate", "verify"]
        return ["analyze", "respond"]

    def _extract_expression(self, content: str) -> Optional[str]:
        """Extract a simple arithmetic expression from the content string."""
        # Strip leading instruction words
        text = re.sub(r"calculate\s*", "", content, flags=re.IGNORECASE).strip()

        # Match patterns like "150 - 45" or "100/4"
        pattern = re.compile(
            r"(-?\d+(?:\.\d+)?)\s*([\+\-\*/])\s*(-?\d+(?:\.\d+)?)"
        )
        match = pattern.search(text)
        if match:
            return f"{match.group(1)}{match.group(2)}{match.group(3)}"

        # Detect word-form operations
        numbers = re.findall(r"\d+(?:\.\d+)?", content)
        if len(numbers) >= 2:
            if "minus" in content.lower() or "subtract" in content.lower():
                return f"{numbers[0]}-{numbers[1]}"
            if (
                "plus" in content.lower()
                or "total" in content.lower()
                or "sum" in content.lower()
            ):
                return f"{numbers[0]}+{numbers[1]}"
            if "divided" in content.lower():
                return f"{numbers[0]}/{numbers[1]}"
            if "times" in content.lower() or "multiplied" in content.lower():
                return f"{numbers[0]}*{numbers[1]}"

        return None

    def decide_action(self, observation: Observation) -> Action:
        """
        Choose the next action based on observation and current plan.
        A real agent would call an LLM here.
        """
        content = observation.content

        # Hard termination: iteration limit exceeded
        if self.state.iteration_count >= self.max_iterations:
            return Action(
                tool_name=None,
                tool_input={"answer": "Maximum iterations reached"},
                reasoning="Terminating due to iteration limit",
            )

        # Try to extract and calculate an arithmetic expression
        expr = self._extract_expression(content)
        if expr:
            return Action(
                tool_name="calculator",
                tool_input={"expression": expr},
                reasoning=f"Arithmetic expression detected: {expr}",
            )

        # Default: answer directly without a tool
        return Action(
            tool_name=None,
            tool_input={"answer": f"Processed: {content}"},
            reasoning="No arithmetic detected, giving direct response",
        )

    def should_terminate(self, action: Action) -> bool:
        """Determine whether the agent should stop after this action."""
        if action.tool_name is None:
            return True
        if self.state.iteration_count >= self.max_iterations:
            return True
        return False

    def run(self, query: str) -> Dict:
        """Execute the agent loop until a termination condition is met."""
        print(f"Starting agent with query: '{query}'\n")

        observation = Observation(content=query, source="user", step=0)

        while True:
            self.state.iteration_count += 1
            step_num = self.state.iteration_count

            print(f"Iteration {step_num}")
            print(f"  Observation: {observation.content[:70]}")

            # Planning phase: generate plan on first iteration
            if not self.state.current_plan:
                self.state.current_plan = self.plan(observation)
                print(f"  Plan: {' -> '.join(self.state.current_plan)}")

            # Reasoning phase: decide next action
            action = self.decide_action(observation)
            print(f"  Reasoning: {action.reasoning}")

            if action.tool_name:
                print(f"  Action: {action.tool_name}({action.tool_input})")
            else:
                print("  Action: Final answer")

            # Termination check
            if self.should_terminate(action):
                step = AgentStep(
                    observation=observation,
                    thought=action.reasoning,
                    action=action,
                    result=action.tool_input.get("answer", "Done"),
                )
                self.state.steps.append(step)
                print(f"\nTerminated after {step_num} iteration(s)")
                return {
                    "answer": action.tool_input.get("answer"),
                    "iterations": step_num,
                    "history": self.state.steps,
                }

            # Execution phase: call the tool
            tool_func = self.tools.get(action.tool_name, {}).get("function")
            if tool_func:
                try:
                    result = tool_func(**action.tool_input)
                    print(f"  Result: {result}")
                except Exception as e:
                    result = f"Error: {str(e)}"
                    print(f"  Result: {result}")
            else:
                result = f"Unknown tool: {action.tool_name}"
                print(f"  Result: {result}")

            # Update working memory with tool result
            self.state.working_memory[f"step_{step_num}_result"] = result

            # Record the completed step
            step = AgentStep(
                observation=observation,
                thought=action.reasoning,
                action=action,
                result=result,
            )
            self.state.steps.append(step)

            # The tool result becomes the next observation
            observation = Observation(
                content=result, source="tool", step=step_num
            )

            # Advance the plan
            if self.state.current_plan:
                self.state.current_plan.pop(0)

            print()

agent = SimpleAgent(max_iterations=3, tools=TOOLS)
result = agent.run(
    "If I have 150 apples and give away 45, how many remain? Calculate 150 - 45"
)

class StatefulAgent(SimpleAgent):
    """Extended agent with explicit state reporting capabilities."""

    def get_state_summary(self) -> str:
        """Generate a human-readable summary of the agent's current state."""
        lines = [
            f"Iterations completed: {self.state.iteration_count}",
            f"Working memory entries: {len(self.state.working_memory)}",
            f"Conversation length: {len(self.state.steps)} step(s)",
        ]

        if self.state.steps:
            last_step = self.state.steps[-1]
            lines.append(
                f"Last action: {last_step.action.tool_name or 'final_answer'}"
            )
            lines.append(f"Last result: {str(last_step.result)[:60]}")

        if self.state.working_memory:
            lines.append("Working memory contents:")
            for key, val in self.state.working_memory.items():
                lines.append(f"  {key}: {val}")

        return "\n".join(lines)


stateful_agent = StatefulAgent(max_iterations=3, tools=TOOLS)
stateful_result = stateful_agent.run("Calculate 100 / 4")