Memory-Augmented and Persistent Agents

Comprehensive course on memory augmentation and persistent agents in AI.

Table of Contents

• Module 1 — Why Memory Is Essential to Agentic Behavior
  • 1.1 Why Does Memory Matter?
  • 1.2 Stateless Agents vs Memory-Augmented Agents
  • 1.3 Types of Memory in AI Agents
  • 1.4 Memory Management: Limitations and Strategies
• Module 2 — Comparing Persistent Memory Techniques and Frameworks
  • 2.1 Introduction to the Main Frameworks
  • 2.2 Semantic Kernel — Managing Conversation History
  • 2.3 Using a Memory Store with Mem0
  • 2.4 Haystack — Memory-Aware RAG Agent
  • 2.5 MemGPT (Letta) — Agents with Memory Blocks
  • 2.6 LlamaIndex — Multi-Agent AI Engineering Team
• Module 3 — End-to-End Demo: Multi-Agent System with Persistent Memory
  • 3.1 LangChain + LangGraph — Hybrid Memory Strategies
  • 3.2 Short-Term Memory (Thread-Level)
  • 3.3 Long-Term Persistent Memory with ChromaDB
  • 3.4 Performance and Production Scaling
  • 3.5 Testing and Performance Monitoring
• References and Resources

Module 1 — Why Memory Is Essential to Agentic Behavior

1.1 Why Does Memory Matter?

AI agents are designed to execute tasks and achieve collective goals on behalf of users. However, their capabilities are limited without proper memory management.

AI Agent Architecture

flowchart LR
    EXT[External Input]
    SENS[Sensors\n👁️ Eyes & Ears]
    ENV[Environment\nDigital / Physical]
    PROC[Processing &\nLLM Reasoning]
    ACT[Actuators\n✋ Actions]
    OUT[Output / Workflow]

    EXT --> SENS --> ENV --> PROC --> ACT --> OUT
    ACT --> ENV

• Sensors: receive external inputs (messages, clicks, requests, data)
• Environment: the world where the agent operates (digital: app, website, servers; physical: robots, IoT)
• Actuators: execute actions based on the current state of the environment
• Workflow: series of steps orchestrated by multiple agents collaborating to achieve a goal

Fundamental Problem: LLMs Are Stateless by Default

flowchart TD
    A[User request] --> B[LLM]
    B --> C[Response]
    C --> D[Total forgetting ❌]
    D --> E[Next request\nBack to zero]

Language models have no real-time awareness. Without memory management, agents share the same limitations.

1.2 Stateless Agents vs Memory-Augmented Agents

Basic AI Agent Loop

flowchart LR
    I[Perceive\nInputs / Tools / Data] --> R[Reason\nDecide what to do]
    R --> A[Act\nActuators / Workflow]
    A --> I

Stateless Agent — Simple Reflex Agent

The simplest agent has no memory and reacts only to the current percept via fixed condition-action rules.

# M1 — Simple Reflex Agent (no memory)
def simple_reflex_agent(percept):
    """
    Simple reflex agent: selects an action
    based solely on the current percept.
    """
    if percept == "dirty":
        return "clean"
    elif percept == "clean":
        return "do_nothing"
    else:
        return "unknown_action"

# Environment simulation
environment_states = ["dirty", "clean", "dirty"]

for state in environment_states:
    action = simple_reflex_agent(state)
    print(f"Percept: {state} -> Action: {action}")

Output:

Percept: dirty -> Action: clean
Percept: clean -> Action: do_nothing
Percept: dirty -> Action: clean

Comparison Table

1.3 Types of Memory in AI Agents

AI agents use three main types of memory, inspired by human cognition:

mindmap
  root((AI Memory))
    Episodic
      Personal experiences
      Past events
      Time-contextualized
      Ex: yesterday I googled sci-fi movies
    Semantic
      General knowledge
      Universally shared facts
      Non-contextual
      Ex: pandas is a Python library
    Procedural
      How to do things
      Skills and routines
      Automated workflows
      Ex: steps to deploy an app

Episodic vs Semantic Memory Comparison

Procedural Memory

Encodes the how-to — routines and automations:

• Application deployment steps
• Code review workflow
• API integration process

1.4 Memory Management: Limitations and Strategies

LLM Limitations Without Memory

flowchart TD
    LLM[Stateless LLM] --> L1[Fixed context window]
    LLM --> L2[No persistent memory]
    LLM --> L3[Slow API responses]
    LLM --> L4[High token costs]
    LLM --> L5[Poor workflow continuity]
    LLM --> L6[No adaptation]

Memory Management Strategies and Techniques

flowchart LR
    STR[Memory\nStrategies] --> S1[Short-term memory\nThread-level]
    STR --> S2[Long-term memory\nVector Store]
    STR --> S3[Summary & Compression\nPruning]
    STR --> S4[RAG + Tool Calling]
    STR --> S5[Hybrid methods\nSelf-managed Stateful]

Typical Agentic Application Use Cases

• Conversational AI assistants
• Customer and technical support agents
• Multi-step workflow automation
• AI research workflows

