AI Agents for Beginners - 13. Memory for AI Agents

AI Agents for Beginners - 13. Memory for AI Agents

Source: microsoft/ai-agents-for-beginners - 13-agent-memory

Understanding AI Agent Memory

Why is Memory Important?

Types of Memory

Process Flowchart
Mermaid
%%{init: {'look': 'handDrawn', 'flowchart': {'subGraphTitleMargin': {'top': 12, 'bottom': 6}}, 'themeVariables': {'clusterBkg': '#e8efff28', 'clusterBorder': '#aabbcc'}}}%%
flowchart LR
    Root["Memory Types"]
    Root --> WM["Working\nMemory"]
    Root --> ST["Short-term\nMemory"]
    Root --> LT["Long-term\nMemory"]
    Root --> PM["Persona\nMemory"]
    Root --> EM["Episodic\nMemory"]
    Root --> ENT["Entity\nMemory"]
    Root --> SR["Structured\nRAG"]

Implementing and Storing Memory

Specialized Memory Tools

Mem0

Mem0: Two-Stage Memory Pipeline
Mermaid
%%{init: {'look': 'handDrawn'}}%%
flowchart LR
  Msg["Agent Messages"] --> Mem0["Mem0 Service"]
  Mem0 --> LLM["LLM\nSummarize & Extract Conversation"]
  LLM --> Dec{"Add / Modify\n/ Delete?"}
  Dec --> Store["Hybrid Datastore\nVector + Graph + Key-Value"]
  Store -->|"recall"| Agent["Agent"]

Cognee

Cognee: Dual-Store Architecture
Mermaid
%%{init: {'look': 'handDrawn'}}%%
flowchart LR
  Data["Structured /\nUnstructured Data"]
  Data -->|"cognee.add()"| VEC["Vector Store\nSimilarity Search"]
  Data -->|"cognee.cognify()"| GR["Graph Store\nRelations & Structure"]
  VEC & GR --> HQ["Hybrid Query\n(GRAPH_COMPLETION)"]
  HQ -->|"relevant context"| Agent["Agent"]

Storing Memory with RAG

Storing Memory with RAG (Azure AI Search)
Mermaid
%%{init: {'look': 'handDrawn'}}%%
flowchart LR
  D1["Conversation History"] & D2["Emails"] & D3["Images"] --> AIS["Azure AI Search"]
  AIS -->|"Structure & Index"| IDX["Structured Index\nDense Information Extraction"]
  IDX -->|"Precise Query"| Agent["Agent\nSuperhuman Precision & Recall"]

Working Memory with Sessions

setup.py
python

working_memory.py
python

new_session.py
python

Process Flowchart
Mermaid
%%{init: {'look': 'handDrawn'}}%%
flowchart LR
    U1["User: 'Love beaches, budget $3000'"] -->|"session"| A["TravelMemoryAgent"]
    U2["User: 'What did I say my budget was?'"] -->|"same session"| A
    A -->|"Reference previous turn"| R["Accurately recalls $3000"]

    U3["User: 'What is my budget?'"] -->|"new_session"| B["TravelMemoryAgent\n(New session)"]
    B -->|"No previous conversation"| NR["Unknown — New session"]

Long-Term Memory Pattern

memory_tools.py
python

Process Flowchart
Mermaid
%%{init: {'look': 'handDrawn'}}%%
flowchart LR
    subgraph Agent["MAF Agent (LLM)"]
        direction TB
        LLM["LLM"]
    end
    subgraph Tools["@tool functions"]
        direction TB
        SP["save_preference()"]
        GP["get_preferences()"]
        SH["search_hotels()"]
    end
    subgraph Store["Preference Store (dict)"]
        direction TB
        PS["user_id → [prefs]"]
    end
    LLM -->|"call"| Tools
    SP & GP --> PS
    PS -->|"persists across\nsessions"| PS
    SH -->|"hotel DB"| LLM

Scenario 1 — First-time user books an anniversary trip

scenario1_session1.py
python

scenario1_session1_followup.py
python

Scenario 2 — Sarah returns weeks later

scenario2_session2.py
python

scenario2_followup.py
python

Process Flowchart
Mermaid
%%{init: {'look': 'handDrawn'}}%%
flowchart TD
    subgraph S1["Session 1 — First Visit"]
        U1["Sarah: 'Romantic, spa, accessible\nBudget $700-800/night'"]
        U2["Sarah: 'Vegetarian, nut allergy'"]
        U1 & U2 -->|"save_preference()"| STORE["preference_store\n['romantic', 'spa', 'accessible',\n 'vegetarian', 'nut allergy']"]
    end
    subgraph S2["Session 2 — Returns Weeks Later"]
        U3["Sarah: 'Recommend a good hotel'\n(New session, no working memory)"]
        U3 -->|"get_preferences()"| STORE
        STORE -->|"Return saved preferences"| A2["TravelBookingAssistant\nPersonalized recommendations"]
    end

