Skip to content
Codeloom
AI

AI Agent Architecture: Tools, Memory, and Planning

Learn how to build AI agents with tool use, memory systems, and planning capabilities using Python and LLMs for autonomous task completion.

·10 min read · By Codeloom
Intermediate 15 min read

What you'll learn

  • Core components of an AI agent: perception, reasoning, action
  • How to implement tool use for agents to interact with external systems
  • Memory architectures: short-term, long-term, and episodic
  • Planning strategies: ReAct, chain-of-thought, and task decomposition

Prerequisites

  • Python fundamentals
  • Basic understanding of LLM APIs
  • Familiarity with function calling

An AI agent is more than a chatbot. While a chatbot responds to messages, an agent perceives its environment, reasons about goals, takes actions, and learns from results. This guide covers the architecture behind practical AI agents, from tool integration to memory systems to planning strategies.

The Agent Loop

Every agent follows a fundamental loop: observe, think, act, repeat. The LLM serves as the reasoning engine, deciding what to do next based on the current state.

from openai import OpenAI
import json
from typing import Any

client = OpenAI()

class SimpleAgent:
    """A minimal agent that can use tools in a loop."""
    
    def __init__(self, tools: dict, system_prompt: str, max_steps: int = 10):
        self.tools = tools
        self.system_prompt = system_prompt
        self.max_steps = max_steps
        self.messages = [{"role": "system", "content": system_prompt}]
    
    def run(self, user_input: str) -> str:
        """Run the agent loop until completion or max steps."""
        self.messages.append({"role": "user", "content": user_input})
        
        for step in range(self.max_steps):
            response = client.chat.completions.create(
                model="gpt-4o",
                messages=self.messages,
                tools=self._get_tool_schemas(),
                tool_choice="auto"
            )
            
            message = response.choices[0].message
            self.messages.append(message)
            
            # If no tool calls, the agent is done
            if not message.tool_calls:
                return message.content
            
            # Execute each tool call
            for tool_call in message.tool_calls:
                result = self._execute_tool(
                    tool_call.function.name,
                    json.loads(tool_call.function.arguments)
                )
                self.messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": json.dumps(result)
                })
            
            print(f"Step {step + 1}: Called {[tc.function.name for tc in message.tool_calls]}")
        
        return "Agent reached maximum steps without completing."
    
    def _execute_tool(self, name: str, args: dict) -> Any:
        """Execute a tool by name with given arguments."""
        if name not in self.tools:
            return {"error": f"Unknown tool: {name}"}
        try:
            return self.tools[name](**args)
        except Exception as e:
            return {"error": str(e)}
    
    def _get_tool_schemas(self) -> list:
        """Return OpenAI-format tool schemas."""
        # In practice, generate these from function signatures
        return self._tool_schemas

This loop is deceptively simple but forms the backbone of systems like AutoGPT, BabyAGI, and production agents at companies like Anthropic, Google, and OpenAI.

Tool Architecture

Tools give agents the ability to interact with the outside world. A well-designed tool system is the difference between a toy demo and a useful agent.

import requests
from datetime import datetime

def search_web(query: str) -> dict:
    """Search the web and return results."""
    # Using a search API (e.g., SerpAPI, Brave Search)
    response = requests.get(
        "https://api.search.brave.com/res/v1/web/search",
        params={"q": query, "count": 5},
        headers={"X-Subscription-Token": "YOUR_API_KEY"}
    )
    results = response.json().get("web", {}).get("results", [])
    return {
        "results": [
            {"title": r["title"], "url": r["url"], "snippet": r.get("description", "")}
            for r in results[:5]
        ]
    }

def read_file(file_path: str) -> dict:
    """Read contents of a file."""
    try:
        with open(file_path, "r") as f:
            content = f.read()
        return {"content": content, "size": len(content)}
    except FileNotFoundError:
        return {"error": f"File not found: {file_path}"}

def run_python(code: str) -> dict:
    """Execute Python code and return the output."""
    import subprocess
    try:
        result = subprocess.run(
            ["python", "-c", code],
            capture_output=True, text=True, timeout=30
        )
        return {
            "stdout": result.stdout,
            "stderr": result.stderr,
            "return_code": result.returncode
        }
    except subprocess.TimeoutExpired:
        return {"error": "Code execution timed out (30s limit)"}

def get_current_time() -> dict:
    """Get the current date and time."""
    now = datetime.now()
    return {"datetime": now.isoformat(), "timezone": "UTC"}

# Tool registry with schemas
TOOLS = {
    "search_web": search_web,
    "read_file": read_file,
    "run_python": run_python,
    "get_current_time": get_current_time,
}

TOOL_SCHEMAS = [
    {
        "type": "function",
        "function": {
            "name": "search_web",
            "description": "Search the web for current information",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Search query"}
                },
                "required": ["query"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "run_python",
            "description": "Execute Python code and return output",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {"type": "string", "description": "Python code to execute"}
                },
                "required": ["code"]
            }
        }
    }
]

Key design principles for tools: keep them focused on a single action, return structured data (not raw HTML or unformatted text), include error handling, and write clear descriptions because the LLM reads them to decide when and how to use each tool.

