akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions
headroom/tests/test_memory_integration.py
chopratejas da74341858 Add hierarchical memory system with graph + vector storage 
Implement comprehensive memory system supporting:
- Local backend (SQLite + FTS5 + HNSW) for zero-dependency operation
- Mem0 backends (Neo4j + Qdrant) for production graph memory
- DirectMem0Adapter for optimized pre-extracted data (bypasses LLM)
- Memory extraction with facts, entities, and relationships
- Proxy integration with --memory flag for automatic memory injection

Key components:
- headroom/memory/backends/: LocalBackend, Mem0Backend, DirectMem0Adapter
- headroom/memory/system.py: MemorySystem with tool-based interface
- headroom/memory/extraction.py: Entity and relationship extraction
- headroom/proxy/memory_handler.py: Proxy integration layer
- headroom/prediction/feature_extractor.py: Content analysis features

Testing:
- 217 new memory system tests covering all backends
- LoCoMo evaluation framework for memory quality assessment
- Integration tests for proxy memory functionality

Also removes deprecated example files in favor of focused test coverage.
2026-01-26 21:58:47 -08:00

627 lines
24 KiB
Python
Raw Permalink Blame History

"""Integration tests for Headroom Memory System.
These tests use REAL API calls - no mocks.
Tests verify the full flow from LLM tool calls to memory storage.
Requirements:
- OPENAI_API_KEY environment variable must be set
- Run with: pytest tests/test_memory_integration.py -v -s
"""
from __future__ import annotations
import os
import tempfile
import uuid
import pytest
from openai import OpenAI
# API keys must be set externally via environment variables
# Tests will be skipped if OPENAI_API_KEY is not available
@pytest.mark.skipif(
not os.environ.get("OPENAI_API_KEY"),
reason="OPENAI_API_KEY environment variable not set",
)
class TestMemoryIntegration:
"""Integration tests for the memory system with real LLM calls."""
@pytest.fixture
def openai_client(self):
"""Create an OpenAI client."""
return OpenAI()
@pytest.fixture
def temp_db_path(self):
"""Create a temporary database path."""
with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
yield f.name
# Cleanup
try:
os.unlink(f.name)
except OSError:
pass
@pytest.fixture
def user_id(self):
"""Generate a unique user ID for test isolation."""
return f"test_user_{uuid.uuid4().hex[:8]}"
# =========================================================================
# Test 1: Verify optimized tools include pre-extraction fields
# =========================================================================
def test_optimized_tools_have_extraction_fields(self):
"""Verify that optimized tools include pre-extraction fields."""
from headroom.memory.tools import get_memory_tools, get_memory_tools_optimized
# Standard tools should NOT have facts/extracted_entities
standard_tools = get_memory_tools()
memory_save = next(t for t in standard_tools if t["function"]["name"] == "memory_save")
props = memory_save["function"]["parameters"]["properties"]
assert "facts" not in props, "Standard tools should not have 'facts'"
assert "extracted_entities" not in props, (
"Standard tools should not have 'extracted_entities'"
)
# Optimized tools SHOULD have facts/extracted_entities/extracted_relationships
optimized_tools = get_memory_tools_optimized()
memory_save_opt = next(t for t in optimized_tools if t["function"]["name"] == "memory_save")
props_opt = memory_save_opt["function"]["parameters"]["properties"]
assert "facts" in props_opt, "Optimized tools should have 'facts'"
assert "extracted_entities" in props_opt, "Optimized tools should have 'extracted_entities'"
assert "extracted_relationships" in props_opt, (
"Optimized tools should have 'extracted_relationships'"
)
assert "background" in props_opt, "Optimized tools should have 'background'"
# =========================================================================
# Test 2: Verify wrapper uses correct tools based on optimized flag
# =========================================================================
def test_wrapper_uses_correct_tools(self, openai_client, temp_db_path, user_id):
"""Verify wrapper uses standard vs optimized tools correctly."""
from headroom.memory import with_memory_tools
from headroom.memory.backends.local import LocalBackend, LocalBackendConfig
config = LocalBackendConfig(db_path=temp_db_path)
backend = LocalBackend(config)
# Create non-optimized wrapper
wrapper_standard = with_memory_tools(
openai_client, backend=backend, user_id=user_id, optimized=False
)
# Create optimized wrapper
wrapper_optimized = with_memory_tools(
openai_client, backend=backend, user_id=user_id, optimized=True
)
# Verify internal flags are set correctly
assert wrapper_standard._optimized is False
assert wrapper_optimized._optimized is True
assert wrapper_optimized._inject_extraction_prompt is True
# =========================================================================
# Test 3: Verify extraction prompt is injected in optimized mode
# =========================================================================
def test_extraction_prompt_injection(self, openai_client, temp_db_path, user_id):
"""Verify extraction prompt is injected into system message."""
from headroom.memory import with_memory_tools
from headroom.memory.backends.local import LocalBackend, LocalBackendConfig
from headroom.memory.extraction import EXTRACTION_SYSTEM_PROMPT
config = LocalBackendConfig(db_path=temp_db_path)
backend = LocalBackend(config)
wrapper = with_memory_tools(
openai_client,
backend=backend,
user_id=user_id,
optimized=True,
inject_extraction_prompt=True,
)
# Get the completions object
completions = wrapper.chat.completions
# Test _prepare_messages with existing system message
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello"},
]
prepared = completions._prepare_messages(messages)
# Verify system message has extraction prompt appended
assert len(prepared) == 2
assert EXTRACTION_SYSTEM_PROMPT in prepared[0]["content"]
assert "You are a helpful assistant." in prepared[0]["content"]
# Test _prepare_messages without existing system message
messages_no_system = [{"role": "user", "content": "Hello"}]
prepared_no_system = completions._prepare_messages(messages_no_system)
# Verify system message was inserted
assert len(prepared_no_system) == 2
assert prepared_no_system[0]["role"] == "system"
assert EXTRACTION_SYSTEM_PROMPT.strip() in prepared_no_system[0]["content"]
# =========================================================================
# Test 4: LocalBackend accepts pre-extraction fields
# =========================================================================
@pytest.mark.asyncio
async def test_local_backend_pre_extraction(self, temp_db_path, user_id):
"""Test LocalBackend save_memory with pre-extraction fields."""
from headroom.memory.backends.local import LocalBackend, LocalBackendConfig
config = LocalBackendConfig(db_path=temp_db_path)
backend = LocalBackend(config)
# Save with pre-extraction fields
# Note: relationships must reference entities that are in extracted_entities
memory = await backend.save_memory(
content="John works at Netflix using Python and TensorFlow.",
user_id=user_id,
importance=0.8,
facts=["John works at Netflix", "John uses Python", "John uses TensorFlow"],
extracted_entities=[
{"entity": "John", "entity_type": "person"},
{"entity": "Netflix", "entity_type": "organization"},
{"entity": "Python", "entity_type": "technology"},
{"entity": "TensorFlow", "entity_type": "technology"},
],
extracted_relationships=[
{
"source": "John",
"relationship": "works_at",
"destination": "Netflix",
},
{"source": "John", "relationship": "uses", "destination": "Python"},
{
"source": "John",
"relationship": "uses",
"destination": "TensorFlow",
},
],
)
# Verify memory was created
assert memory is not None
assert memory.user_id == user_id
assert memory.metadata.get("_pre_extracted") is True
assert memory.metadata.get("_fact_count") == 3
# Verify entities were added to graph
graph = await backend.get_graph()
netflix_entity = await graph.get_entity_by_name(user_id, "Netflix")
assert netflix_entity is not None
assert netflix_entity.entity_type == "organization"
python_entity = await graph.get_entity_by_name(user_id, "Python")
assert python_entity is not None
assert python_entity.entity_type == "technology"
john_entity = await graph.get_entity_by_name(user_id, "John")
assert john_entity is not None
assert john_entity.entity_type == "person"
# Verify relationships were added
# Get all relationships by counting via primary storage
all_rels = [r for r in graph._relationships.values() if r.user_id == user_id]
assert len(all_rels) >= 3, f"Expected at least 3 relationships, got {len(all_rels)}"
# Verify specific relationship exists (John -> Netflix)
john_id = john_entity.id
john_outgoing_rel_ids = graph._relationships_by_source.get(john_id, set())
assert len(john_outgoing_rel_ids) >= 3, (
"Expected John to have at least 3 outgoing relationships"
)
await backend.close()
# =========================================================================
# Test 5: End-to-end with real LLM - Standard Mode
# =========================================================================
def test_e2e_standard_mode_llm_call(self, openai_client, temp_db_path, user_id):
"""Test end-to-end flow with real LLM call in standard mode."""
from headroom.memory import with_memory_tools
from headroom.memory.backends.local import LocalBackend, LocalBackendConfig
config = LocalBackendConfig(db_path=temp_db_path)
backend = LocalBackend(config)
client = with_memory_tools(
openai_client,
backend=backend,
user_id=user_id,
optimized=False, # Standard mode
)
# Make a real LLM call that should trigger memory_save
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{
"role": "system",
"content": "You are a helpful assistant that remembers important user information. When the user shares personal information, save it to memory using the memory_save tool.",
},
{
"role": "user",
"content": "Hi! My name is Alex and I work as a data scientist at Google.",
},
],
)
# Verify response was generated
assert response is not None
assert response.choices is not None
assert len(response.choices) > 0
# Check if memory tool was called
message = response.choices[0].message
if message.tool_calls:
# Verify memory_save was called
tool_names = [tc.function.name for tc in message.tool_calls]
print(f"Tools called: {tool_names}")
# Check if auto-handled
if hasattr(response, "_memory_tool_results"):
print(f"Memory tool results: {response._memory_tool_results}")
assert len(response._memory_tool_results) > 0
# =========================================================================
# Test 6: End-to-end with real LLM - Optimized Mode
# =========================================================================
def test_e2e_optimized_mode_llm_call(self, openai_client, temp_db_path, user_id):
"""Test end-to-end flow with real LLM call in optimized mode."""
from headroom.memory import with_memory_tools
from headroom.memory.backends.local import LocalBackend, LocalBackendConfig
config = LocalBackendConfig(db_path=temp_db_path)
backend = LocalBackend(config)
client = with_memory_tools(
openai_client,
backend=backend,
user_id=user_id,
optimized=True, # Optimized mode - should extract facts/entities
)
# Make a real LLM call
response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{
"role": "user",
"content": "I'm Sarah, a software engineer at Microsoft. I use Python, React, and PostgreSQL daily.",
},
],
)
# Verify response was generated
assert response is not None
assert response.choices is not None
# Check if memory tool was called with pre-extraction
message = response.choices[0].message
if message.tool_calls:
for tc in message.tool_calls:
if tc.function.name == "memory_save":
import json
args = json.loads(tc.function.arguments)
print(f"memory_save arguments: {json.dumps(args, indent=2)}")
# In optimized mode, LLM SHOULD include facts/entities
# (depends on LLM following the extraction prompt)
if "facts" in args:
print(f"Pre-extracted facts: {args['facts']}")
if "extracted_entities" in args:
print(f"Pre-extracted entities: {args['extracted_entities']}")
if "extracted_relationships" in args:
print(f"Pre-extracted relationships: {args['extracted_relationships']}")
# Check auto-handled results
if hasattr(response, "_memory_tool_results"):
print(f"Memory tool results: {response._memory_tool_results}")
# =========================================================================
# Test 7: Verify memory search works after save
# =========================================================================
@pytest.mark.asyncio
async def test_memory_search_after_save(self, temp_db_path, user_id):
"""Test that saved memories can be searched."""
from headroom.memory.backends.local import LocalBackend, LocalBackendConfig
config = LocalBackendConfig(db_path=temp_db_path)
backend = LocalBackend(config)
# Save some memories
await backend.save_memory(
content="User prefers Python for backend development",
user_id=user_id,
importance=0.9,
entities=["Python"],
extracted_entities=[{"entity": "Python", "entity_type": "technology"}],
)
await backend.save_memory(
content="User works at Netflix as a senior engineer",
user_id=user_id,
importance=0.8,
entities=["Netflix"],
extracted_entities=[{"entity": "Netflix", "entity_type": "organization"}],
)
# Search for memories
results = await backend.search_memories(
query="What programming language does the user prefer?",
user_id=user_id,
top_k=5,
)
assert len(results) > 0, "Expected at least one search result"
print(f"Search results: {[(r.memory.content, r.score) for r in results]}")
# Search with entity filter
results_netflix = await backend.search_memories(
query="Where does the user work?",
user_id=user_id,
entities=["Netflix"],
top_k=5,
)
# Should find the Netflix-related memory
assert any("Netflix" in r.memory.content for r in results_netflix), (
"Expected Netflix in results"
)
await backend.close()
# =========================================================================
# Test 8: Test include_related graph expansion
# =========================================================================
@pytest.mark.asyncio
async def test_include_related_graph_expansion(self, temp_db_path, user_id):
"""Test that include_related expands results via graph."""
from headroom.memory.backends.local import LocalBackend, LocalBackendConfig
config = LocalBackendConfig(db_path=temp_db_path)
backend = LocalBackend(config)
# Save memories with related entities
await backend.save_memory(
content="Alice is a data scientist",
user_id=user_id,
importance=0.8,
entities=["Alice"],
extracted_entities=[{"entity": "Alice", "entity_type": "person"}],
)
await backend.save_memory(
content="Alice works at Acme Corp",
user_id=user_id,
importance=0.8,
entities=["Alice", "Acme Corp"],
extracted_entities=[
{"entity": "Alice", "entity_type": "person"},
{"entity": "Acme Corp", "entity_type": "organization"},
],
extracted_relationships=[
{
"source": "Alice",
"relationship": "works_at",
"destination": "Acme Corp",
}
],
)
await backend.save_memory(
content="Acme Corp is a tech company in San Francisco",
user_id=user_id,
importance=0.7,
entities=["Acme Corp", "San Francisco"],
extracted_entities=[
{"entity": "Acme Corp", "entity_type": "organization"},
{"entity": "San Francisco", "entity_type": "location"},
],
)
# Search for Alice - should expand to related memories via graph
results_with_related = await backend.search_memories(
query="Tell me about Alice",
user_id=user_id,
top_k=10,
include_related=True,
)
# Search without related
results_without_related = await backend.search_memories(
query="Tell me about Alice",
user_id=user_id,
top_k=10,
include_related=False,
)
print(f"With related: {[r.memory.content for r in results_with_related]}")
print(f"Without related: {[r.memory.content for r in results_without_related]}")
# With related should potentially include the Acme Corp memory via Alice connection
# (This depends on graph expansion finding the connection)
assert len(results_with_related) >= len(results_without_related), (
"include_related should return same or more results"
)
await backend.close()
# =========================================================================
# Test 9: Test MemorySystem tool dispatch
# =========================================================================
@pytest.mark.asyncio
async def test_memory_system_tool_dispatch(self, temp_db_path, user_id):
"""Test MemorySystem processes tool calls correctly."""
from headroom.memory.backends.local import LocalBackend, LocalBackendConfig
from headroom.memory.system import MemorySystem
config = LocalBackendConfig(db_path=temp_db_path)
backend = LocalBackend(config)
system = MemorySystem(backend, user_id=user_id)
# Test memory_save dispatch
save_result = await system.process_tool_call(
"memory_save",
{
"content": "User likes dark mode",
"importance": 0.7,
"facts": ["Prefers dark mode"],
"extracted_entities": [{"entity": "dark mode", "entity_type": "preference"}],
},
)
assert save_result["success"] is True
assert "memory_id" in save_result or "data" in save_result
print(f"Save result: {save_result}")
# Test memory_search dispatch
search_result = await system.process_tool_call(
"memory_search", {"query": "dark mode preferences", "top_k": 5}
)
assert search_result["success"] is True
print(f"Search result: {search_result}")
await backend.close()
# =========================================================================
# Test 10: Full flow - LLM saves, then retrieves via search
# =========================================================================
def test_full_flow_save_then_search(self, openai_client, temp_db_path, user_id):
"""Test complete flow: LLM saves memory, then searches for it."""
import json
from headroom.memory import with_memory_tools
from headroom.memory.backends.local import LocalBackend, LocalBackendConfig
config = LocalBackendConfig(db_path=temp_db_path)
backend = LocalBackend(config)
client = with_memory_tools(
openai_client,
backend=backend,
user_id=user_id,
optimized=True,
)
# First: Have LLM save some information
save_response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{
"role": "user",
"content": "Remember this: My favorite programming language is Rust and I'm working on a CLI tool called headroom.",
},
],
)
print(f"Save response: {save_response.choices[0].message}")
# Process tool calls if any
if save_response.choices[0].message.tool_calls:
print(
f"Tool calls made: {[tc.function.name for tc in save_response.choices[0].message.tool_calls]}"
)
if hasattr(save_response, "_memory_tool_results"):
print(f"Results: {save_response._memory_tool_results}")
# Second: Ask LLM to recall the information
recall_response = client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{
"role": "user",
"content": "What is my favorite programming language? Search your memory.",
},
],
)
print(f"Recall response: {recall_response.choices[0].message}")
# Check if search was invoked
if recall_response.choices[0].message.tool_calls:
for tc in recall_response.choices[0].message.tool_calls:
print(f"Tool: {tc.function.name}, Args: {tc.function.arguments}")
if hasattr(recall_response, "_memory_tool_results"):
results = recall_response._memory_tool_results.get(tc.id, {})
print(f"Tool result: {json.dumps(results, indent=2, default=str)}")
class TestExtractionPrompts:
"""Tests for extraction prompt templates."""
def test_extraction_prompts_exist_and_valid(self):
"""Verify extraction prompts are defined and non-empty."""
from headroom.memory.extraction import (
ENTITY_EXTRACTION_PROMPT,
EXTRACTION_SYSTEM_PROMPT,
FACT_EXTRACTION_PROMPT,
RELATIONSHIP_EXTRACTION_PROMPT,
)
assert len(EXTRACTION_SYSTEM_PROMPT) > 100, "System prompt should be substantial"
assert len(FACT_EXTRACTION_PROMPT) > 100, "Fact prompt should be substantial"
assert len(ENTITY_EXTRACTION_PROMPT) > 100, "Entity prompt should be substantial"
assert len(RELATIONSHIP_EXTRACTION_PROMPT) > 100, (
"Relationship prompt should be substantial"
)
# Verify they mention key concepts
assert "facts" in EXTRACTION_SYSTEM_PROMPT.lower()
assert "entities" in EXTRACTION_SYSTEM_PROMPT.lower()
assert "relationships" in EXTRACTION_SYSTEM_PROMPT.lower()
class TestWrapperToolsModule:
"""Tests for wrapper_tools.py module."""
def test_wrapper_tools_imports(self):
"""Verify all necessary imports work."""
from headroom.memory.wrapper_tools import (
MemoryToolsChatCompletions,
MemoryToolsCompletions,
MemoryToolsWrapper,
with_memory_tools,
)
assert with_memory_tools is not None
assert MemoryToolsWrapper is not None
assert MemoryToolsChatCompletions is not None
assert MemoryToolsCompletions is not None
def test_with_memory_tools_accepts_optimized_param(self):
"""Verify with_memory_tools accepts optimized parameter."""
import inspect
from headroom.memory.wrapper_tools import with_memory_tools
sig = inspect.signature(with_memory_tools)
params = list(sig.parameters.keys())
assert "optimized" in params, "with_memory_tools should accept 'optimized' param"
assert "inject_extraction_prompt" in params, (
"with_memory_tools should accept 'inject_extraction_prompt' param"
)
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])
Reference in a new issue View git blame Copy permalink