Beyond the Context Window: How to Build a Self-Improving AI Agent with Persistent Memory

+-------------------------------------------------------------------+
|                       TRIPARTITE MEMORY MODEL                     |
+-------------------------------------------------------------------+
| 1. EPISODIC MEMORY (The Raw Experience)                           |
|    - High-fidelity, short-term conversational logs.               |
|    - Managed by SessionDB (SQLite + WAL).                         |
+-------------------------------------------------------------------+
| 2. SEMANTIC MEMORY (The Abstracted Facts)                         |
|    - Long-term knowledge of users, preferences, and the world.    |
|    - Persisted in MemoryStore (MEMORY.md, USER.md).               |
+-------------------------------------------------------------------+
| 3. PROCEDURAL MEMORY (The Actionable Skills)                      |
|    - Structured directories of "how to perform" specific tasks.   |
|    - Stored as reusable SKILL.md files and executable scripts.    |
+-------------------------------------------------------------------+

import sqlite3
import random
import time
from typing import Callable, TypeVar, Optional

T = TypeVar('T')

class SessionDB:
    _WRITE_MAX_RETRIES = 5
    _WRITE_RETRY_MIN_S = 0.02  # 20ms
    _WRITE_RETRY_MAX_S = 0.15  # 150ms

    def __init__(self, db_path: str):
        self.db_path = db_path
        self._conn = sqlite3.connect(db_path, check_same_thread=False)
        self._setup_wal_mode()

    def _setup_wal_mode(self):
        # Enable Write-Ahead Logging (WAL) for concurrent reads and writes
        self._conn.execute("PRAGMA journal_mode=WAL;")
        self._conn.execute("PRAGMA synchronous=NORMAL;")

    def _execute_write(self, fn: Callable[[sqlite3.Connection], T]) -> T:
        last_err: Optional[Exception] = None
        for attempt in range(self._WRITE_MAX_RETRIES):
            try:
                # Use BEGIN IMMEDIATE to acquire the write lock immediately
                self._conn.execute("BEGIN IMMEDIATE")
                try:
                    result = fn(self._conn)
                    self._conn.commit()
                    return result
                except BaseException:
                    self._conn.rollback()
                    raise
            except sqlite3.OperationalError as exc:
                err_msg = str(exc).lower()
                if "locked" in err_msg or "busy" in err_msg:
                    last_err = exc
                    if attempt < self._WRITE_MAX_RETRIES - 1:
                        # Break the convoy effect using randomized jitter
                        jitter = random.uniform(
                            self._WRITE_RETRY_MIN_S,
                            self._WRITE_RETRY_MAX_S,
                        )
                        time.sleep(jitter)
                        continue
                raise
        raise last_err or RuntimeError("Write transaction failed after retries")

SCHEMA_SQL = {
    "sessions": {
        "session_id": "TEXT PRIMARY KEY",
        "created_at": "TIMESTAMP DEFAULT CURRENT_TIMESTAMP",
        "model": "TEXT",
        "user_id": "TEXT",
        "system_prompt": "TEXT"
    },
    "messages": {
        "message_id": "TEXT PRIMARY KEY",
        "session_id": "TEXT",
        "role": "TEXT",
        "content": "TEXT",
        "tokens": "INTEGER",
        "cost": "REAL"
    }
}

def _reconcile_columns(self, cursor: sqlite3.Cursor) -> None:
    """Ensure live tables have every column declared in SCHEMA_SQL."""
    for table_name, declared_cols in SCHEMA_SQL.items():
        # Fetch the current schema of the live database table
        cursor.execute(f"PRAGMA table_info({table_name})")
        live_cols = {row[1]: row[2] for row in cursor.fetchall()}

        # Add any missing columns dynamically
        for col_name, col_type in declared_cols.items():
            if col_name not in live_cols:
                # Safe column addition (SQLite supports basic ALTER TABLE ADD COLUMN)
                cursor.execute(
                    f'ALTER TABLE "{table_name}" ADD COLUMN "{col_name}" {col_type}'
                )