Memory Architecture

Agents need memory to maintain context, learn from past interactions, and avoid repeating mistakes. There are three types of memory that matter.

Short-Term Memory (Working Memory)

This is the conversation history within a single session. The simplest implementation is the message array, but as conversations grow, you need to manage its size.

class WorkingMemory:
    """Manages short-term conversation memory with compaction."""
    
    def __init__(self, max_tokens: int = 50000):
        self.messages = []
        self.max_tokens = max_tokens
    
    def add(self, message: dict):
        """Add a message and compact if needed."""
        self.messages.append(message)
        self._compact_if_needed()
    
    def _compact_if_needed(self):
        """Summarize old messages when memory gets too large."""
        total_tokens = self._estimate_tokens()
        
        if total_tokens <= self.max_tokens:
            return
        
        # Keep system message and last N messages
        system_msg = self.messages[0] if self.messages[0]["role"] == "system" else None
        recent = self.messages[-10:]
        old = self.messages[1:-10] if system_msg else self.messages[:-10]
        
        if not old:
            return
        
        # Summarize old messages
        summary = self._summarize(old)
        
        self.messages = []
        if system_msg:
            self.messages.append(system_msg)
        self.messages.append({
            "role": "system",
            "content": f"Summary of earlier conversation:\n{summary}"
        })
        self.messages.extend(recent)
    
    def _summarize(self, messages: list) -> str:
        """Use LLM to summarize old messages."""
        conversation = "\n".join(
            f"{m['role']}: {m.get('content', '[tool call]')}" 
            for m in messages
        )
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "user",
                "content": f"Summarize this conversation concisely, preserving key decisions and facts:\n\n{conversation}"
            }],
            temperature=0.2
        )
        return response.choices[0].message.content
    
    def _estimate_tokens(self) -> int:
        content = json.dumps(self.messages)
        return len(content) // 4  # Rough estimate

Long-Term Memory (Semantic Memory)

Long-term memory persists across sessions. It stores facts, user preferences, and learned information using a vector database for retrieval.

class LongTermMemory:
    """Persistent memory using embeddings for retrieval."""
    
    def __init__(self, collection_name: str = "agent_memory"):
        import chromadb
        self.chroma = chromadb.PersistentClient(path="./agent_memory_db")
        self.collection = self.chroma.get_or_create_collection(collection_name)
    
    def store(self, content: str, metadata: dict = None):
        """Store a memory with automatic embedding."""
        import uuid
        memory_id = str(uuid.uuid4())
        self.collection.add(
            documents=[content],
            metadatas=[metadata or {}],
            ids=[memory_id]
        )
        return memory_id
    
    def recall(self, query: str, top_k: int = 5) -> list[str]:
        """Retrieve relevant memories for a query."""
        results = self.collection.query(
            query_texts=[query],
            n_results=top_k
        )
        return results["documents"][0] if results["documents"] else []
    
    def store_interaction_summary(self, user_input: str, agent_output: str):
        """Store a summary of an interaction for future reference."""
        summary_response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "user",
                "content": f"Extract key facts and decisions from this interaction in 2-3 bullet points:\nUser: {user_input}\nAgent: {agent_output}"
            }],
            temperature=0.1
        )
        summary = summary_response.choices[0].message.content
        self.store(summary, metadata={
            "type": "interaction_summary",
            "timestamp": datetime.now().isoformat()
        })

Episodic Memory

Episodic memory stores complete past experiences that can be retrieved and replayed to help with similar future tasks.

class EpisodicMemory:
    """Stores and retrieves complete task episodes."""
    
    def __init__(self):
        self.episodes = []
    
    def record_episode(self, task: str, steps: list[dict], outcome: str, success: bool):
        """Record a complete task episode."""
        self.episodes.append({
            "task": task,
            "steps": steps,
            "outcome": outcome,
            "success": success,
            "timestamp": datetime.now().isoformat()
        })
    
    def find_similar_episodes(self, current_task: str, top_k: int = 3) -> list[dict]:
        """Find past episodes similar to the current task."""
        if not self.episodes:
            return []
        
        # Use LLM to find relevant episodes
        episode_summaries = "\n".join(
            f"{i}. [{e['outcome']}] {e['task']}"
            for i, e in enumerate(self.episodes)
        )
        
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "user",
                "content": f"Which of these past tasks are most similar to: '{current_task}'?\n\n{episode_summaries}\n\nReturn the indices as a JSON list."
            }],
            response_format={"type": "json_object"},
            temperature=0
        )
        
        indices = json.loads(response.choices[0].message.content).get("indices", [])
        return [self.episodes[i] for i in indices[:top_k] if i < len(self.episodes)]

Planning Strategies

Planning determines how an agent breaks down complex tasks and decides what to do next.

ReAct Pattern (Reasoning + Acting)

The ReAct pattern interleaves thinking and acting. The agent explicitly reasons about what to do before each action.

