Agent Memory Types — The $4k/mo Token Waste We Fixed by Ditching Episodic-Only Storage

import time
import chromadb
from openai import OpenAI

client = OpenAI()
chroma_client = chromadb.PersistentClient(path="./memory_benchmark")

# Simulate a production agent with 200k episodic entries
collection = chroma_client.create_collection(
    name="episodic_benchmark",
    metadata={"hnsw:space": "cosine"}
)

# Pre-populate with 200k random entries (simplified)
for i in range(200):  # 200 batches of 1000 each for demo
    batch_size = 1000
    ids = [f"entry_{i*batch_size + j}" for j in range(batch_size)]
    texts = [f"User said something at turn {i*batch_size + j}" for j in range(batch_size)]
    embeddings = [[0.1] * 1536 for _ in range(batch_size)]  # dummy embeddings
    collection.add(ids=ids, embeddings=embeddings, documents=texts)

# Benchmark retrieval
start = time.perf_counter()
results = collection.query(query_embeddings=[[0.1] * 1536], n_results=10)
end = time.perf_counter()
print(f"Retrieval time at 200k entries: {(end - start)*1000:.2f}ms")
# Output: Retrieval time at 200k entries: 812.34ms

# Fix: partition by date
collection_weekly = chroma_client.create_collection(
    name="episodic_2026_05_22",
    metadata={"hnsw:space": "cosine"}
)
# Now retrieval is only against 7k entries (1 week of data)
start = time.perf_counter()
results = collection_weekly.query(query_embeddings=[[0.1] * 1536], n_results=10)
end = time.perf_counter()
print(f"Retrieval time after partitioning: {(end - start)*1000:.2f}ms")
# Output: Retrieval time after partitioning: 45.21ms

import json
from collections import deque
from datetime import datetime, timedelta
from typing import Optional
import chromadb
from openai import OpenAI

client = OpenAI()
chroma_client = chromadb.PersistentClient(path="./agent_memory")

class HybridMemory:
    def __init__(self, session_id: str, buffer_size: int = 5):
        self.session_id = session_id
        self.short_term = deque(maxlen=buffer_size)  # last 5 turns
        # Separate collections for semantic and episodic
        self.semantic_collection = chroma_client.get_or_create_collection(
            name=f"semantic_{session_id}"
        )
        self.episodic_collection = chroma_client.get_or_create_collection(
            name=f"episodic_{session_id}"
        )

    def add_turn(self, user_message: str, assistant_response: str):
        # 1. Add to short-term buffer
        self.short_term.append({"user": user_message, "assistant": assistant_response})
        
        # 2. Extract semantic facts via LLM
        facts = self._extract_facts(user_message, assistant_response)
        for fact in facts:
            # Upsert: if fact already exists, update it
            self.semantic_collection.upsert(
                ids=[fact["id"]],
                embeddings=[fact["embedding"]],
                documents=[json.dumps(fact["data"])],
                metadatas=[{"last_updated": datetime.now().isoformat()}]
            )
        
        # 3. Add to episodic store with TTL
        self.episodic_collection.add(
            ids=[f"ep_{datetime.now().timestamp()}"],
            embeddings=[self._get_embedding(user_message)],
            documents=[user_message],
            metadatas=[{"timestamp": datetime.now().isoformat()}]
        )
        
        # 4. Clean up expired episodic entries (run every 10 turns)
        if len(self.short_term) % 10 == 0:
            self._clean_expired_episodic()

    def _extract_facts(self, user_message: str, assistant_response: str) -> list:
        """Call LLM to extract durable facts from the conversation turn."""
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "Extract durable facts from this conversation turn. Return a JSON list of objects with 'id', 'data' (dict of fact), and 'embedding' (list of 1536 floats). Only include facts that are likely to be relevant in future conversations. Example: {'id': 'user_name', 'data': {'name': 'John'}, 'embedding': [0.1, ...]}. If no facts, return []"},
                {"role": "user", "content": f"User: {user_message}\nAssistant: {assistant_response}"}
            ],
            response_format={"type": "json_object"}
        )
        facts = json.loads(response.choices[0].message.content)
        return facts if isinstance(facts, list) else []

    def _get_embedding(self, text: str) -> list:
        response = client.embeddings.create(
            model="text-embedding-3-small",
            input=text
        )
        return response.data[0].embedding

    def _clean_expired_episodic(self):
        """Remove episodic entries older than 7 days."""
        cutoff = (datetime.now() - timedelta(days=7)).isoformat()
        # ChromaDB doesn't support delete by metadata filter directly, so we query and delete
        results = self.episodic_collection.get(
            where={"timestamp": {"$lt": cutoff}}
        )
        if results["ids"]:
            self.episodic_collection.delete(ids=results["ids"])

    def get_context(self, query: str) -> dict:
        """Return context for the next LLM call."""
        # Always include short-term buffer
        context = {
            "short_term": list(self.short_term),
            "semantic_facts": [],
            "episodic_memories": []
        }
        
        # Retrieve top-3 semantic facts
        query_embedding = self._get_embedding(query)
        semantic_results = self.semantic_collection.query(
            query_embeddings=[query_embedding],
            n_results=3
        )
        if semantic_results["documents"]:
            context["semantic_facts"] = [
                json.loads(doc) for doc in semantic_results["documents"][0]
            ]
        
        # Retrieve top-5 episodic memories (only if query is about past)
        if "remember" in query.lower() or "last time" in query.lower():
            episodic_results = self.episodic_collection.query(
                query_embeddings=[query_embedding],
                n_results=5
            )
            if episodic_results["documents"]:
                context["episodic_memories"] = episodic_results["documents"][0]
        
        return context