Module 2 — Comparing Persistent Memory Techniques and Frameworks

2.1 Introduction to the Main Frameworks

mindmap
  root((Agentic Memory Frameworks))
    Semantic Kernel
      Microsoft
      ChatCompletionAgent
      ChatHistoryAgentThread
    Haystack
      deepset
      Native RAG
      Agents + Mem0
    MemGPT / Letta
      Memory Blocks
      Self-managed
      Auto-update
    LlamaIndex
      FunctionAgent
      AgentWorkflow
      ChromaDB
    LangChain + LangGraph
      InMemorySaver
      StateGraph
      LangSmith

Project goal: transform stateless agents into stateful super-agents with augmented memory to automate software engineering tasks.

2.2 Semantic Kernel — Managing Conversation History

Semantic Kernel is a Microsoft framework that offers a simple way to maintain and persist conversation context via a thread.

Semantic Kernel Architecture

sequenceDiagram
    participant U as User
    participant A as ChatCompletionAgent
    participant K as Kernel
    participant T as Thread (ChatHistory)
    participant LLM as OpenAI GPT-4

    U->>A: User message
    A->>T: Retrieve history
    T-->>A: Past context
    A->>K: Prepare request
    K->>LLM: API call
    LLM-->>K: Response
    K-->>A: Enriched response
    A->>T: Save to thread
    A-->>U: Final response

Implementation Steps

1. Installation

pip install semantic-kernel python-dotenv colorama

2. Create the ChatCompletion service

from semantic_kernel.connectors.ai.open_ai import OpenAIChatCompletion
import os
from dotenv import load_dotenv

load_dotenv()

chat_completion_service = OpenAIChatCompletion(
    ai_model_id="gpt-4",
    api_key=os.getenv("OPENAI_API_KEY")
)

3. Initialize the Kernel and add the service

from semantic_kernel import Kernel
from semantic_kernel.connectors.ai.chat_completion_client_base import ChatCompletionClientBase

kernel = Kernel()
kernel.add_service(chat_completion_service, ChatCompletionClientBase)

4. Create the ChatCompletionAgent

from semantic_kernel.agents.chat_completion.chat_completion_agent import (
    ChatCompletionAgent, 
    ChatHistoryAgentThread
)

prompt = """
You are a Senior Developer Assistant. 
You think like a senior engineer with strong software architecture judgment.
You help with:
  - coding and debugging
  - architecture decisions
  - refactoring & performance optimization
  - reviewing pull requests
  - troubleshooting production issues
"""

agent = ChatCompletionAgent(
    kernel=kernel,
    name="developer_agent",
    instructions=prompt,
)

5. Manage conversation history with a Thread

# Create a thread with a unique ID to maintain history
thread = ChatHistoryAgentThread(thread_id="dev_session_thread")

Mock Test — Maintaining Conversation History

from colorama import Fore, Style

async def save_to_thread(agent, thread, messages):
    """Sends messages and updates the thread correctly."""
    for msg in messages:
        print(f"{Fore.BLUE}USER → {msg}{Style.RESET_ALL}")
        response = await agent.get_response(messages=msg, thread=thread)
        updated_thread = response.thread
        print(f"ASSISTANT → {response.content}\n")
    return updated_thread

async def count_messages(thread):
    """Counts and retrieves messages from the thread."""
    count = 0
    retrieved = []
    async for message in thread.get_messages():
        count += 1
        content = message.content
        if isinstance(content, list):
            content = " ".join(getattr(c, "content", str(c)) for c in content)
        elif not isinstance(content, str):
            content = getattr(content, "content", str(content))
        retrieved.append(f"{str(message.role).upper()}: {content}")
    return count, retrieved

async def test_thread_memory(agent, thread):
    """Test: save, count and retrieve messages in the thread."""
    test_messages = [
        "Write a Python function that converts a string into a URL-friendly slug",
        "Refactor previous Python function to improve its performance and readability",
        "Create a function that loads and validates a JSON configuration file",
        "Update loading and validation function with error handling and logging",
        "Write a function that reads a CSV file and returns a list of dictionaries"
    ]

    updated_thread = await save_to_thread(agent, thread, test_messages)
    count, messages = await count_messages(updated_thread)
    print(f"\n✅ Total messages in thread: {count}")

    assert count >= len(test_messages), "Messages missing in thread!"
    assert any("json" in m.lower() for m in messages), "JSON task missing!"
    assert any("csv" in m.lower() for m in messages), "CSV task missing!"

    return updated_thread

# Execution
new_thread = await test_thread_memory(agent, thread)

Interactive Conversation Loop

async def run():
    while True:
        await print_thread_summary(thread)
        user_text = input("You: ").strip()
        if user_text.lower() in {"exit", "quit"}:
            break

        response = await agent.get_response(messages=user_text, thread=thread)
        thread = response.thread
        print(f"\n{response.role}: {response.content}")

Resources: Semantic Kernel Agent Framework | ChatHistoryAgentThread

2.3 Using a Memory Store with Mem0

Mem0 is a managed memory layer designed for AI applications and agentic systems. The goal is to enable AI agents to remember, retrieve, and learn from past interactions.

