The 3-Layer Memory Architecture Every AI Agent Needs

def run_agent(user_message: str) -> str:
    response = llm.complete(
        system="You are a helpful assistant that manages tasks.",
        user=user_message
    )
    return response

class SemanticMemory:
    def __init__(self, storage_path: str):
        self.path = Path(storage_path)

    def load(self) -> dict:
        if not self.path.exists():
            return {}
        with open(self.path) as f:
            return json.load(f)

    def update(self, key: str, value: any):
        data = self.load()
        data[key] = {
            "value": value,
            "updated_at": datetime.utcnow().isoformat()
        }
        with open(self.path, "w") as f:
            json.dump(data, f, indent=2)

    def to_prompt_block(self) -> str:
        data = self.load()
        if not data:
            return ""
        lines = ["## Background Context"]
        for key, entry in data.items():
            lines.append(f"- {key}: {entry['value']}")
        return "\n".join(lines)

# Usage
memory = SemanticMemory("agent_memory.json")
memory.update("user_company", "Acme Corp, B2B SaaS, 50-person team")
memory.update("user_goal", "Reduce customer churn from 12% to 8% by Q3")
memory.update("preferred_output", "Bullet points with action items at the bottom")

class WorkingMemory:
    def __init__(self, max_tokens: int = 4000):
        self.messages = []
        self.max_tokens = max_tokens

    def add(self, role: str, content: str):
        self.messages.append({"role": role, "content": content})
        self._trim_if_needed()

    def _trim_if_needed(self):
        total_chars = sum(len(m["content"]) for m in self.messages)
        estimated_tokens = total_chars / 4

        while estimated_tokens > self.max_tokens and len(self.messages) > 2:
            removed = self.messages.pop(1)  # Keep first message
            total_chars -= len(removed["content"])
            estimated_tokens = total_chars / 4

    def summarize_and_compress(self, llm) -> str:
        """Ask the LLM to compress old messages into a summary."""
        if len(self.messages) < 6:
            return ""

        to_compress = self.messages[1:len(self.messages)//2]
        compress_prompt = f"""Summarize these conversation turns into a concise paragraph
        capturing all key decisions, facts learned, and context established:
        {json.dumps(to_compress, indent=2)}"""

        summary = llm.complete(compress_prompt)

        # Replace compressed messages with summary
        self.messages = (
            [self.messages[0]] +
            [{"role": "system", "content": f"[Earlier in this session]: {summary}"}] +
            self.messages[len(self.messages)//2:]
        )
        return summary

import chromadb

class EpisodicMemory:
    def __init__(self, collection_name: str = "agent_episodes"):
        self.client = chromadb.Client()
        self.collection = self.client.get_or_create_collection(collection_name)

    def store_episode(self, episode_id: str, content: str, metadata: dict = None):
        """Store a memory of what happened in a session."""
        self.collection.add(
            documents=[content],
            ids=[episode_id],
            metadatas=[metadata or {}]
        )

    def recall(self, query: str, n_results: int = 3) -> list[str]:
        """Retrieve relevant past episodes based on semantic similarity."""
        results = self.collection.query(
            query_texts=[query],
            n_results=n_results
        )
        return results["documents"][0] if results["documents"] else []

    def to_prompt_block(self, current_context: str) -> str:
        memories = self.recall(current_context)
        if not memories:
            return ""
        lines = ["## Relevant Past Context"]
        for memory in memories:
            lines.append(f"- {memory}")
        return "\n".join(lines)

# After each session, store what happened
episodic = EpisodicMemory()
episodic.store_episode(
    episode_id="session_2026_02_26_001",
    content="User discussed Q2 pricing strategy. Decided to hold prices but add a premium tier at $499/mo.",
    metadata={"date": "2026-02-26", "topic": "pricing", "outcome": "decision_made"}
)

class MemoryAwareAgent:
    def __init__(self):
        self.semantic = SemanticMemory("semantic.json")
        self.working = WorkingMemory(max_tokens=3000)
        self.episodic = EpisodicMemory()
        self.llm = YourLLMClient()

    def build_system_prompt(self, user_message: str) -> str:
        semantic_block = self.semantic.to_prompt_block()
        episodic_block = self.episodic.to_prompt_block(user_message)

        return f"""You are a strategic assistant with memory across sessions.

{semantic_block}

{episodic_block}

Use this context naturally. Don't mention the memory system — just be helpful."""

    def respond(self, user_message: str) -> str:
        system = self.build_system_prompt(user_message)
        self.working.add("user", user_message)

        response = self.llm.complete(
            system=system,
            messages=self.working.get_messages()
        )

        self.working.add("assistant", response)
        return response

    def end_session(self, session_summary: str):
        """Call this when a session ends to persist what happened."""
        self.episodic.store_episode(
            episode_id=f"session_{datetime.utcnow().isoformat()}",
            content=session_summary,
            metadata={"date": datetime.utcnow().date().isoformat()}
        )
        # Extract semantic memory updates...