LLM Memory Management — How a $4k/month Token Leak Nearly Broke Our Chatbot

import json
from openai import OpenAI

client = OpenAI()

def extract_memories(conversation: list[dict]) -> list[dict]:
    """
    Extract structured memories from a conversation.
    Returns list of dicts with keys: content, importance, category.
    """
    prompt = f"""
Extract important facts about the user from this conversation.
Only extract facts that are explicitly stated by the user, not inferred.
Return a JSON array of objects with keys: content (str), importance (int 1-10), category (str: preference|background|goal|other).

Conversation:
{json.dumps(conversation, indent=2)}

Memories:
"""
    response = client.chat.completions.create(
        model="gpt-4o-mini",  # cheaper model for extraction
        messages=[
            {"role": "system", "content": "You are a memory extraction assistant. Be precise and conservative."},
            {"role": "user", "content": prompt}
        ],
        response_format={"type": "json_object"},
        temperature=0.0  # deterministic output
    )
    # Parse the JSON response
    raw = json.loads(response.choices[0].message.content)
    # Validate structure
    memories = raw.get("memories", raw if isinstance(raw, list) else [])
    for m in memories:
        if "content" not in m:
            m["content"] = m.get("fact", "")  # fallback for different key names
    return memories

import chromadb
from chromadb.utils import embedding_functions
import numpy as np
from typing import Optional

# Initialize ChromaDB with persistent storage
client = chromadb.PersistentClient(path="./memory_store")

# Use a local embedding model (no API calls for embedding)
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="all-MiniLM-L6-v2"
)

def get_or_create_collection(user_id: str):
    """Get or create a collection for a user. Each user gets their own collection."""
    collection_name = f"user_{user_id}"
    try:
        return client.get_collection(name=collection_name, embedding_function=sentence_transformer_ef)
    except ValueError:
        return client.create_collection(name=collection_name, embedding_function=sentence_transformer_ef)

def add_memory_with_dedup(user_id: str, memory: dict, similarity_threshold: float = 0.85):
    """
    Add a memory to the user's collection, but only if it's not a near-duplicate.
    """
    collection = get_or_create_collection(user_id)
    
    # Compute embedding for the new memory
    new_embedding = sentence_transformer_ef([memory["content"]])[0]
    
    # Query for similar existing memories
    results = collection.query(
        query_embeddings=[new_embedding],
        n_results=5,
        include=["distances", "metadatas"]
    )
    
    # Check if any existing memory is too similar
    if results["distances"] and len(results["distances"][0]) > 0:
        min_distance = min(results["distances"][0])
        if (1 - min_distance) > similarity_threshold:  # cosine distance to similarity
            print(f"Skipping duplicate: {memory['content']} (similarity: {1 - min_distance:.2f})")
            return
    
    # Add the new memory
    collection.add(
        documents=[memory["content"]],
        metadatas=[{
            "importance": memory.get("importance", 5),
            "category": memory.get("category", "other"),
            "timestamp": memory.get("timestamp", ""),
            "hit_count": 0
        }],
        ids=[f"mem_{user_id}_{hash(memory['content'])}"]
    )
    print(f"Added memory: {memory['content']}")

def should_use_memory(config: dict) -> bool:
    """
    Decision function for whether to enable long-term memory.
    Returns False if memory would add cost without benefit.
    """
    # If average session length is 1 interaction, skip memory
    if config["avg_session_length"] <= 1:
        return False
    
    # If latency budget is under 200ms, skip memory (adds 50-100ms)
    if config["latency_budget_ms"] < 200:
        return False
    
    # If dealing with PII and no GDPR compliance path, skip memory
    if config["has_pii"] and not config["gdpr_compliant"]:
        print("Warning: Memory with PII without GDPR compliance is risky")
        return False
    
    # If context window is large enough for full conversation, skip memory
    if config["max_tokens"] >= config["avg_conversation_tokens"] * 2:
        return False
    
    return True

import hashlib
from typing import List

# Shard configuration: map user_id to shard number
SHARD_COUNT = 4

def get_shard(user_id: str) -> int:
    """Consistent hashing to determine which ChromaDB instance to use."""
    hash_val = int(hashlib.md5(user_id.encode()).hexdigest(), 16)
    return hash_val % SHARD_COUNT

def batch_extract(conversation_buffer: List[dict]) -> List[dict]:
    """
    Extract memories from a batch of messages.
    Called every 5 messages or when buffer reaches 2000 tokens.
    """
    if len(conversation_buffer) < 3:
        return []  # Not enough context for extraction
    
    # Call extraction LLM (simplified)
    memories = extract_memories(conversation_buffer)
    
    # Clear buffer after extraction
    conversation_buffer.clear()
    return memories

def lazy_retrieve(user_id: str, current_embedding: List[float], threshold: float = 0.3) -> List[dict]:
    """
    Only retrieve memories if the current query is semantically different from the last one.
    """
    # Get last query embedding from cache (Redis)
    last_embedding = cache.get(f"last_embedding_{user_id}")
    
    if last_embedding is not None:
        # Compute cosine similarity
        similarity = np.dot(current_embedding, last_embedding) / (
            np.linalg.norm(current_embedding) * np.linalg.norm(last_embedding)
        )
        if similarity > (1 - threshold):  # If similar, skip retrieval
            return cache.get(f"cached_memories_{user_id}") or []
    
    # Retrieve fresh memories
    memories = retrieve_memories(user_id, current_embedding)
    
    # Update cache
    cache.set(f"last_embedding_{user_id}", current_embedding, ttl=300)
    cache.set(f"cached_memories_{user_id}", memories, ttl=300)
    
    return memories