Mem0 Memory Lifecycle

flowchart LR
    AGENT[AI Agent] -->|store| MEM0[Mem0\nMemory Store]
    MEM0 -->|retrieve| AGENT
    MEM0 --> EXT[Extraction]
    MEM0 --> CON[Consolidation]
    MEM0 --> RET[Retrieval]
    
    subgraph Operations
        EXT
        CON
        RET
    end

Mem0 Advantages

• Transforms stateless agents into agents that accumulate intelligence
• Memory shareable across multiple agents
• Monitoring dashboard with latency and logs
• Namespace per user_id to structure memories

Initial Configuration

from mem0 import MemoryClient
import os

# Initialize the Mem0 client
memory = MemoryClient(api_key=os.getenv("MEM0_API_KEY"))
MEMORY_USER_ID = "engineering_user"
MEMORY_TYPE = "engineering_feedback"

filters = {
    "OR": [{"user_id": MEMORY_USER_ID}]
}

Helper Functions to Save and Retrieve

def store_feedback(message, memory_type: str = MEMORY_TYPE):
    """Saves feedback to memory with a specific type."""
    memory.add(
        message,
        user_id=MEMORY_USER_ID,
        metadata={"type": memory_type}
    )

def retrieve_memory(user_input, limit=5):
    """Retrieves feedback from memory based on user input."""
    memories = memory.search(
        user_input,
        filters=filters,
        limit=limit
    )

    if isinstance(memories, dict):
        items = memories.get("results", [])
    elif isinstance(memories, list):
        items = memories
    else:
        items = []

    filtered_items = [
        m for m in items
        if isinstance(m, dict)
        and m.get("metadata", {}).get("type") == MEMORY_TYPE
    ]

    return "\n".join(
        m.get("memory", "") for m in filtered_items if m.get("memory")
    )

Required Environment Variables

MEM0_API_KEY=<your-mem0-api-key>
MEM0_USER_ID=<user-identifier>
OPENAI_API_KEY=<your-openai-key>

Resources: Mem0 Dashboard | Mem0 API Keys

2.4 Haystack — Memory-Aware RAG Agent

Goal: build an AI engineering agent that can:

• Fix bugs
• Follow code standards
• Avoid repeated mistakes
• Learn from past reviews

Architecture: Haystack + RAG + Mem0

flowchart TD
    U[User] --> AGT[Haystack Agent\nengineering_specialist]
    AGT --> MT[Tool: engineering_memory\nMem0 Memory Store]
    AGT --> ST[Tool: coding_standards\nRAG Vector Store]
    MT -->|Retrieve history| AGT
    ST -->|Retrieve standards| AGT
    AGT --> LLM[OpenAI GPT-4o-mini]
    LLM --> R[Augmented response]
    R --> SM[Store Feedback\nMem0]

Installation

pip install haystack-ai haystack-experimental mem0ai colorama

Step 1 — Loading Code Standards (RAG)

from haystack import Document
from haystack.document_stores.in_memory import InMemoryDocumentStore

doc_store = InMemoryDocumentStore()

def load_standards(file_path: str):
    """Loads code standards from a text file."""
    with open(file_path, "r", encoding="utf-8") as f:
        lines = [
            line.strip()
            for line in f.readlines()
            if line.strip() and not line.startswith("#")
        ]
    return lines

standards = load_standards("coding_standards.txt")
docs = [
    Document(content=standard, meta={"type": "coding_standard"})
    for standard in standards
]
doc_store.write_documents(docs)
print(f"Loaded {len(docs)} coding standards.")

Step 2 — Creating Embeddings and Retriever

from haystack.components.retrievers import InMemoryEmbeddingRetriever
from haystack.components.embedders import (
    SentenceTransformersDocumentEmbedder,
    SentenceTransformersTextEmbedder,
)
from haystack.document_stores.types import DuplicatePolicy

embedder = SentenceTransformersDocumentEmbedder()
embedded_docs = embedder.run(documents=docs)["documents"]
doc_store.write_documents(embedded_docs, policy=DuplicatePolicy.OVERWRITE)

query_embedder = SentenceTransformersTextEmbedder()
retriever = InMemoryEmbeddingRetriever(doc_store)

Step 3 — Defining Functions and Tools

# Retrieve from memory of past queries
def memory_function(query: str):
    return retrieve_memory(query)

# Retrieve from vector store standards
def standards_function(query: str):
    query_embedding = query_embedder.run(text=query)["embedding"]
    result = retriever.run(query_embedding=query_embedding, top_k=3)
    return "\n".join(doc.content for doc in result["documents"])

from haystack.tools import Tool

memory_tool = Tool(
    name="engineering_memory",
    description="Retrieve past engineering issues and lessons learned.",
    parameters={
        "type": "object",
        "properties": {
            "query": {"type": "string", "description": "Query to search past issues"}
        },
        "required": ["query"],
        "additionalProperties": False
    },
    function=memory_function
)