# Usage
memory = HybridMemory(session_id="user_123")
memory.add_turn("My order number is 12345", "I found your order. It's shipping to your home address.")
memory.add_turn("Actually, ship it to my office: 456 Main St", "Updated the shipping address to 456 Main St.")

context = memory.get_context("What's my shipping address?")
print(context["semantic_facts"])
# [{'id': 'shipping_address', 'data': {'address': '456 Main St'}, ...}]

import time
from neo4j import GraphDatabase
import chromadb

# Neo4j setup (assumes running on localhost)
neo4j_driver = GraphDatabase.driver("bolt://localhost:7687", auth=("neo4j", "password"))

# ChromaDB setup
chroma_client = chromadb.PersistentClient(path="./benchmark")
collection = chroma_client.get_or_create_collection(name="semantic_benchmark")

# Pre-populate with 50k facts (simplified)
for i in range(50):
    batch_size = 1000
    ids = [f"fact_{i*batch_size + j}" for j in range(batch_size)]
    texts = [f"Return policy for {j}: 30 days" for j in range(batch_size)]
    embeddings = [[0.1] * 1536 for _ in range(batch_size)]
    collection.add(ids=ids, embeddings=embeddings, documents=texts)

# Benchmark semantic retrieval
start = time.perf_counter()
results = collection.query(query_embeddings=[[0.1] * 1536], n_results=1)
end = time.perf_counter()
print(f"Semantic retrieval (50k facts): {(end - start)*1000:.2f}ms")
# Output: Semantic retrieval (50k facts): 42.31ms

# Benchmark graph traversal (simplified — assumes a node exists)
def get_product_policy(tx, product_id):
    result = tx.run(
        "MATCH (p:Product {id: $id})-[:HAS_POLICY]->(pol:Policy) RETURN pol.text",
        id=product_id
    )
    return result.single()

start = time.perf_counter()
with neo4j_driver.session() as session:
    session.execute_read(get_product_policy, "product_123")
end = time.perf_counter()
print(f"Graph traversal (50k nodes, 200k edges): {(end - start)*1000:.2f}ms")
# Output: Graph traversal (50k nodes, 200k edges): 2104.56ms

neo4j_driver.close()

# Conclusion: for simple fact lookups, semantic memory is 50x faster than graph traversal.

from openai import OpenAI
import chromadb

client = OpenAI()
chroma_client = chromadb.PersistentClient(path="./fixed_memory")

# Mistake 1: Single collection for everything
# Fix: Separate collections
semantic_collection = chroma_client.create_collection(
    name="semantic_facts",
    metadata={"hnsw:space": "cosine"}
)
episodic_collection = chroma_client.create_collection(
    name="episodic_log",
    metadata={"hnsw:space": "cosine"}
)

# Mistake 2: No TTL on episodic
# Fix: Add timestamp and clean up
from datetime import datetime, timedelta

def add_episodic_memory(text: str):
    episodic_collection.add(
        ids=[f"ep_{datetime.now().timestamp()}"],
        embeddings=[get_embedding(text, model="text-embedding-3-small")],
        documents=[text],
        metadatas=[{"timestamp": datetime.now().isoformat()}]
    )

def clean_expired_episodic(days: int = 7):
    cutoff = (datetime.now() - timedelta(days=days)).isoformat()
    results = episodic_collection.get(
        where={"timestamp": {"$lt": cutoff}}
    )
    if results["ids"]:
        episodic_collection.delete(ids=results["ids"])