Storing Memory with Cognee Knowledge Graphs

Part 1 — Building the Knowledge Base

cognee_setup.py
python

cognee_data.py
python

cognee_ingest.py
python

Process Flowchart
Mermaid
%%{init: {'look': 'handDrawn'}}%%
flowchart TD
    D1["Developer Profile\n(developer_intro)"]
    D2["Past Conversations\n(human_agent_conversations)"]
    D3["Python Principles\n(python_zen_principles)"]
    D1 & D2 -->|"cognee.add(node_set=['developer_data'])"| ADD["cognee.add()"]
    D3 -->|"cognee.add(node_set=['principles_data'])"| ADD
    ADD -->|"cognee.cognify()"| KG["Knowledge Graph\n+ Vector Embeddings"]
    KG -->|"cognee.memify()"| Rules["Enriched Memory\nPatterns, Rules & Relations"]
    KG -->|"visualize_graph()"| HTML["cognee_graph.html\nInteractive Visualization"]

Part 2 — MAF Agent with Cognee Tools

cognee_tools.py
python

cognee_agent.py
python

Process Flowchart
Mermaid
%%{init: {'look': 'handDrawn'}}%%
flowchart LR
    subgraph MAF["MAF Agent (AgentSession)"]
        direction TB
        LLM2["CodingAssistant\n(LLM)"]
        Sess["create_session()\nWorking Memory"]
        LLM2 <--> Sess
    end
    subgraph CogneeTools["@tool functions"]
        direction TB
        SK["search_knowledge()\nGRAPH_COMPLETION"]
        SP2["search_principles()\nprinciples_data only"]
    end
    subgraph KG2["Cognee Knowledge Graph"]
        direction TB
        VEC["Vector Embeddings"]
        GR["Graph Relations"]
    end
    LLM2 -->|"call"| CogneeTools
    CogneeTools -->|"cognee.search()"| KG2
    KG2 -->|"Return relevant knowledge"| LLM2

Making AI Agents Self-Improve

1. Identify valuable information: Determines which parts of the conversation are worth storing as general knowledge or specific user preferences
2. Extract and summarize: Extracts and summarizes essential learnings or preferences from the conversation
3. Store in a knowledge base: Stores the extracted information in a vector database or similar storage
4. Augment future queries: When the user initiates a new query, retrieves relevant stored information and appends it to the user prompt (similar to RAG)

Process Flowchart
Mermaid
%%{init: {'look': 'handDrawn'}}%%
flowchart TD
    User["User"] -->|"New query"| KA["Knowledge Agent\n(Observer)"]
    KA -->|"Retrieve relevant knowledge\n(RAG)"| VDB["Vector DB\nKnowledge Base"]
    VDB -->|"Add stored context"| PA["Primary Agent"]
    PA -->|"Response"| User
    PA -->|"Observe conversation"| KA
    KA -->|"Valuable info\nExtract, summarize & store"| VDB

Optimizations for Memory

Summary

Process Flowchart
Mermaid
%%{init: {'look': 'handDrawn'}}%%
flowchart LR
    Root["Memory for\nAI Agents"]
    Root --> Types["Memory Types"]
    Root --> Tools["Memory Tools"]
    Root --> Pattern["Self-Improving\nPattern"]

    Types --> T1["Working\ncreate_session()"]
    Types --> T2["Short-term\nIn-session context"]
    Types --> T3["Long-term\n@tool persistent store"]
    Types --> T4["Persona·Episodic\nEntity·Structured RAG"]

    Tools --> TL1["Mem0\nExtract + Update Pipeline"]
    Tools --> TL2["Cognee\nKnowledge Graph + Vector"]
    Tools --> TL3["Azure AI Search\nStructured RAG Backend"]

    Pattern --> P1["Knowledge Agent\nObserve & Extract Conversation"]
    Pattern --> P2["RAG Augmentation\nAdd context to future queries"]

• Provides working memory via agent.create_session() — the agent sees the full conversation history within the session
• New sessions lose context — without long-term memory, the agent cannot recall past conversations
• @tool functions bridge the gap — allowing the agent to store and retrieve information from a persistent store
• Mem0: Transitions stateless agents into stateful ones using a two-stage pipeline (extraction → update)
• Cognee: Relation-aware long-term memory via knowledge graphs + vector embeddings — understands relationships between concepts using GRAPH_COMPLETION search
• Forms a self-improving loop where the more preferences and history are stored, the more refined the agent's recommendations become