Agent Memory & State Management

Agents need memory to function effectively across steps and sessions. Without memory, an agent forgets what it has already done, loses track of user preferences, and repeats work. This lesson covers working memory, long-term memory, episodic memory, entity memory, and how to build memory-augmented agents.

---

Types of Agent Memory

┌─────────────────────────────────────────────────────────────────┐
│                      AGENT MEMORY                               │
│                                                                 │
│  ┌───────────────┐  ┌──────────────────┐  ┌─────────────────┐  │
│  │  WORKING      │  │  EPISODIC        │  │  SEMANTIC       │  │
│  │  MEMORY       │  │  MEMORY          │  │  MEMORY         │  │
│  │               │  │                  │  │                 │  │
│  │  Current task │  │  Past experiences│  │  Facts & entities│ │
│  │  context,     │  │  stored as       │  │  extracted from │  │
│  │  scratchpad   │  │  episodes        │  │  interactions   │  │
│  └───────────────┘  └──────────────────┘  └─────────────────┘  │
│                                                                 │
│  ┌───────────────┐  ┌──────────────────┐                       │
│  │  PROCEDURAL   │  │  ENTITY          │                       │
│  │  MEMORY       │  │  MEMORY          │                       │
│  │               │  │                  │                       │
│  │  Learned      │  │  Knowledge about │                       │
│  │  procedures   │  │  specific people,│                       │
│  │  and patterns │  │  places, things  │                       │
│  └───────────────┘  └──────────────────┘                       │
└─────────────────────────────────────────────────────────────────┘

Memory Type	Scope	Persistence	Example
Working	Current task	Until task ends	Scratchpad, intermediate results
Episodic	Past interactions	Long-term (DB)	"Last time, the user preferred CSV"
Semantic	Facts and knowledge	Long-term (DB)	"User's company is Acme Corp"
Entity	Specific entities	Long-term (DB)	People, projects, concepts the agent tracks
Procedural	How to do things	Long-term (code)	Learned workflows and patterns

---

Working Memory

Working memory is the agent's scratchpad for the current task. It holds intermediate results, the current plan, and recent observations.

from dataclasses import dataclass, field
from typing import Any

@dataclass
class WorkingMemory:
    """Short-term memory for the current task."""
    task: str = ""
    plan: list[str] = field(default_factory=list)
    current_step: int = 0
    scratchpad: list[dict] = field(default_factory=list)
    variables: dict[str, Any] = field(default_factory=dict)

    def add_observation(self, step: str, observation: str):
        self.scratchpad.append({
            "step": step,
            "observation": observation,
        })

    def set_variable(self, key: str, value: Any):
        self.variables[key] = value

    def get_variable(self, key: str, default: Any = None) -> Any:
        return self.variables.get(key, default)

    def to_context_string(self) -> str:
        parts = [f"Current task: {self.task}"]
        if self.plan:
            parts.append("Plan: " + " -> ".join(self.plan))
            parts.append(f"Current step: {self.current_step + 1}/{len(self.plan)}")
        if self.scratchpad:
            parts.append("Observations:")
            for entry in self.scratchpad[-5:]:  # Last 5 observations
                parts.append(f"  - {entry['step']}: {entry['observation'][:200]}")
        if self.variables:
            parts.append("Variables: " + str(self.variables))
        return "\n".join(parts)

---

Episodic Memory

Episodic memory stores past interactions as discrete episodes that can be recalled when relevant.

import json
import sqlite3
from datetime import datetime

class EpisodicMemory:
    """Stores and retrieves past interaction episodes."""

    def __init__(self, db_path: str = "episodic_memory.db"):
        self.conn = sqlite3.connect(db_path)
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS episodes (
                id INTEGER PRIMARY KEY,
                timestamp TEXT,
                task TEXT,
                outcome TEXT,
                steps_taken INTEGER,
                success BOOLEAN,
                summary TEXT
            )
        """)
        self.conn.commit()

    def store_episode(self, task: str, outcome: str,
                      steps_taken: int, success: bool, summary: str):
        self.conn.execute(
            "INSERT INTO episodes (timestamp, task, outcome, steps_taken, success, summary) "
            "VALUES (?, ?, ?, ?, ?, ?)",
            (datetime.now().isoformat(), task, outcome, steps_taken, success, summary),
        )
        self.conn.commit()

    def recall_similar(self, current_task: str, limit: int = 3) -> list[dict]:
        """Recall episodes similar to the current task (simple keyword match)."""
        keywords = current_task.lower().split()
        rows = self.conn.execute(
            "SELECT task, outcome, success, summary FROM episodes ORDER BY id DESC LIMIT 50"
        ).fetchall()