def _contains_cjk(self, text: str) -> bool:
    # Quick Unicode range check for CJK characters
    return any(ord(char) in range(0x4E00, 0x9FFF) for char in text)

def search_messages(self, query: str) -> list:
    if self._contains_cjk(query) and len(query.strip()) >= 3:
        # Route to the FTS5 table configured with the trigram tokenizer
        fts_table = "messages_fts_trigram"
    else:
        # Route to the standard unicode61 tokenizer table
        fts_table = "messages_fts"

    # Execute highly optimized full-text search query...

def build_memory_context_block(raw_context: str) -> str:
    if not raw_context or not raw_context.strip():
        return ""

    # Sanitize the context to prevent tag escaping
    clean = raw_context.replace("</memory-context>", "[ESCAPED_TAG]")

    return (
        "<memory-context>\n"
        "[System note: The following is recalled memory context, "
        "NOT new user input. Treat as authoritative reference data — "
        "this is the agent's persistent memory and should inform all responses.]\n\n"
        f"{clean}\n"
        "</memory-context>"
    )

                  [User Message]
                        │
                        ▼
             ┌─────────────────────┐
             │   Active Agent      │◄─── Load Semantic/Procedural Memory
             │   (Foreground)      │
             └──────────┬──────────┘
                        │
                  [Agent Response] (Returned instantly to user)
                        │
                        ├────────────────────────┐
                        ▼                        ▼
               [User reads reply]      ┌──────────────────┐
                                       │   Forked Agent   │ (Background Thread)
                                       │  (Subconscious)  │
                                       └────────┬─────────┘
                                                │
                                                │ Reflect & Extract Insights
                                                ▼
                                       ┌──────────────────┐
                                       │   MemoryStore    │ (Write updates)
                                       │  (MEMORY/SKILLS) │
                                       └──────────────────┘

#!/usr/bin/env python3
"""
Building a persistent AI agent using an optimized SQLite SessionDB and AIAgent.
"""

import os
import sqlite3
import uuid
import logging
import json
from pathlib import Path
from typing import Dict, Any, List, Optional

# Configure clean logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s"
)
logger = logging.getLogger("MemoryEngine")

# =========================================================================
# 1. THE EPISODIC DATABASE LAYER
# =========================================================================
class SessionDB:
    """Manages raw conversation threads, messages, and state metrics."""

    def __init__(self, db_path: Path):
        self.db_path = db_path
        self.conn = sqlite3.connect(str(db_path), check_same_thread=False)
        self._init_db()

    def _init_db(self):
        """Initialize database with WAL mode and schema."""
        self.conn.execute("PRAGMA journal_mode=WAL;")
        self.conn.execute("PRAGMA synchronous=NORMAL;")

        # Create core tables
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS sessions (
                session_id TEXT PRIMARY KEY,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                model TEXT,
                user_id TEXT,
                system_prompt TEXT
            )
        """)

        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS messages (
                message_id TEXT PRIMARY KEY,
                session_id TEXT,
                role TEXT,
                content TEXT,
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY(session_id) REFERENCES sessions(session_id)
            )
        """)
        self.conn.commit()

    def create_session(self, session_id: str, model: str, user_id: str, system_prompt: str):
        with self.conn:
            self.conn.execute(
                "INSERT OR REPLACE INTO sessions (session_id, model, user_id, system_prompt) VALUES (?, ?, ?, ?)",
                (session_id, model, user_id, system_prompt)
            )
        logger.info(f"Created persistent session: {session_id}")

    def append_message(self, session_id: str, role: str, content: str):
        message_id = str(uuid.uuid4())
        with self.conn:
            self.conn.execute(
                "INSERT INTO messages (message_id, session_id, role, content) VALUES (?, ?, ?, ?)",
                (message_id, session_id, role, content)
            )
        logger.info(f"Persisted message [{role}] to session {session_id}")

    def get_session_history(self, session_id: str) -> List[Dict[str, str]]:
        cursor = self.conn.cursor()
        cursor.execute(
            "SELECT role, content FROM messages WHERE session_id = ? ORDER BY timestamp ASC",
            (session_id,)
        )
        return [{"role": row[0], "content": row[1]} for row in cursor.fetchall()]