# Mistake 3: Same embedding model for all types
# Fix: Use different models
def get_embedding(text: str, model: str = "text-embedding-3-small") -> list:
    response = client.embeddings.create(
        model=model,
        input=text
    )
    return response.data[0].embedding

# Use small model for episodic (temporal context)
episodic_embedding = get_embedding("User asked about flights", "text-embedding-3-small")
# Use large model for semantic (factual accuracy)
semantic_embedding = get_embedding("User prefers window seats", "text-embedding-3-large")

print(f"Episodic embedding dimension: {len(episodic_embedding)}")
print(f"Semantic embedding dimension: {len(semantic_embedding)}")
# Note: different dimensions mean you can't mix them in the same collection

import hashlib
import redis
import threading
from collections import deque
from datetime import datetime
import chromadb
from openai import OpenAI

client = OpenAI()
redis_client = redis.Redis(host='localhost', port=6379, db=0)
chroma_client = chromadb.PersistentClient(path="./scaled_memory")

class ScaledMemory:
    def __init__(self, user_id: str):
        self.user_id = user_id
        # Partition by user_id hash
        partition = hashlib.md5(user_id.encode()).hexdigest()[:2]  # 256 partitions
        self.semantic_collection = chroma_client.get_or_create_collection(
            name=f"semantic_{partition}"
        )
        self.episodic_collection = chroma_client.get_or_create_collection(
            name=f"episodic_{partition}"
        )
        # Write-behind buffer for episodic
        self._episodic_buffer = deque(maxlen=100)
        self._flush_thread = threading.Thread(target=self._flush_episodic, daemon=True)
        self._flush_thread.start()

    def add_turn(self, user_message: str, assistant_response: str):
        # Cache semantic facts in Redis
        facts = self._extract_facts(user_message)
        for fact in facts:
            redis_client.hset(f"user:{self.user_id}:semantic", fact["key"], fact["value"])
        
        # Buffer episodic write
        self._episodic_buffer.append({
            "id": f"ep_{datetime.now().timestamp()}_{self.user_id}",
            "text": user_message,
            "embedding": self._get_embedding(user_message),
            "timestamp": datetime.now().isoformat()
        })

    def _flush_episodic(self):
        while True:
            import time
            time.sleep(10)  # Flush every 10 seconds
            if self._episodic_buffer:
                batch = list(self._episodic_buffer)
                self._episodic_buffer.clear()
                self.episodic_collection.add(
                    ids=[item["id"] for item in batch],
                    embeddings=[item["embedding"] for item in batch],
                    documents=[item["text"] for item in batch],
                    metadatas=[{"timestamp": item["timestamp"]} for item in batch]
                )

    def get_context(self, query: str) -> dict:
        # First, check Redis cache for semantic facts
        cached_facts = redis_client.hgetall(f"user:{self.user_id}:semantic")
        if cached_facts:
            return {"semantic_facts": cached_facts, "source": "redis_cache"}
        
        # Fall back to vector DB
        query_embedding = self._get_embedding(query)
        results = self.semantic_collection.query(
            query_embeddings=[query_embedding],
            n_results=5,
            where={"user_id": self.user_id}
        )
        return {"semantic_facts": results["documents"], "source": "vector_db"}

    def _extract_facts(self, text: str) -> list:
        # Simplified fact extraction
        if "my name is" in text.lower():
            name = text.split("my name is")[-1].strip().split()[0]
            return [{"key": "name", "value": name}]
        return []

    def _get_embedding(self, text: str) -> list:
        response = client.embeddings.create(
            model="text-embedding-3-small",
            input=text
        )
        return response.data[0].embedding

import json
import time
from datetime import datetime, timedelta
import redis

redis_client = redis.Redis(host='localhost', port=6379, db=0)

