More Than a Chatbot

User Request: "Fix the failing tests"
        ↓
┌─────────────────────────────────────────┐
│              AGENT LOOP                 │
│                                         │
│  1. LLM Plans                           │
│     ├── Analyze request                 │
│     ├── Consider context                │
│     └── Decide first action             │
│                                         │
│  2. Select Tool                         │
│     ├── Read test file                  │
│     ├── Run pytest                      │
│     └── Search for error                │
│                                         │
│  3. Execute Tool                        │
│     └── Tool runs in environment        │
│                                         │
│  4. Observe Result                      │
│     ├── Tool output                     │
│     ├── Error messages                  │
│     └── File changes                    │
│                                         │
│  5. Decide Next Action                  │
│     ├── More work needed? → Loop        │
│     └── Task complete? → Done           │
│                                         │
└─────────────────────────────────────────┘
        ↓
   Final Output: Code changes + explanation

"""Minimal AI Agent Loop Demo — Safe mock tools only."""

import json
from typing import Callable

# ── Mock Tools (safe, no real execution) ──────────────────

def read_file(path: str) -> str:
    """Mock: reads a file and returns contents."""
    mock_files = {
        "main.py": "def add(a, b):\n    return a + b",
        "test_main.py": "from main import add\n\ndef test_add():\n    assert add(1, 2) == 3",
    }
    return mock_files.get(path, f"Error: {path} not found")


def write_file(path: str, content: str) -> str:
    """Mock: writes content to a file."""
    return f"Written {len(content)} chars to {path}"


def run_command(cmd: str) -> str:
    """Mock: runs a terminal command."""
    if "pytest" in cmd:
        return "1 passed in 0.02s"
    return f"Executed: {cmd}"


# ── Tool Registry ─────────────────────────────────────────

TOOLS: dict[str, Callable] = {
    "read_file": read_file,
    "write_file": write_file,
    "run_command": run_command,
}


# ── Agent Loop ────────────────────────────────────────────

def agent_loop(user_request: str, max_steps: int = 10) -> list[dict]:
    """
    Simulate an AI agent loop:
    1. Plan (mock LLM)
    2. Select tool
    3. Execute
    4. Observe
    5. Repeat or finish
    """
    history: list[dict] = []
    context: list[str] = []

    for step in range(max_steps):
        print(f"\n── Step {step + 1} ──")

        # Step 1: Plan (mock LLM decision)
        plan = mock_llm_plan(user_request, context, history)
        print(f"Plan: {plan['thought']}")

        if plan.get("done"):
            print("Agent: Task complete!")
            break

        # Step 2: Select tool
        tool_name = plan["tool"]
        tool_args = plan["args"]
        print(f"Tool: {tool_name}({tool_args})")

        # Step 3: Execute tool
        tool_fn = TOOLS[tool_name]
        result = tool_fn(**tool_args)
        print(f"Result: {result}")

        # Step 4: Observe and record
        observation = {
            "step": step + 1,
            "thought": plan["thought"],
            "tool": tool_name,
            "args": tool_args,
            "result": result,
        }
        history.append(observation)
        context.append(f"{tool_name}({tool_args}) → {result}")

    return history


def mock_llm_plan(request: str, context: list, history: list) -> dict:
    """Mock LLM that follows a fixed plan for demonstration."""
    steps_done = [h["tool"] for h in history]

    if "read_file" not in steps_done:
        return {
            "thought": "I need to read the test file to understand the failure",
            "tool": "read_file",
            "args": {"path": "test_main.py"},
        }
    elif "run_command" not in steps_done:
        return {
            "thought": "Let me run the tests to see current status",
            "tool": "run_command",
            "args": {"cmd": "pytest test_main.py -v"},
        }
    else:
        return {
            "thought": "Tests pass. Task is complete.",
            "done": True,
        }


# ── Run Demo ──────────────────────────────────────────────

if __name__ == "__main__":
    print("=" * 50)
    print("AI Agent Loop Demo")
    print("=" * 50)

    result = agent_loop("Fix the failing tests")

    print("\n" + "=" * 50)
    print("Execution History:")
    print("=" * 50)
    for obs in result:
        print(f"  Step {obs['step']}: {obs['thought']}")
        print(f"         Tool: {obs['tool']} → {obs['result'][:50]}")

Step 1: Read test_main.py
  State: [test file contents]

Step 2: Read main.py
  State: [test file, source file]

Step 3: Run pytest
  State: [test file, source file, test results]

Step 4: Edit main.py
  State: [test file, source file, test results, edit applied]

Step 5: Run pytest again
  State: [test file, source file, test results, edit applied, tests pass]

Step 6: Done — all tests pass

# Save the code as agent_demo.py
python agent_demo.py