REACT_SYSTEM_PROMPT = """You are an AI agent that solves tasks step by step.

For each step, you MUST follow this format:
Thought: [Your reasoning about what to do next]
Action: [The tool to call, or "finish" if done]
Action Input: [The input for the tool]

After receiving an observation from a tool, reason about the result before deciding the next action.

Available tools: {tool_descriptions}
"""

def react_agent(task: str, tools: dict, max_steps: int = 10) -> str:
    """Run a ReAct-style agent."""
    messages = [
        {"role": "system", "content": REACT_SYSTEM_PROMPT.format(
            tool_descriptions="\n".join(f"- {name}: {func.__doc__}" for name, func in tools.items())
        )},
        {"role": "user", "content": f"Task: {task}"}
    ]
    
    for step in range(max_steps):
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
            temperature=0.2
        )
        
        reply = response.choices[0].message.content
        messages.append({"role": "assistant", "content": reply})
        print(f"\n--- Step {step + 1} ---\n{reply}")
        
        # Check if agent decided to finish
        if "Action: finish" in reply.lower():
            return reply
        
        # Parse and execute the action
        action, action_input = parse_react_output(reply)
        if action and action in tools:
            observation = tools[action](**json.loads(action_input))
            messages.append({
                "role": "user",
                "content": f"Observation: {json.dumps(observation)}"
            })
    
    return "Max steps reached."

Task Decomposition

For complex tasks, the agent first creates a plan, then executes each sub-task.

def plan_and_execute(task: str, tools: dict) -> str:
    """Decompose a task into subtasks, then execute each one."""
    
    # Planning phase
    plan_response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "Break down the given task into a numbered list of concrete subtasks. Each subtask should be achievable with the available tools."},
            {"role": "user", "content": f"Task: {task}\n\nAvailable tools: {list(tools.keys())}"}
        ],
        response_format={"type": "json_object"},
        temperature=0.3
    )
    
    plan = json.loads(plan_response.choices[0].message.content)
    subtasks = plan.get("subtasks", [])
    print(f"Plan created with {len(subtasks)} subtasks")
    
    # Execution phase
    results = []
    for i, subtask in enumerate(subtasks):
        print(f"\nExecuting subtask {i+1}: {subtask}")
        
        # Each subtask gets its own mini agent loop
        result = react_agent(
            task=f"Complete this specific subtask: {subtask}\n\nContext from previous subtasks: {json.dumps(results)}",
            tools=tools,
            max_steps=5
        )
        results.append({"subtask": subtask, "result": result})
    
    # Synthesis phase
    synthesis = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "Synthesize the results of all subtasks into a final comprehensive answer."},
            {"role": "user", "content": f"Original task: {task}\n\nSubtask results: {json.dumps(results)}"}
        ]
    )
    
    return synthesis.choices[0].message.content

Error Handling and Recovery

Production agents must handle failures gracefully. Tools fail, LLMs hallucinate tool names, and plans hit dead ends.

class ResilientAgent(SimpleAgent):
    """Agent with error handling and recovery."""
    
    def _execute_tool(self, name: str, args: dict) -> Any:
        """Execute with retries and error recovery."""
        max_retries = 2
        
        for attempt in range(max_retries + 1):
            try:
                result = self.tools[name](**args)
                return result
            except Exception as e:
                if attempt < max_retries:
                    print(f"Tool {name} failed (attempt {attempt + 1}), retrying...")
                    continue
                return {
                    "error": f"Tool '{name}' failed after {max_retries + 1} attempts: {str(e)}",
                    "suggestion": "Try an alternative approach or different tool."
                }
    
    def run(self, user_input: str) -> str:
        """Run with stuck detection."""
        self.messages.append({"role": "user", "content": user_input})
        last_actions = []
        
        for step in range(self.max_steps):
            response = client.chat.completions.create(
                model="gpt-4o",
                messages=self.messages,
                tools=self._get_tool_schemas(),
                tool_choice="auto"
            )
            
            message = response.choices[0].message
            self.messages.append(message)
            
            if not message.tool_calls:
                return message.content
            
            # Detect loops (same tool called 3+ times in a row)
            current_actions = [tc.function.name for tc in message.tool_calls]
            last_actions.append(tuple(current_actions))
            
            if len(last_actions) >= 3 and len(set(last_actions[-3:])) == 1:
                self.messages.append({
                    "role": "user",
                    "content": "You appear to be stuck in a loop. Try a completely different approach or explain what's blocking you."
                })
            
            for tool_call in message.tool_calls:
                result = self._execute_tool(
                    tool_call.function.name,
                    json.loads(tool_call.function.arguments)
                )
                self.messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": json.dumps(result)
                })
        
        return "Agent reached maximum steps. Here is what was accomplished so far."

Wrapping Up

AI agent architecture rests on three pillars: tools for interacting with the world, memory for maintaining context and learning, and planning for breaking down complex tasks. The agent loop itself is simple, but the quality of each component determines whether your agent is useful or frustrating. Start with a basic loop and a few well-designed tools. Add working memory compaction when conversations get long. Introduce long-term memory when users return across sessions. Layer in planning when tasks require multiple steps. And always build error handling and loop detection from the start, because agents will surprise you with creative ways to get stuck.