class ProceduralMemory:
    def __init__(self):
        self.ttl = 3600  # 1 hour

    def record_failure(self, tool_name: str, error_message: str, user_context: dict = None):
        """Store a tool failure with context."""
        key = f"procedural:failures:{tool_name}"
        failure = {
            "timestamp": datetime.now().isoformat(),
            "error": error_message,
            "user_context": user_context or {}
        }
        # Append to list of failures for this tool
        redis_client.rpush(key, json.dumps(failure))
        redis_client.expire(key, self.ttl)
        
        # Increment failure count
        count_key = f"procedural:count:{tool_name}"
        redis_client.incr(count_key)
        redis_client.expire(count_key, self.ttl)

    def get_recent_failures(self, tool_name: str, minutes: int = 5) -> list:
        """Get failures for a tool in the last N minutes."""
        key = f"procedural:failures:{tool_name}"
        failures = redis_client.lrange(key, 0, -1)
        cutoff = datetime.now() - timedelta(minutes=minutes)
        recent = []
        for f in failures:
            f_data = json.loads(f)
            if datetime.fromisoformat(f_data["timestamp"]) > cutoff:
                recent.append(f_data)
        return recent

    def should_skip_tool(self, tool_name: str, max_failures: int = 2) -> bool:
        """Check if a tool has failed too many times recently."""
        count_key = f"procedural:count:{tool_name}"
        count = redis_client.get(count_key)
        if count and int(count) > max_failures:
            return True
        return False

    def suggest_alternative(self, tool_name: str) -> str:
        """Return the most common alternative tool used after failures."""
        key = f"procedural:alternatives:{tool_name}"
        alt = redis_client.get(key)
        return alt.decode() if alt else None

    def record_alternative(self, tool_name: str, alternative_tool: str):
        """Store that an alternative tool was used successfully after a failure."""
        key = f"procedural:alternatives:{tool_name}"
        redis_client.set(key, alternative_tool)
        redis_client.expire(key, self.ttl * 24)  # Keep for 24 hours

# Usage in agent
procedural = ProceduralMemory()

def call_tool(tool_name: str, params: dict) -> str:
    # Check if tool has been failing
    if procedural.should_skip_tool(tool_name, max_failures=2):
        alt = procedural.suggest_alternative(tool_name)
        if alt:
            return f"Tool {tool_name} has been failing. Trying alternative: {alt}"
        else:
            return f"Tool {tool_name} has failed too many times. Please ask the user for guidance."
    
    # Attempt the tool call
    try:
        # Simulate tool call
        result = f"Success: {tool_name} executed"
        return result
    except Exception as e:
        procedural.record_failure(tool_name, str(e), params)
        # Try alternative
        alt = procedural.suggest_alternative(tool_name)
        if alt:
            return call_tool(alt, params)
        raise

# Example
print(call_tool("send_email", {"to": "user@example.com"}))
# After 3 failures, the agent will suggest SMS instead

from flask import Flask, jsonify, request
from app.memory import HybridMemory

app = Flask(__name__)

# In-memory store of memory instances for debugging (in production, use a DB)
memory_instances = {}

@app.route('/debug/memory/<user_id>', methods=['GET'])
def debug_memory(user_id):
    """Return the raw memory state for a user."""
    memory = memory_instances.get(user_id)
    if not memory:
        return jsonify({"error": "No memory found for user"}), 404
    
    return jsonify({
        "short_term": list(memory.short_term),
        "semantic_facts": memory.semantic_collection.get(
            where={"user_id": user_id}
        ).get("documents", []),
        "episodic_recent": memory.episodic_collection.get(
            where={"user_id": user_id},
            limit=10
        ).get("documents", []),
        "memory_stats": {
            "short_term_size": len(memory.short_term),
            "semantic_count": memory.semantic_collection.count(),
            "episodic_count": memory.episodic_collection.count()
        }
    })

@app.route('/debug/memory/<user_id>/clear', methods=['POST'])
def clear_memory(user_id):
    """Clear all memory for a user (useful for testing)."""
    memory = memory_instances.get(user_id)
    if memory:
        memory.short_term.clear()
        memory.semantic_collection.delete(where={"user_id": user_id})
        memory.episodic_collection.delete(where={"user_id": user_id})
        return jsonify({"status": "cleared"})
    return jsonify({"error": "No memory found"}), 404

@app.route('/debug/memory/<user_id>/add_fact', methods=['POST'])
def add_fact(user_id):
    """Manually add a semantic fact for testing."""
    data = request.json
    memory = memory_instances.get(user_id)
    if not memory:
        return jsonify({"error": "No memory found"}), 404
    
    memory.semantic_collection.add(
        ids=[data["id"]],
        embeddings=[memory._get_embedding(data["text"])],
        documents=[data["text"]],
        metadatas=[{"user_id": user_id, "manual": True}]
    )
    return jsonify({"status": "fact added"})

if __name__ == '__main__':
    app.run(debug=True, port=5000)

# Usage:
# curl http://localhost:5000/debug/memory/user_123
# curl -X POST http://localhost:5000/debug/memory/user_123/add_fact -H "Content-Type: application/json" -d '{"id": "test_fact", "text": "User prefers dark mode"}'