standards_tool = Tool(
    name="coding_standards",
    description="Retrieve coding standards and best practices.",
    parameters={
        "type": "object",
        "properties": {
            "query": {"type": "string", "description": "Query to retrieve coding standards"}
        },
        "required": ["query"],
        "additionalProperties": False
    },
    function=standards_function
)

Step 4 — Create the Agent with LLM and Instructions

from haystack.components.generators.chat.openai import OpenAIChatGenerator
from haystack.components.agents import Agent

llm = OpenAIChatGenerator(model="gpt-4o-mini")

engineering_agent = Agent(
    chat_generator=llm,
    tools=[memory_tool, standards_tool],
    system_prompt="""
You are a senior engineering specialist.

Use tools when useful.

Goals:
• recall past issues
• apply coding standards
• prevent failures
• propose scalable solutions

Always consider:
- performance
- scalability
- reliability
- production readiness
"""
)

Step 5 — Execution with Streaming

from haystack.dataclasses import ChatMessage
from colorama import Fore, Style

def run(question: str):
    store_feedback(f"User asked: {question}")
    query_with_context = retrieve_memory(question) or question

    result = engineering_agent.run(
        messages=[ChatMessage.from_user(query_with_context)],
        memory_store_kwargs={"user_id": MEMORY_USER_ID},
        stream=True,
    )

    answer = result["last_message"].text
    store_feedback(f"Agent response: {answer}")
    print(answer)
    return answer

while True:
    user_input = input("Ask the engineering agent (or 'exit' to quit): ")
    if user_input.lower() == "exit":
        break
    if user_input:
        response = run(user_input)

2.5 MemGPT (Letta) — Agents with Memory Blocks

MemGPT (now Letta) is a memory-augmented framework designed to create stateful agents that can remember, learn, and improve over time.

Core Concept: Memory Separation

flowchart TD
    MEMGPT[MemGPT Agent] --> WM[Working Memory\nShort-term memory]
    MEMGPT --> LTM[Long-Term Memory\nLong-term memory]
    
    WM --> CTX[Active context\nToken window]
    LTM --> VEC[Vector storage\nRetrieved on demand]
    
    MEMGPT --> CTRL[Self-managed\nAutonomous control]
    CTRL --> |When context is full| ARCH[Archive to LTM]
    CTRL --> |When context needed| LOAD[Load from LTM]

Memory Blocks Architecture

Each memory is divided into labeled blocks:

Installation and Configuration

pip install letta-client colorama

from dotenv import load_dotenv
from letta_client import Letta
import os

load_dotenv()
client = Letta(api_key=os.getenv("LETTA_API_KEY"))

Create an Agent with Memory Blocks

messages = [
    {
        "label": "human",
        "value": "Name: Sandy\nOccupation: Software Engineer\nHobbies: Hiking, Cooking, Traveling"
    },
    {
        "label": "persona",
        "value": "You are self-improving superintelligence who acts as a senior AI engineer."
    }
]

agent_state = client.agents.create(
    model="openai/gpt-4o-mini",
    memory_blocks=messages
)

print(f"Agent created with ID: {agent_state.id}")

Sending Messages to the Agent

from colorama import Fore

def send_message(user_input):
    """Sends a message to the agent and returns the response."""
    return client.agents.messages.create(
        agent_id=agent_state.id,
        input=user_input
    )

while True:
    user_input = input("You: ")
    if user_input.lower() in ["exit", "quit"]:
        break

    response = send_message(user_input)
    print(f"\n{Fore.BLUE}🤖 Agent:")
    for message in response.messages:
        print(f"{message.content}{Fore.RESET}")

Autonomous and Self-Improving Operation

sequenceDiagram
    participant U as User
    participant A as MemGPT Agent
    participant WM as Working Memory
    participant LTM as Long-Term Memory

    U->>A: "Write a Python function to generate a random number between 1 and 10"
    A->>WM: Consults current context
    A->>LTM: Searches past interactions
    A-->>U: Responds with solution

    U->>A: "Update it for 1 to 100"
    A->>WM: Retrieves previous request
    Note over A: Auto-decision: update memory
    A->>LTM: Archives updated knowledge
    A-->>U: "Here is the updated Python function..."

The process is autonomous, self-managed, adaptive, and self-improving — the agent decides when and how to update its memory.

Resources: Letta Quickstart | Memory Blocks API

2.6 LlamaIndex — Multi-Agent AI Engineering Team

Goal: build a team of AI agents that can:

• Process code projects end-to-end
• Learn and improve over time with memory-augmented RAG
• Share context across interactions

Engineering Workflow Architecture

flowchart TD
    U[User Task]
    DEV[Developer Agent\nWrites the code]
    SR[Senior Developer Agent\nRefactors the code]
    REV[Reviewer Agent\nCode review]
    LEAD[Lead Engineer\nApproves / Rejects]
    SUM[Summarization & Pruning\nMemory optimization]
    MEM[Memory Storage\n+ Vector Index ChromaDB]

    U --> DEV --> SR --> REV --> LEAD --> SUM --> MEM
    LEAD -.->|Rejection| DEV

Installation

pip install llama-index llama-index-llms-openai chromadb tavily-python