# Mistake 1: Extracting after every turn (bad)
# This costs $0.0015 per call * 10K calls/day = $15/day
for message in conversation:
    memories = extract_memories([message])  # Too frequent!
    store_memories(user_id, memories)

# Fix: Batch extraction
conversation_buffer = []
for message in conversation:
    conversation_buffer.append(message)
    if len(conversation_buffer) >= 5:  # Extract every 5th message
        memories = extract_memories(conversation_buffer)
        store_memories(user_id, memories)
        conversation_buffer.clear()

# Mistake 2: No deduplication (bad)
# User ends up with 200 memories, 80% duplicates
def store_memories_no_dedup(user_id, memories):
    for m in memories:
        collection.add(documents=[m["content"]], ids=[f"mem_{hash(m['content'])}"])

# Fix: Add deduplication
from sklearn.metrics.pairwise import cosine_similarity

def store_memories_with_dedup(user_id, memories, collection):
    for m in memories:
        # Compute embedding
        emb = embed_model.encode([m["content"]])
        # Query for similar
        results = collection.query(query_embeddings=emb, n_results=1)
        if results["distances"] and len(results["distances"][0]) > 0:
            if (1 - results["distances"][0][0]) > 0.85:
                continue  # Skip duplicate
        collection.add(documents=[m["content"]], ids=[f"mem_{hash(m['content'])}"])

# Mistake 3: No embedding versioning (bad)
# Old embeddings become incompatible after model change
EMBEDDING_MODEL = "text-embedding-3-small"  # Hardcoded, no version tracking

# Fix: Version your embeddings
def add_memory_with_version(user_id, memory, model_version="v1"):
    collection.add(
        documents=[memory["content"]],
        metadatas={"embedding_model": model_version},
        ids=[f"mem_{hash(memory['content'])}"]
    )

# On model change, re-embed all memories
if current_model_version != stored_model_version:
    re_embed_all_memories(new_model)

# LangMem example (quick start)
from langmem import MemoryManager

manager = MemoryManager()
# This handles extraction and storage automatically, but hard to customize
manager.add_conversation(conversation, user_id="user_123")
memories = manager.get_relevant_memories("What does the user like?", user_id="user_123")

# Mem0 example (turnkey)
from mem0 import Memory

m = Memory()
# Black box: extraction, storage, retrieval all in one call
m.add("User likes Python", user_id="user_123")
memories = m.search("programming preferences", user_id="user_123")

# Custom system (full control)
# We control every step:
# 1. Extraction prompt
memories = extract_memories(conversation)  # custom prompt
# 2. Deduplication
for mem in memories:
    add_memory_with_dedup(user_id, mem)  # custom logic
# 3. Retrieval with re-ranking
memories = retrieve_and_rerank(user_id, query, top_k=5)  # custom strategy

from prometheus_client import Counter, Histogram, Gauge
import time

# Metrics
MEMORY_STORE_SIZE = Gauge('memory_store_size_per_user', 'Number of memories per user', ['user_id'])
EXTRACTION_CALLS = Counter('extraction_calls_total', 'Total extraction calls', ['user_id'])
INJECTION_TOKENS = Histogram('injection_tokens_per_prompt', 'Tokens injected per prompt', buckets=[100, 500, 1000, 2000, 5000])
RETRIEVAL_LATENCY = Histogram('retrieval_latency_seconds', 'Latency of memory retrieval', buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0])

def monitored_retrieve(user_id: str, query: str) -> list[dict]:
    start = time.time()
    memories = retrieve_memories(user_id, query)
    duration = time.time() - start
    RETRIEVAL_LATENCY.observe(duration)
    
    # Update gauge for memory store size
    MEMORY_STORE_SIZE.labels(user_id=user_id).set(len(memories))
    
    return memories

def monitored_extract(user_id: str, conversation: list[dict]) -> list[dict]:
    EXTRACTION_CALLS.labels(user_id=user_id).inc()
    memories = extract_memories(conversation)
    
    # Track injection tokens (estimate)
    total_tokens = sum(len(m["content"].split()) for m in memories) * 1.3  # rough estimate
    INJECTION_TOKENS.observe(total_tokens)
    
    return memories

// io.thecodeforge
import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;

public class EmbeddingMigration {
    private static final int MODEL_DIMENSION = 768; // Fixed for all-MiniLM-L6-v2
    private static final int BATCH_SIZE = 256;

    public void reindexOnModelChange(String oldModel, String newModel) {
        if (!oldModel.equals(newModel)) {
            ChromaEmbeddingStore store = new ChromaEmbeddingStore("http://localhost:8000");
            // Fetch all existing embeddings, re-embed with new model, replace
            store.collection().get().ids().forEach(id -> {
                List<Embedding> batch = fetchBatch(id, BATCH_SIZE);
                store.addAll(batch); // Replace old vectors
            });
        }
    }

    private List<Embedding> fetchBatch(String cursor, int size) { /* ... */ }
}

// io.thecodeforge
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;

public class DistilledQueryRetriever {
    private static final int TOP_K = 5;
    private static final int MIN_RESULTS = 3;
    private ChromaEmbeddingStore store;
    private OpenAiChatModel queryModel;

    public List<Memory> retrieve(String conversationTurn) {
        // Step 1: Distill the user's intent
        String query = queryModel.generate("Extract the user's intent from this turn (max 10 words): " + conversationTurn);
        // Step 2: Embed and search
        List<Memory> results = store.search(embed(query), TOP_K);
        // Step 3: Fallback if too few results
        if (results.size() < MIN_RESULTS) {
            results.addAll(store.search(embed(conversationTurn), TOP_K));
        }
        return results;
    }
}