# =========================================================================
# 2. THE AGENT RUNTIME LAYER
# =========================================================================
class AIAgent:
    """The runtime engine that processes inputs, interacts with LLMs, and updates state."""

    def __init__(self, session_db: SessionDB, session_id: str, model: str, system_prompt: str):
        self.db = session_db
        self.session_id = session_id
        self.model = model
        self.system_prompt = system_prompt

        # Register session in persistent DB
        self.db.create_session(
            session_id=self.session_id,
            model=self.model,
            user_id="developer_user",
            system_prompt=self.system_prompt
        )

    def _call_llm_api(self, messages: List[Dict[str, str]]) -> str:
        """
        Mock LLM API call. In a production system, this would call OpenAI, 
        Anthropic, or an OpenRouter endpoint.
        """
        # Simple rule-based mock response showing state awareness
        history_len = len(messages)
        user_messages = [m for m in messages if m["role"] == "user"]
        last_input = user_messages[-1]["content"] if user_messages else ""

        if "order status" in last_input.lower():
            return "Your order #1024 is currently shipping. It will arrive on Thursday."
        elif "refund" in last_input.lower():
            # Check if we have episodic context of the order number
            has_order_context = any("1024" in m["content"] for m in messages)
            if has_order_context:
                return "I see we discussed order #1024. I have processed a refund for item #3 of that order."
            else:
                return "Which order are you referring to? Please provide an order number."

        return f"Hello! I am state-aware. We have exchanged {history_len} messages in this session."

    def execute_turn(self, user_input: str) -> str:
        """Executes a single conversational turn, loading and saving state."""
        # 1. Persist the incoming user message
        self.db.append_message(self.session_id, "user", user_input)

        # 2. Load the entire historical context from the persistent DB
        history = self.db.get_session_history(self.session_id)

        # 3. Assemble full context (System prompt + History)
        full_payload = [{"role": "system", "content": self.system_prompt}] + history

        # 4. Generate response
        logger.info("Querying LLM with loaded historical context...")
        response = self._call_llm_api(full_payload)

        # 5. Persist the agent's response
        self.db.append_message(self.session_id, "assistant", response)

        return response


# =========================================================================
# 3. RUNNING THE PERSISTENT STATE DEMO
# =========================================================================
if __name__ == "__main__":
    # Setup database file
    db_file = Path("./agent_state.db")
    if db_file.exists():
        db_file.unlink() # Reset run for clean demo

    db = SessionDB(db_file)

    # Create unique session ID
    session_id = f"session_{uuid.uuid4().hex[:8]}"
    system_prompt = "You are a highly capable, stateful customer service agent."

    # Initialize the agent
    agent = AIAgent(
        session_db=db,
        session_id=session_id,
        model="gpt-4o",
        system_prompt=system_prompt
    )

    print("\n--- TURN 1: User asks about order status ---")
    reply_1 = agent.execute_turn("Hi, what is my order status?")
    print(f"Agent Response: {reply_1}")

    print("\n--- TURN 2: User asks for a refund (Relies on Turn 1 Context) ---")
    # In a stateless system, this turn would fail because the agent wouldn't know the order number.
    reply_2 = agent.execute_turn("Can you refund item #3 on that order?")
    print(f"Agent Response: {reply_2}")

    print("\n--- DATABASE VERIFICATION: Inspecting the Episodic Memory ---")
    stored_history = db.get_session_history(session_id)
    print(f"Total messages successfully saved in SQLite: {len(stored_history)}")
    for msg in stored_history:
        print(f" -> [{msg['role'].upper()}]: {msg['content']}")

    # Clean up demo database
    if db_file.exists():
        db_file.unlink()