Step 1 — LLM Configuration and Agent Prompts

from llama_index.llms.openai import OpenAI
from dotenv import load_dotenv

load_dotenv()
llm = OpenAI(model="gpt-4o-mini")

DEV_PROMPT = """
You are a software engineer.
Write clean, working code to implement the requested feature.
Focus on correctness and clarity.
Use:
1. Past issues from memory: {memory_context}
2. Coding standards: {standards_context}
3. Security best practices
"""

SENIOR_DEV_PROMPT = """
You are a senior software engineer and Refactor Agent.
Refactor, optimize, and improve existing code.
Use:
1. Past issues from memory: {memory_context}
2. Coding standards: {standards_context}
"""

REVIEW_PROMPT = """
You are a strict reviewer.
Review the code for bugs, security risks, performance issues.
Use:
1. Past issues from memory: {memory_context}
2. Coding standards: {standards_context}
"""

LEAD_DEV_PROMPT = """
You are the Lead Engineer.
Evaluate: requirements fulfillment, code quality, security.

OUTPUT FORMAT (MANDATORY):
    ✅ APPROVE  → ready for production
    ❌ REJECT   → major issues
    ⚠️ NEEDS IMPROVEMENT → revisions required
"""

Step 2 — RAG for Code Standards

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

docs = SimpleDirectoryReader("standards_docs").load_data()
index = VectorStoreIndex.from_documents(docs)
retriever = index.as_retriever(similarity_top_k=2)

def get_standards_context(query: str):
    nodes = retriever.retrieve(query or "python best practices security")
    return "\n".join(n.node.text for n in nodes)

Step 3 — Memory Storage with ChromaDB

from llama_index.core import VectorStoreIndex, Document
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core.storage.storage_context import StorageContext
import chromadb

client = chromadb.Client()
collection = client.get_or_create_collection("engineering_memory")

