151 lines
5.5 KiB
Python
151 lines
5.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test MemoryManager orchestration
|
|
|
|
Verifies unified memory interface works correctly.
|
|
"""
|
|
import asyncio
|
|
import sys
|
|
from datetime import datetime
|
|
|
|
sys.path.insert(0, '/app')
|
|
|
|
from src.memory import MemoryManager, get_memory_manager, MessageRole, TokenUsage
|
|
|
|
|
|
async def test_memory_manager():
|
|
"""Test MemoryManager orchestration"""
|
|
print("\n" + "="*60)
|
|
print("MEMORY MANAGER TEST")
|
|
print("="*60)
|
|
|
|
test_conv_id = f"manager_test_{int(datetime.utcnow().timestamp())}"
|
|
|
|
try:
|
|
# Initialize manager
|
|
manager = get_memory_manager()
|
|
print(f"✓ MemoryManager initialized")
|
|
|
|
# Test 1: Add turns through manager
|
|
print("\n1. Adding turns via MemoryManager...")
|
|
turn1 = await manager.add_turn(
|
|
conversation_id=test_conv_id,
|
|
role=MessageRole.USER,
|
|
content="Hello, how are you?",
|
|
tokens=TokenUsage(prompt=5, completion=0, total=5)
|
|
)
|
|
assert turn1.turn_number == 1
|
|
print(f" ✓ Turn 1 added: {turn1.content[:30]}...")
|
|
|
|
turn2 = await manager.add_turn(
|
|
conversation_id=test_conv_id,
|
|
role=MessageRole.ASSISTANT,
|
|
content="I'm doing great! How can I help you today?",
|
|
tokens=TokenUsage(prompt=5, completion=10, total=15)
|
|
)
|
|
assert turn2.turn_number == 2
|
|
print(f" ✓ Turn 2 added: {turn2.content[:30]}...")
|
|
|
|
# Test 2: Get recent turns (from buffer)
|
|
print("\n2. Getting recent turns from buffer...")
|
|
recent = await manager.get_recent_turns(test_conv_id, limit=10)
|
|
assert len(recent) == 2
|
|
assert recent[0].turn_number == 1
|
|
assert recent[1].turn_number == 2
|
|
print(f" ✓ Retrieved {len(recent)} recent turns from buffer")
|
|
|
|
# Test 3: Add more turns to trigger consolidation (threshold = 10)
|
|
print("\n3. Adding turns to trigger auto-consolidation...")
|
|
for i in range(3, 11): # Add turns 3-10
|
|
await manager.add_turn(
|
|
conversation_id=test_conv_id,
|
|
role=MessageRole.USER if i % 2 == 1 else MessageRole.ASSISTANT,
|
|
content=f"Test message number {i}",
|
|
tokens=TokenUsage(prompt=5, completion=5, total=10)
|
|
)
|
|
print(f" ✓ Added 8 more turns (total: 10)")
|
|
|
|
# Check if consolidation happened (turn 10 should trigger it)
|
|
print("\n4. Verifying auto-consolidation...")
|
|
stats = await manager.get_conversation_stats(test_conv_id)
|
|
print(f" Buffer turns: {stats['buffer_turns']}")
|
|
print(f" Qdrant turns: {stats['qdrant_turns']}")
|
|
print(f" Exists in buffer: {stats['exists_in_buffer']}")
|
|
print(f" Exists in Qdrant: {stats['exists_in_qdrant']}")
|
|
|
|
if stats['qdrant_turns'] > 0:
|
|
print(f" ✓ Auto-consolidation triggered! {stats['qdrant_turns']} turns in Qdrant")
|
|
else:
|
|
print(f" ⚠ No auto-consolidation yet (threshold may not be reached)")
|
|
|
|
# Test 4: Manual consolidation
|
|
print("\n5. Testing manual consolidation...")
|
|
consolidated = await manager.consolidate(test_conv_id)
|
|
print(f" ✓ Manually consolidated {consolidated} turns")
|
|
|
|
# Test 5: Get full history (buffer + Qdrant)
|
|
print("\n6. Getting full conversation history...")
|
|
full_history = await manager.get_full_history(test_conv_id)
|
|
print(f" ✓ Retrieved {len(full_history)} total turns")
|
|
assert len(full_history) == 10, f"Expected 10 turns, got {len(full_history)}"
|
|
print(f" ✓ Full history verified (10 turns)")
|
|
|
|
# Test 6: Semantic search
|
|
print("\n7. Testing semantic search...")
|
|
search_results = await manager.search_conversations(
|
|
query="greeting hello",
|
|
conversation_id=test_conv_id,
|
|
limit=3
|
|
)
|
|
if len(search_results) > 0:
|
|
print(f" ✓ Semantic search found {len(search_results)} matches")
|
|
print(f" Top: '{search_results[0]['content'][:40]}...' (score: {search_results[0]['score']:.3f})")
|
|
else:
|
|
print(f" ⚠ No semantic search results (may need more data)")
|
|
|
|
# Test 7: Clear conversation
|
|
print("\n8. Clearing conversation...")
|
|
await manager.clear_conversation(test_conv_id)
|
|
stats_after = await manager.get_conversation_stats(test_conv_id)
|
|
assert stats_after['buffer_turns'] == 0
|
|
assert stats_after['qdrant_turns'] == 0
|
|
print(f" ✓ Conversation cleared from all tiers")
|
|
|
|
print("\n" + "="*60)
|
|
print("✅ MEMORY MANAGER TEST: PASSED")
|
|
print("="*60)
|
|
print("\nMemoryManager verified:")
|
|
print(" ✓ Add turns with auto turn numbering")
|
|
print(" ✓ Get recent turns from buffer")
|
|
print(" ✓ Auto-consolidation (when threshold reached)")
|
|
print(" ✓ Manual consolidation")
|
|
print(" ✓ Get full history (buffer + Qdrant)")
|
|
print(" ✓ Semantic search")
|
|
print(" ✓ Clear conversation")
|
|
print(" ✓ Conversation stats")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"\n❌ MEMORY MANAGER TEST: FAILED")
|
|
print(f"Error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
# Cleanup on error
|
|
try:
|
|
await manager.clear_conversation(test_conv_id)
|
|
except:
|
|
pass
|
|
|
|
return False
|
|
|
|
|
|
def main():
|
|
"""Run the test"""
|
|
success = asyncio.run(test_memory_manager())
|
|
return 0 if success else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
exit(main())
|