mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
9c9bb30ded
Features: - with_fast_memory(): Zero-latency inline extraction (Letta-style) - Memory extracted as part of LLM response, no extra API calls - Semantic retrieval with local embeddings (sub-50ms) - with_memory(): Background extraction for non-blocking memory - SQLite + FTS5 storage with vector similarity search - Multi-user isolation by user_id Memory enables temporal compression - extract key facts instead of carrying full conversation history (4000 tokens → 50 tokens). Includes: - Comprehensive test suite (71 new tests) - Documentation (docs/memory.md) - Benchmark examples comparing approaches - E2E test with LLM-as-judge evaluation
595 lines
21 KiB
Python
595 lines
21 KiB
Python
#!/usr/bin/env python3
|
|
"""End-to-end memory system test with LLM-as-judge evaluation.
|
|
|
|
This program tests the memory extraction and retrieval system with:
|
|
1. Multi-turn conversations containing embedded memory nuggets
|
|
2. Real OpenAI API calls for extraction and conversation
|
|
3. LLM-as-judge evaluation of memory quality
|
|
|
|
Usage:
|
|
export OPENAI_API_KEY="sk-..."
|
|
python examples/memory_e2e_test.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
# Add parent to path for local development
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from openai import OpenAI
|
|
|
|
from headroom.memory import with_memory
|
|
from headroom.memory.store import SQLiteMemoryStore
|
|
|
|
# =============================================================================
|
|
# Test Scenarios - Conversations with embedded memory nuggets
|
|
# =============================================================================
|
|
|
|
|
|
@dataclass
|
|
class MemoryNugget:
|
|
"""A fact that should be remembered from the conversation."""
|
|
|
|
content: str
|
|
category: str # preference, fact, context
|
|
importance: float # 0.0-1.0
|
|
turn_index: int # Which turn contains this nugget
|
|
|
|
|
|
@dataclass
|
|
class TestScenario:
|
|
"""A test scenario with conversation and expected memories."""
|
|
|
|
name: str
|
|
description: str
|
|
entity_id: str # user_id or agent_id
|
|
conversation: list[tuple[str, str]] # List of (user_msg, expected_response_topic)
|
|
expected_nuggets: list[MemoryNugget]
|
|
retrieval_queries: list[tuple[str, list[str]]] # (query, expected_keywords_in_memory)
|
|
|
|
|
|
# Scenario 1: Software Developer User
|
|
DEVELOPER_SCENARIO = TestScenario(
|
|
name="software_developer",
|
|
description="A software developer discussing their preferences and projects",
|
|
entity_id="dev_alice",
|
|
conversation=[
|
|
# Turn 0 - Preference nugget
|
|
(
|
|
"Hi! I'm starting a new backend project. I strongly prefer Python over "
|
|
"JavaScript for backend work because of its cleaner syntax.",
|
|
"backend_project_advice",
|
|
),
|
|
# Turn 1 - Fact nugget
|
|
(
|
|
"Good point. I work at a fintech startup called PayFlow where we handle "
|
|
"high-volume payment processing.",
|
|
"fintech_architecture",
|
|
),
|
|
# Turn 2 - Context nugget
|
|
(
|
|
"We're currently migrating from a monolith to microservices. It's been "
|
|
"challenging but necessary for scale.",
|
|
"migration_advice",
|
|
),
|
|
# Turn 3 - Preference nugget
|
|
(
|
|
"For databases, I always use PostgreSQL. I've tried MongoDB but found "
|
|
"relational databases more reliable for financial data.",
|
|
"database_choice",
|
|
),
|
|
# Turn 4 - Casual (no nugget expected)
|
|
("Thanks for all the help today!", "closing"),
|
|
],
|
|
expected_nuggets=[
|
|
MemoryNugget("Prefers Python over JavaScript for backend", "preference", 0.8, 0),
|
|
MemoryNugget("Works at fintech startup PayFlow", "fact", 0.9, 1),
|
|
MemoryNugget("Handles high-volume payment processing", "fact", 0.7, 1),
|
|
MemoryNugget("Migrating from monolith to microservices", "context", 0.8, 2),
|
|
MemoryNugget("Prefers PostgreSQL over MongoDB", "preference", 0.8, 3),
|
|
MemoryNugget("Works with financial data", "fact", 0.7, 3),
|
|
],
|
|
retrieval_queries=[
|
|
# FTS5 is keyword-based, so queries must contain matching words
|
|
("Python backend", ["Python", "backend"]),
|
|
("PostgreSQL database", ["PostgreSQL", "database"]),
|
|
("PayFlow fintech", ["PayFlow", "fintech"]),
|
|
("microservices migration", ["microservices", "monolith"]),
|
|
],
|
|
)
|
|
|
|
# Scenario 2: AI Research Agent
|
|
AGENT_SCENARIO = TestScenario(
|
|
name="research_agent",
|
|
description="An AI agent discussing its capabilities and constraints",
|
|
entity_id="agent_researcher",
|
|
conversation=[
|
|
# Turn 0 - Capability fact
|
|
(
|
|
"I'm Agent-7, specialized in scientific literature analysis. I can process "
|
|
"up to 50 papers per hour and identify cross-domain connections.",
|
|
"agent_intro",
|
|
),
|
|
# Turn 1 - Constraint context
|
|
(
|
|
"My knowledge cutoff is March 2025, so I may not have the latest preprints. "
|
|
"I work best with structured abstracts.",
|
|
"limitations",
|
|
),
|
|
# Turn 2 - Preference
|
|
(
|
|
"When summarizing papers, I prefer to use the IMRaD structure - Introduction, "
|
|
"Methods, Results, and Discussion. It's more systematic.",
|
|
"summary_format",
|
|
),
|
|
# Turn 3 - Configuration fact
|
|
(
|
|
"I'm currently configured to prioritize papers from Nature, Science, and Cell "
|
|
"journals, with a citation threshold of 10+.",
|
|
"configuration",
|
|
),
|
|
# Turn 4 - Context about ongoing task
|
|
(
|
|
"Right now I'm tracking the emerging field of mechanistic interpretability "
|
|
"in neural networks. It's my primary research focus.",
|
|
"current_focus",
|
|
),
|
|
],
|
|
expected_nuggets=[
|
|
MemoryNugget("Agent-7 specialized in scientific literature", "fact", 0.9, 0),
|
|
MemoryNugget("Can process 50 papers per hour", "fact", 0.7, 0),
|
|
MemoryNugget("Knowledge cutoff March 2025", "context", 0.8, 1),
|
|
MemoryNugget("Prefers IMRaD structure for summaries", "preference", 0.8, 2),
|
|
MemoryNugget("Prioritizes Nature, Science, Cell journals", "fact", 0.7, 3),
|
|
MemoryNugget("Citation threshold of 10+", "fact", 0.6, 3),
|
|
MemoryNugget("Focus on mechanistic interpretability", "context", 0.9, 4),
|
|
],
|
|
retrieval_queries=[
|
|
# FTS5 keyword-based queries
|
|
("scientific papers analysis", ["papers", "scientific", "literature"]),
|
|
("IMRaD summary structure", ["IMRaD", "structure"]),
|
|
("Nature Science Cell journals", ["Nature", "Science", "Cell"]),
|
|
("mechanistic interpretability neural", ["mechanistic", "interpretability"]),
|
|
],
|
|
)
|
|
|
|
# Scenario 3: Multi-session customer
|
|
CUSTOMER_SCENARIO = TestScenario(
|
|
name="returning_customer",
|
|
description="A customer across multiple support interactions",
|
|
entity_id="customer_bob",
|
|
conversation=[
|
|
# Turn 0 - Account fact
|
|
(
|
|
"Hi, I'm Bob Chen, account number AC-789456. I've been a premium member since 2021.",
|
|
"account_lookup",
|
|
),
|
|
# Turn 1 - Preference
|
|
(
|
|
"Please always contact me via email at bob.chen@email.com, never by phone. "
|
|
"I work odd hours as a night shift nurse.",
|
|
"contact_preference",
|
|
),
|
|
# Turn 2 - Issue context
|
|
(
|
|
"I've had recurring issues with billing - this is the third time this month "
|
|
"I've been double-charged.",
|
|
"billing_issue",
|
|
),
|
|
# Turn 3 - Product preference
|
|
(
|
|
"I mainly use your enterprise plan for the API access. The dashboard features "
|
|
"I never touch.",
|
|
"usage_pattern",
|
|
),
|
|
],
|
|
expected_nuggets=[
|
|
MemoryNugget("Bob Chen, account AC-789456", "fact", 0.9, 0),
|
|
MemoryNugget("Premium member since 2021", "fact", 0.7, 0),
|
|
MemoryNugget("Prefers email contact, never phone", "preference", 0.9, 1),
|
|
MemoryNugget("Works as night shift nurse", "fact", 0.6, 1),
|
|
MemoryNugget("Recurring billing/double-charge issues", "context", 0.8, 2),
|
|
MemoryNugget("Uses enterprise plan for API access", "fact", 0.7, 3),
|
|
],
|
|
retrieval_queries=[
|
|
# FTS5 keyword-based queries
|
|
("Bob Chen account premium", ["Bob", "account", "premium"]),
|
|
("email contact phone", ["email", "phone"]),
|
|
("billing double charged", ["billing", "charged"]),
|
|
("enterprise API plan", ["API", "enterprise"]),
|
|
],
|
|
)
|
|
|
|
ALL_SCENARIOS = [DEVELOPER_SCENARIO, AGENT_SCENARIO, CUSTOMER_SCENARIO]
|
|
|
|
|
|
# =============================================================================
|
|
# Conversation Simulator
|
|
# =============================================================================
|
|
|
|
|
|
class ConversationSimulator:
|
|
"""Simulates realistic conversations using OpenAI."""
|
|
|
|
def __init__(self, client: OpenAI):
|
|
self.client = client
|
|
self.model = "gpt-4o-mini"
|
|
|
|
def generate_response(self, user_message: str, topic_hint: str) -> str:
|
|
"""Generate a realistic assistant response."""
|
|
response = self.client.chat.completions.create(
|
|
model=self.model,
|
|
messages=[
|
|
{
|
|
"role": "system",
|
|
"content": (
|
|
"You are a helpful assistant. Respond naturally and briefly "
|
|
"to the user's message. Keep responses under 100 words."
|
|
),
|
|
},
|
|
{"role": "user", "content": user_message},
|
|
],
|
|
max_tokens=150,
|
|
temperature=0.7,
|
|
)
|
|
return response.choices[0].message.content
|
|
|
|
|
|
# =============================================================================
|
|
# LLM-as-Judge Evaluator
|
|
# =============================================================================
|
|
|
|
|
|
@dataclass
|
|
class EvaluationResult:
|
|
"""Result of LLM judge evaluation."""
|
|
|
|
scenario_name: str
|
|
extraction_score: float # 0-1, how many expected nuggets were captured
|
|
retrieval_score: float # 0-1, how well queries retrieved relevant memories
|
|
overall_score: float
|
|
extracted_memories: list[dict]
|
|
missing_nuggets: list[str]
|
|
retrieval_results: list[dict]
|
|
judge_reasoning: str
|
|
|
|
|
|
class LLMJudge:
|
|
"""Uses LLM to evaluate memory extraction and retrieval quality."""
|
|
|
|
def __init__(self, client: OpenAI):
|
|
self.client = client
|
|
self.model = "gpt-4o" # Use stronger model for judging
|
|
|
|
def evaluate_extraction(
|
|
self,
|
|
scenario: TestScenario,
|
|
extracted_memories: list[dict],
|
|
) -> tuple[float, list[str], str]:
|
|
"""Evaluate if extracted memories capture expected nuggets."""
|
|
|
|
prompt = f"""You are evaluating a memory extraction system.
|
|
|
|
The system processed this conversation and extracted memories.
|
|
|
|
## Expected Information to Remember:
|
|
{json.dumps([{"content": n.content, "category": n.category, "importance": n.importance} for n in scenario.expected_nuggets], indent=2)}
|
|
|
|
## Actually Extracted Memories:
|
|
{json.dumps(extracted_memories, indent=2)}
|
|
|
|
## Evaluation Task:
|
|
1. For each expected nugget, determine if it was captured (exact match not required - semantic similarity counts)
|
|
2. Calculate what percentage of expected nuggets were captured
|
|
3. Identify which nuggets were MISSING
|
|
|
|
Return a JSON object:
|
|
{{
|
|
"captured_count": <number of expected nuggets that were captured>,
|
|
"total_expected": {len(scenario.expected_nuggets)},
|
|
"score": <0.0 to 1.0>,
|
|
"missing_nuggets": ["list of expected nuggets that were not captured"],
|
|
"reasoning": "Brief explanation of the evaluation"
|
|
}}
|
|
|
|
Return ONLY valid JSON."""
|
|
|
|
response = self.client.chat.completions.create(
|
|
model=self.model,
|
|
messages=[{"role": "user", "content": prompt}],
|
|
response_format={"type": "json_object"},
|
|
temperature=0.0,
|
|
)
|
|
|
|
result = json.loads(response.choices[0].message.content)
|
|
return (
|
|
result.get("score", 0.0),
|
|
result.get("missing_nuggets", []),
|
|
result.get("reasoning", ""),
|
|
)
|
|
|
|
def evaluate_retrieval(
|
|
self,
|
|
query: str,
|
|
expected_keywords: list[str],
|
|
retrieved_memories: list[dict],
|
|
) -> tuple[float, str]:
|
|
"""Evaluate if retrieval returned relevant memories."""
|
|
|
|
prompt = f"""You are evaluating a memory retrieval system.
|
|
|
|
## Query: "{query}"
|
|
|
|
## Expected Keywords in Results: {expected_keywords}
|
|
|
|
## Retrieved Memories:
|
|
{json.dumps(retrieved_memories, indent=2)}
|
|
|
|
## Evaluation Task:
|
|
Determine if the retrieved memories are relevant to the query and contain the expected information.
|
|
|
|
Return a JSON object:
|
|
{{
|
|
"score": <0.0 to 1.0>,
|
|
"keywords_found": ["list of expected keywords that appeared in results"],
|
|
"reasoning": "Brief explanation"
|
|
}}
|
|
|
|
Return ONLY valid JSON."""
|
|
|
|
response = self.client.chat.completions.create(
|
|
model=self.model,
|
|
messages=[{"role": "user", "content": prompt}],
|
|
response_format={"type": "json_object"},
|
|
temperature=0.0,
|
|
)
|
|
|
|
result = json.loads(response.choices[0].message.content)
|
|
return result.get("score", 0.0), result.get("reasoning", "")
|
|
|
|
|
|
# =============================================================================
|
|
# Main Test Runner
|
|
# =============================================================================
|
|
|
|
|
|
class MemoryE2ETest:
|
|
"""End-to-end test runner for the memory system."""
|
|
|
|
def __init__(self, api_key: str):
|
|
self.raw_client = OpenAI(api_key=api_key)
|
|
self.simulator = ConversationSimulator(self.raw_client)
|
|
self.judge = LLMJudge(self.raw_client)
|
|
self.results: list[EvaluationResult] = []
|
|
|
|
def run_scenario(self, scenario: TestScenario, db_path: Path) -> EvaluationResult:
|
|
"""Run a complete test scenario."""
|
|
print(f"\n{'=' * 60}")
|
|
print(f"Running scenario: {scenario.name}")
|
|
print(f"Description: {scenario.description}")
|
|
print(f"{'=' * 60}")
|
|
|
|
# Create memory-wrapped client
|
|
store = SQLiteMemoryStore(db_path)
|
|
memory_client = with_memory(
|
|
self.raw_client,
|
|
user_id=scenario.entity_id,
|
|
db_path=db_path,
|
|
_store=store,
|
|
)
|
|
|
|
# Run conversation through memory-wrapped client
|
|
print(f"\n--- Running {len(scenario.conversation)} conversation turns ---")
|
|
for i, (user_msg, _topic_hint) in enumerate(scenario.conversation):
|
|
print(f"\nTurn {i + 1}:")
|
|
print(f" User: {user_msg[:80]}...")
|
|
|
|
# Send through memory-wrapped client - this:
|
|
# 1. Retrieves relevant memories (if any)
|
|
# 2. Injects them into user message
|
|
# 3. Calls the actual API
|
|
# 4. Queues extraction of (original_query, response) in background
|
|
response = memory_client.chat.completions.create(
|
|
model="gpt-4o-mini",
|
|
messages=[
|
|
{"role": "user", "content": user_msg},
|
|
],
|
|
)
|
|
assistant_response = response.choices[0].message.content
|
|
print(f" Assistant: {assistant_response[:80]}...")
|
|
|
|
# Small delay to allow background extraction
|
|
time.sleep(0.5)
|
|
|
|
# Flush all extractions (force immediate processing)
|
|
print("\n--- Flushing background extractions ---")
|
|
start_time = time.time()
|
|
success = memory_client.flush_extractions(timeout=60.0)
|
|
elapsed = time.time() - start_time
|
|
|
|
if success:
|
|
print(f" All extractions complete in {elapsed:.1f}s")
|
|
else:
|
|
print(f" WARNING: Flush timed out after {elapsed:.1f}s")
|
|
pending = store.get_pending_extractions(limit=10)
|
|
still_pending = [p for p in pending if p.status == "pending"]
|
|
if still_pending:
|
|
print(f" {len(still_pending)} extractions still pending")
|
|
for p in still_pending[:2]:
|
|
print(f" - Query: {p.query[:50]}...")
|
|
|
|
# If no memories extracted, try direct extraction for debugging
|
|
all_memories = store.get_all(scenario.entity_id)
|
|
if not all_memories and scenario.conversation:
|
|
print("\n DEBUG: Attempting direct extraction for first turn...")
|
|
from headroom.memory.extractor import MemoryExtractor
|
|
|
|
extractor = MemoryExtractor(self.raw_client)
|
|
first_query, _ = scenario.conversation[0]
|
|
test_response = "Acknowledged, I understand."
|
|
direct_memories = extractor.extract(first_query, test_response)
|
|
print(f" DEBUG: Direct extraction got {len(direct_memories)} memories")
|
|
for m in direct_memories[:3]:
|
|
print(f" - [{m.category}] {m.content[:50]}...")
|
|
|
|
# Get all extracted memories (refresh)
|
|
extracted = [
|
|
{"content": m.content, "category": m.category, "importance": m.importance}
|
|
for m in all_memories
|
|
]
|
|
print(f"\n--- Extracted {len(extracted)} memories ---")
|
|
for m in extracted:
|
|
print(f" [{m['category']}] {m['content'][:60]}...")
|
|
|
|
# Evaluate extraction quality
|
|
print("\n--- Evaluating extraction quality ---")
|
|
extraction_score, missing, extraction_reasoning = self.judge.evaluate_extraction(
|
|
scenario, extracted
|
|
)
|
|
print(f" Extraction Score: {extraction_score:.2f}")
|
|
if missing:
|
|
print(f" Missing nuggets: {len(missing)}")
|
|
for m in missing[:3]:
|
|
print(f" - {m[:60]}...")
|
|
|
|
# Test retrieval queries
|
|
print("\n--- Testing retrieval queries ---")
|
|
retrieval_results = []
|
|
retrieval_scores = []
|
|
|
|
for query, expected_keywords in scenario.retrieval_queries:
|
|
results = store.search(scenario.entity_id, query, top_k=5)
|
|
retrieved = [{"content": m.content, "category": m.category} for m in results]
|
|
|
|
score, reasoning = self.judge.evaluate_retrieval(query, expected_keywords, retrieved)
|
|
retrieval_scores.append(score)
|
|
|
|
retrieval_results.append(
|
|
{
|
|
"query": query,
|
|
"expected_keywords": expected_keywords,
|
|
"retrieved_count": len(retrieved),
|
|
"score": score,
|
|
"reasoning": reasoning,
|
|
}
|
|
)
|
|
print(f" Query: '{query[:40]}...' -> Score: {score:.2f}, Found: {len(retrieved)}")
|
|
|
|
avg_retrieval_score = (
|
|
sum(retrieval_scores) / len(retrieval_scores) if retrieval_scores else 0
|
|
)
|
|
|
|
# Calculate overall score
|
|
overall_score = (extraction_score * 0.6) + (avg_retrieval_score * 0.4)
|
|
|
|
result = EvaluationResult(
|
|
scenario_name=scenario.name,
|
|
extraction_score=extraction_score,
|
|
retrieval_score=avg_retrieval_score,
|
|
overall_score=overall_score,
|
|
extracted_memories=extracted,
|
|
missing_nuggets=missing,
|
|
retrieval_results=retrieval_results,
|
|
judge_reasoning=extraction_reasoning,
|
|
)
|
|
|
|
print("\n--- Scenario Complete ---")
|
|
print(f" Extraction Score: {extraction_score:.2f}")
|
|
print(f" Retrieval Score: {avg_retrieval_score:.2f}")
|
|
print(f" Overall Score: {overall_score:.2f}")
|
|
|
|
return result
|
|
|
|
def run_all_scenarios(self) -> list[EvaluationResult]:
|
|
"""Run all test scenarios."""
|
|
print("\n" + "=" * 60)
|
|
print("MEMORY SYSTEM END-TO-END TEST")
|
|
print("=" * 60)
|
|
print(f"Running {len(ALL_SCENARIOS)} scenarios with LLM-as-judge evaluation")
|
|
|
|
results = []
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
for scenario in ALL_SCENARIOS:
|
|
db_path = Path(tmpdir) / f"{scenario.name}.db"
|
|
result = self.run_scenario(scenario, db_path)
|
|
results.append(result)
|
|
|
|
self.results = results
|
|
return results
|
|
|
|
def print_summary(self):
|
|
"""Print summary of all test results."""
|
|
print("\n" + "=" * 60)
|
|
print("FINAL SUMMARY")
|
|
print("=" * 60)
|
|
|
|
total_extraction = 0
|
|
total_retrieval = 0
|
|
total_overall = 0
|
|
|
|
for r in self.results:
|
|
print(f"\n{r.scenario_name}:")
|
|
print(f" Extraction: {r.extraction_score:.2f}")
|
|
print(f" Retrieval: {r.retrieval_score:.2f}")
|
|
print(f" Overall: {r.overall_score:.2f}")
|
|
if r.missing_nuggets:
|
|
print(f" Missing: {len(r.missing_nuggets)} nuggets")
|
|
|
|
total_extraction += r.extraction_score
|
|
total_retrieval += r.retrieval_score
|
|
total_overall += r.overall_score
|
|
|
|
n = len(self.results)
|
|
print(f"\n{'=' * 60}")
|
|
print("AGGREGATE SCORES")
|
|
print(f"{'=' * 60}")
|
|
print(f" Avg Extraction: {total_extraction / n:.2f}")
|
|
print(f" Avg Retrieval: {total_retrieval / n:.2f}")
|
|
print(f" Avg Overall: {total_overall / n:.2f}")
|
|
|
|
# Overall assessment
|
|
avg_overall = total_overall / n
|
|
if avg_overall >= 0.8:
|
|
verdict = "EXCELLENT - Memory system working well"
|
|
elif avg_overall >= 0.6:
|
|
verdict = "GOOD - Memory system functional with room for improvement"
|
|
elif avg_overall >= 0.4:
|
|
verdict = "FAIR - Memory system needs tuning"
|
|
else:
|
|
verdict = "POOR - Memory system needs significant work"
|
|
|
|
print(f"\nVERDICT: {verdict}")
|
|
print("=" * 60)
|
|
|
|
return avg_overall
|
|
|
|
|
|
def main():
|
|
"""Main entry point."""
|
|
# Check for API key
|
|
api_key = os.environ.get("OPENAI_API_KEY")
|
|
if not api_key:
|
|
print("ERROR: OPENAI_API_KEY environment variable not set")
|
|
print("Usage: export OPENAI_API_KEY='sk-...' && python examples/memory_e2e_test.py")
|
|
sys.exit(1)
|
|
|
|
# Run tests
|
|
tester = MemoryE2ETest(api_key)
|
|
tester.run_all_scenarios()
|
|
avg_score = tester.print_summary()
|
|
|
|
# Exit with appropriate code
|
|
sys.exit(0 if avg_score >= 0.5 else 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|