vector_store = ChromaVectorStore(chroma_collection=collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
memory_index = VectorStoreIndex([], storage_context=storage_context)
memory_retriever = memory_index.as_retriever(similarity_top_k=3)

def store_review_feedback(text: str, doc_id: str = "latest_feedback"):
    """Saves review feedback to memory."""
    doc = Document(text=text, doc_id=doc_id)
    memory_index.insert(doc)

def retrieve_past_issues(query: str = "past issues and feedback"):
    return memory_retriever.retrieve(query)

def retrieve_print_memory(query: str = "recent issues and feedback"):
    memories = memory_retriever.retrieve(query)
    return "\n".join([n.node.text for n in memories])

Step 4 — Memory Pruning and Summarization

async def summarize_memory(llm, keep_recent=5):
    messages = retrieve_past_issues("recent issues and feedback")
    if len(messages) <= keep_recent:
        return

    old_messages = messages[:-keep_recent]
    old_text = "\n".join([f"{m.role}: {m.content}" for m in old_messages])
    summary_prompt = f"Summarize these older messages concisely, keeping key decisions:\n{old_text}"

    summary_response = await llm.invoke({"messages": [{"role": "user", "content": summary_prompt}]})
    summary_text = summary_response["messages"][-1].content

    for _ in old_messages:
        memory_index.delete()
    memory_index.insert(summary_text)

async def prune_memory_if_needed(llm, max_tokens=4000, keep_recent=5):
    messages = retrieve_past_issues()
    current_tokens = sum(len(m.node.text) // 4 for m in messages)
    if current_tokens > max_tokens:
        print(f"⚠️ Memory token count {current_tokens} exceeds {max_tokens}. Summarizing...")
        await summarize_memory(llm, keep_recent)
    else:
        print(f"✅ Memory token count {current_tokens} within limit.")

Step 5 — Create Agents with FunctionAgent

from llama_index.core.agent.workflow import FunctionAgent

dev_agent = FunctionAgent(
    name="Developer",
    description="Writes code",
    system_prompt=DEV_PROMPT,
    llm=llm,
    tools=[search_web, record_notes],
    can_handoff_to=["Senior_Developer"],
    streaming=True,
)

senior_dev_agent = FunctionAgent(
    name="Senior_Developer",
    description="Refactors code",
    system_prompt=SENIOR_DEV_PROMPT,
    llm=llm,
    tools=[refactor_code, search_web, record_notes],
    can_handoff_to=["Reviewer"],
    streaming=True,
)

review_agent = FunctionAgent(
    name="Reviewer",
    description="Reviews code",
    system_prompt=REVIEW_PROMPT,
    llm=llm,
    tools=[write_review_report],
    can_handoff_to=["Lead"],
    streaming=True,
)

lead_agent = FunctionAgent(
    name="Lead",
    description="Approves code and review",
    system_prompt=LEAD_DEV_PROMPT,
    llm=llm,
    tools=[review_report],
    streaming=True,
)

Step 6 — Build and Execute the Workflow

from llama_index.core.agent.workflow import AgentWorkflow
from llama_index.core.workflow import Context

workflow = AgentWorkflow(
    agents=[dev_agent, senior_dev_agent, review_agent, lead_agent],
    root_agent="Developer",
)
ctx = Context(workflow)

Mock Test — Memory Pruning and Summarization

from collections import deque
from llama_index.core.llms import ChatMessage

class MockMemory:
    def __init__(self):
        self.messages = deque()

    def put_messages(self, messages):
        self.messages.extend(messages)

    def get(self):
        return list(self.messages)

    def clear(self):
        self.messages.clear()

    def __len__(self):
        return len(self.messages)

def mock_summarize_memory(keep_last_n=5):
    """Simulates summarization: keeps only the last N messages."""
    chat_history = memory.get()
    if len(chat_history) <= keep_last_n:
        return
    summarized = chat_history[-keep_last_n:]
    memory.clear()
    memory.put_messages(summarized)
    print(f"🔹 Memory summarized, kept last {keep_last_n} messages.")

def mock_prune_memory_if_needed(max_messages=5):
    """Triggers pruning if memory exceeds max_messages."""
    if len(memory.get()) > max_messages:
        mock_summarize_memory(keep_last_n=max_messages)

Module 3 — End-to-End Demo: Multi-Agent System with Persistent Memory

3.1 LangChain + LangGraph — Hybrid Memory Strategies

Final goal: build a collaborative, scalable multi-agent AI system with hybrid memory strategies (short-term + long-term).

Technology Stack

Complete Workflow Architecture

flowchart TD
    START([START]) --> TASK[User Task]
    TASK --> DEV[Developer Agent\nwrite_code]
    DEV --> REV[Reviewer Agent\nreview_code]
    REV --> LEAD[Lead Agent\napprove_code]
    LEAD --> ROUTER[Decision Router]
    
    ROUTER -->|approved| SUM[Summarizer Node]
    ROUTER -->|rejected| DEV

    SUM --> STM[Short Memory Maintenance\nPruning & Compression]
    STM --> LTM[Long Term Maintenance\nChromaDB Pruning]
    LTM --> END([END])

    subgraph Memory
        STM
        LTM
    end
    
    subgraph Agents
        DEV
        REV
        LEAD
    end

Workflow State (WorkflowState)

from typing import TypedDict

class WorkflowState(TypedDict, total=False):
    task: str
    code: str
    review: str
    approval: str
    config: dict
    short_term_memory: list[str]
    summary: str

3.2 Short-Term Memory (Thread-Level)

Short-term memory (thread-level persistence) allows agents to track multi-turn conversations.

Implementation with InMemorySaver

from langgraph.checkpoint.memory import InMemorySaver
import uuid

# Checkpointer for thread persistence
checkpointer = InMemorySaver()
shared_config = {"configurable": {"thread_id": "multi-agent-engineering-thread"}}

# Unique IDs per agent
dev_thread       = str(uuid.uuid4())
reviewer_thread  = str(uuid.uuid4())
lead_thread      = str(uuid.uuid4())

dev_config      = {"configurable": {"thread_id": dev_thread}}
reviewer_config = {"configurable": {"thread_id": reviewer_thread}}
lead_config     = {"configurable": {"thread_id": lead_thread}}

Conversation Buffer

MAX_SHORT_TERM = 6
convo_buffer: list[str] = []

def update_short_term_memory(state: WorkflowState, entry: str):
    """Adds an entry to short-term memory and maintains the size limit."""
    memory = state.get("short_term_memory", [])
    memory.append(entry)
    state["short_term_memory"] = memory[-MAX_SHORT_TERM:]

The simplest form: the entire conversation history is stored and passed as context on each interaction. Size-limited to avoid context bloat.

3.3 Long-Term Persistent Memory with ChromaDB

Long-term memory uses ChromaDB as a persistent vector store to store and retrieve knowledge via semantic search.

Vector Store Configuration

from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
import uuid

embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

# RAG — Code standards
vector_store = Chroma(
    collection_name="coding_standards",
    embedding_function=embeddings,
    persist_directory="./chroma_db"
)

# Long-term memory
memory_store = Chroma(
    collection_name="engineering_memory",
    embedding_function=embeddings,
    persist_directory="./chroma_db"
)

Memory Management Functions

def store_memory(text: str):
    """Stores memory in ChromaDB with metadata."""
    memory_store.add_texts(
        texts=[text],
        metadatas=[{"source": "agent_memory"}],
        ids=[str(uuid.uuid4())]
    )
    memory_store.persist()

def count_all_memories():
    """Counts total memories in ChromaDB."""
    data = memory_store.get()
    return len(data['ids'])

def retrieve_memory(query: str, k: int = 3):
    """Retrieves relevant memories via semantic search."""
    docs = memory_store.similarity_search(query, k=k)
    memories_count = count_all_memories()
    return [doc.page_content for doc in docs], memories_count

Loading Code Standards (RAG)

from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

loader = TextLoader("coding_standards.txt")
docs = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
all_splits = text_splitter.split_documents(docs)

_ = vector_store.add_documents(documents=all_splits)
vector_store.persist()

3.4 Performance and Production Scaling

To build a production-grade AI system, two key techniques are required:

Optimization Techniques

mindmap
  root((Production Optimization))
    Memory pruning
      Max document threshold
      Automatic compression
      Summary of old entries
    Summarization
      Short-term over 500 chars
      Long-term over max_docs
      LLM summary
    Caching
      SHA256 hash of message
      Before LLM call
      Returns cached response
    Middleware
      before_model check_cache
      after_model summarize_log_cache

Short-Term Memory Compression and Summarization

def summarize_compress_memory(text: str):
    """Summarizes text if it exceeds 500 characters."""
    if len(text) < 500:
        return text

    prompt = (
        "Summarize and compress the following engineering knowledge "
        "into reusable insights:\n\n" + text
    )
    return llm.invoke(prompt).content

def short_memory_maintenance_node(state: WorkflowState):
    """Maintains short-term memory: summarizes when the limit is reached."""
    memory = state.get("short_term_memory", [])

    if len(memory) >= MAX_MEMORY_BEFORE_SUMMARY:
        combined = "\n".join(memory)
        summary = summarize_compress_memory(combined)
        return {
            **state,
            "short_term_memory": [summary],
            "summary": summary
        }
    return state

Long-Term Memory Pruning

def prune_long_term_memory(max_docs: int = 100):
    """
    Prunes long-term memory in ChromaDB when it exceeds max_docs.
    Strategy:
        - Retrieve the oldest documents
        - Summarize them
        - Delete the originals
        - Store the compressed summary
    """
    all_items = memory_store._collection.get(include=["metadatas"])
    ids = all_items.get("ids", [])

    if len(ids) <= max_docs:
        return "No pruning required"

    excess = len(ids) - max_docs
    prune_ids = ids[:excess]

    prune_docs = memory_store._collection.get(ids=prune_ids, include=["documents"])
    documents = prune_docs.get("documents", [])
    combined_text = "\n".join(documents)

    # Summarize entries to be deleted
    summary = llm.invoke(
        "Summarize and compress the following engineering knowledge:\n" + combined_text
    ).content

    # Delete old entries
    memory_store._collection.delete(ids=prune_ids)

    # Store the summary
    memory_store.add_texts(
        texts=[summary],
        metadatas=[{"source": "memory_pruned_summary"}],
        ids=[str(uuid.uuid4())]
    )
    memory_store.persist()

    return f"Pruned {excess} memories and stored compressed summary."

Response Caching (Middleware)

import hashlib
from langchain_core.messages import AIMessage

response_cache = {}

def check_cache(state, runtime):
    """Returns the cached response if the request has already been processed."""
    if not state["messages"]:
        return None
    last_msg = state["messages"][-1].content
    key = hashlib.sha256(last_msg.encode()).hexdigest()
    if key in response_cache:
        print("⚡ Returning cached response")
        return {
            "messages": [AIMessage(content=response_cache[key])],
            "jump_to": "end"
        }
    return None

def summarize_log_cache(state, runtime):
    """Summarizes, stores in memory, logs and caches the response."""
    if not state["messages"]:
        return None

    last_content = state["messages"][-1].content
    summary = last_content

    update_short_term_memory(state, summary)
    store_memory(summary)

    key = hashlib.sha256(state["messages"][-1].content.encode()).hexdigest()
    response_cache[key] = last_content

    print(f"📝 Post-model state update:")
    print(f"Last message: {last_content}")
    print("-" * 60)
    return None

Workflow Summarizer Node

def summarizer_node(state: WorkflowState):
    """
    Generates a structured summary of the engineering workflow.
    This summary is optimized for long-term memory storage.
    """
    task = state.get("task", "")
    approval = state.get("approval", "")
    short_term_memory = state.get("short_term_memory", [])

    combined_context = f"""
    Task: {task}
    Final Decision: {approval}
    Recent Actions: {short_term_memory}
    """

    summary = llm.invoke(f"""
        Extract:
        - Key technical decisions
        - Design improvements
        - Bugs or risks identified
        - Performance or security considerations
        
        Compress into reusable long-term engineering insight.
        Context: {combined_context}
    """).content

    return {**state, "summary": summary}

Building the LangGraph Graph

from langgraph.graph import END, StateGraph

graph = StateGraph(WorkflowState)

# Agent nodes
graph.add_node("dev",    dev_node)
graph.add_node("review", review_node)
graph.add_node("lead",   lead_node)

# Memory nodes
graph.add_node("summarizer",               summarizer_node)
graph.add_node("short_memory_maintenance", short_memory_maintenance_node)
graph.add_node("long_term_maintenance",    long_term_maintenance_node)

# Routing node
graph.add_node("decision_router", decision_router_node)

graph.set_entry_point("dev")

# Main flow
graph.add_edge("dev",    "review")
graph.add_edge("review", "lead")
graph.add_edge("lead",   "decision_router")

# Conditional routing
graph.add_conditional_edges(
    "decision_router",
    lambda state: state.get("approval"),
    {
        "approved": "summarizer",
        "rejected": "dev"
    }
)

# Memory flow
graph.add_edge("summarizer",               "short_memory_maintenance")
graph.add_edge("short_memory_maintenance", "long_term_maintenance")
graph.add_edge("long_term_maintenance",    END)

workflow = graph.compile()

Running the Workflow

from colorama import Fore

def main():
    while True:
        task = input("\nTask (exit to quit): ")
        if task == "exit":
            break

        print(f"{Fore.CYAN}START: {task}\n")
        print("-> Running workflow... task in progress\n")

        result = workflow.invoke({
            "task": task,
            "config": shared_config
        })

        print("\nFINAL RESULT:")
        print(result["approval"])

if __name__ == "__main__":
    main()

3.5 Testing and Performance Monitoring

Unit Tests with Mocks

import warnings
from unittest.mock import MagicMock
from pytest import MonkeyPatch
from colorama import Fore, Style, init

init(autoreset=True)
warnings.filterwarnings("ignore")

# --- TEST 1: Engineering workflow ---
def test_engineering_workflow():
    task = "Implement user authentication with JWT"

    code_spec       = write_code.invoke({"feature_description": task})
    review_comments = review_code.invoke({"code_snippet": code_spec})
    approval_result = approve_code.invoke({"code_snippet": code_spec})

    assert isinstance(code_spec, str)
    assert isinstance(review_comments, str)
    assert isinstance(approval_result, str)

    print("✔ Engineering workflow executed successfully")

# --- TEST 2: Long-term memory pruning ---
def test_long_term_memory_pruning():
    # Insert 120 entries
    for i in range(120):
        store_memory(f"Memory entry {i}")

    before_pruning, memories_count = retrieve_memory("Memory entry")
    print(f"Total memories before pruning: {memories_count}")

    result = prune_long_term_memory(max_docs=100)

    after_pruning, memories_count = retrieve_memory("Memory entry")
    print(f"Total memories after pruning: {memories_count}")
    print(f"Pruning result: {result}")

    assert isinstance(result, str)
    print("✔ Memory pruning completed")

# --- TEST 3: Summarizer node ---
def test_summarizer_updates_state(monkeypatch):
    mock_llm = MagicMock()
    mock_llm.invoke.return_value.content = "Structured summary"
    monkeypatch.setattr("utils.llm", mock_llm)

    state = {
        "task": "Refactor code",
        "approval": "approved",
        "short_term_memory": ["Fix bug", "Improve performance"]
    }

    new_state = summarizer_node(state)

    assert "summary" in new_state
    assert new_state["summary"] == "Structured summary"
    print("✔ Summarizer node updated workflow state")

# --- TEST 4: Long-term maintenance ---
def test_long_term_pruning_called(monkeypatch):
    mock_prune = MagicMock(return_value="Pruned")
    monkeypatch.setattr("utils.prune_long_term_memory", mock_prune)

    state = {"task": "Deploy app"}
    long_term_maintenance_node(state)
    mock_prune.assert_called_once()
    print("✔ Long-term pruning was triggered as expected")

Run Tests

python test.py

Observability with LangSmith

flowchart LR
    APP[LangChain Application] -->|traces| LS[LangSmith]
    LS --> DASH[Dashboard]
    DASH --> T[Traces per execution]
    DASH --> L[Latency]
    DASH --> TOK[Token usage]
    DASH --> ERR[Errors]

LangSmith Configuration:

LANGSMITH_TRACING=true
LANGSMITH_API_KEY=<your-api-key>
OPENAI_API_KEY=<your-openai-key>
LANGSMITH_WORKSPACE_ID=<your-workspace-id>

LangSmith enables tracing, debugging and monitoring of each workflow execution in real time.

Code Standards Applied by the Agents

The agents in this course use a coding_standards.txt file as a RAG knowledge base. Key principles:

Final Project Configuration and Prerequisites (Module 3)

Installation

# Create and activate a virtual environment
python -m venv .venv

# Windows
.venv\Scripts\activate
# macOS
source .venv/bin/activate

# Install dependencies
pip install -r requirements.txt

Dependencies (requirements.txt)

python-dotenv==1.1.0
colorama==0.4.6
openai>=1.8.0
langsmith>=0.0.1
langchain>=1.0.0
langchain-community>=0.0.15
langchain-text-splitters>=0.1.9
langchain-core>=0.0.10
langchain-openai>=0.0.15
langgraph>=0.0.12
chromadb>=0.5.0
pytest>=8.0.0
pytest-mock>=3.12.0

Environment Variables

OPENAI_API_KEY=<your-openai-key>
LANGSMITH_TRACING=true
LANGSMITH_API_KEY=<your-langsmith-key>
LANGSMITH_WORKSPACE_ID=<your-workspace-id>

References and Resources

Frameworks

Vector Databases (Production)

Prerequisites

• Basics of vector databases and RAG (Retrieval-Augmented Generation)
• Experience with at least one agentic framework (LangChain, LlamaIndex, CrewAI)
• Intermediate / advanced Python

Search Terms

memory-augmented · persistent · agents · ai · orchestration · artificial · intelligence · generative · memory · agent · architecture · workflow · configuration · rag · conversation · installation · mem0 · standards · blocks · frameworks · functions · llm · management · pruning