270 lines
9.8 KiB
Python
270 lines
9.8 KiB
Python
"""
|
|
Integration tests for Phase 2 Memory System
|
|
|
|
Tests the complete memory stack:
|
|
- Tier 1: ConversationBufferMemory
|
|
- Tier 2/3: QdrantConversationMemory
|
|
- Embedding Client
|
|
"""
|
|
import asyncio
|
|
import pytest
|
|
from datetime import datetime
|
|
from src.memory import (
|
|
ConversationBufferMemory,
|
|
QdrantConversationMemory,
|
|
ConversationTurn,
|
|
MessageRole,
|
|
TokenUsage,
|
|
get_buffer_memory,
|
|
get_qdrant_memory
|
|
)
|
|
from src.models.embeddings import get_embedding_client
|
|
|
|
|
|
class TestEmbeddingClient:
|
|
"""Test embedding generation"""
|
|
|
|
def test_embedding_client_init(self):
|
|
"""Test embedding client initialization"""
|
|
client = get_embedding_client()
|
|
assert client is not None
|
|
assert client.dimension == 384
|
|
print(f"✓ Embedding client initialized: {client.model_name}")
|
|
|
|
def test_single_embedding(self):
|
|
"""Test single text embedding"""
|
|
client = get_embedding_client()
|
|
text = "Hello, this is a test message for embedding generation"
|
|
|
|
embedding = client.embed_text(text)
|
|
|
|
assert isinstance(embedding, list)
|
|
assert len(embedding) == 384
|
|
assert all(isinstance(x, float) for x in embedding)
|
|
print(f"✓ Single embedding generated: {len(embedding)} dimensions")
|
|
|
|
def test_batch_embedding(self):
|
|
"""Test batch text embedding"""
|
|
client = get_embedding_client()
|
|
texts = [
|
|
"First message about Python programming",
|
|
"Second message about machine learning",
|
|
"Third message about data science"
|
|
]
|
|
|
|
embeddings = client.embed_batch(texts)
|
|
|
|
assert len(embeddings) == 3
|
|
assert all(len(emb) == 384 for emb in embeddings)
|
|
print(f"✓ Batch embeddings generated: {len(embeddings)} texts")
|
|
|
|
|
|
class TestQdrantMemory:
|
|
"""Test Qdrant memory storage and retrieval"""
|
|
|
|
@pytest.fixture
|
|
def qdrant_memory(self):
|
|
"""Get Qdrant memory instance"""
|
|
return get_qdrant_memory()
|
|
|
|
@pytest.fixture
|
|
def test_conversation_id(self):
|
|
"""Generate unique test conversation ID"""
|
|
return f"test_conv_{int(datetime.utcnow().timestamp())}"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_qdrant_connection(self, qdrant_memory):
|
|
"""Test Qdrant connection and collection"""
|
|
assert qdrant_memory.client is not None
|
|
assert qdrant_memory.collection_name == "core_api_conversations"
|
|
print(f"✓ Connected to Qdrant: {qdrant_memory.host}:{qdrant_memory.port}")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_turn(self, qdrant_memory, test_conversation_id):
|
|
"""Test adding a turn to Qdrant"""
|
|
turn = ConversationTurn(
|
|
role=MessageRole.USER,
|
|
content="What is Python?",
|
|
turn_number=1,
|
|
tokens=TokenUsage(prompt=10, completion=0, total=10)
|
|
)
|
|
|
|
await qdrant_memory.add_turn(test_conversation_id, turn)
|
|
|
|
# Verify it was stored
|
|
exists = await qdrant_memory.conversation_exists(test_conversation_id)
|
|
assert exists is True
|
|
print(f"✓ Turn stored in Qdrant: {test_conversation_id}")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chronological_retrieval(self, qdrant_memory, test_conversation_id):
|
|
"""Test Tier 2 mode: chronological retrieval"""
|
|
# Add multiple turns
|
|
turns = [
|
|
ConversationTurn(role=MessageRole.USER, content="What is Python?", turn_number=1),
|
|
ConversationTurn(role=MessageRole.ASSISTANT, content="Python is a programming language", turn_number=2),
|
|
ConversationTurn(role=MessageRole.USER, content="How do I learn it?", turn_number=3),
|
|
]
|
|
|
|
for turn in turns:
|
|
await qdrant_memory.add_turn(test_conversation_id, turn)
|
|
|
|
# Retrieve turns chronologically
|
|
retrieved = await qdrant_memory.get_turns(test_conversation_id)
|
|
|
|
assert len(retrieved) == 3
|
|
assert retrieved[0].turn_number == 1
|
|
assert retrieved[1].turn_number == 2
|
|
assert retrieved[2].turn_number == 3
|
|
assert retrieved[0].content == "What is Python?"
|
|
print(f"✓ Chronological retrieval works: {len(retrieved)} turns")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_semantic_search(self, qdrant_memory, test_conversation_id):
|
|
"""Test Tier 3 mode: semantic search"""
|
|
# Add turns with distinct topics
|
|
turns = [
|
|
ConversationTurn(role=MessageRole.USER, content="I love machine learning and neural networks", turn_number=10),
|
|
ConversationTurn(role=MessageRole.USER, content="Pizza is my favorite food", turn_number=11),
|
|
ConversationTurn(role=MessageRole.USER, content="Deep learning models are fascinating", turn_number=12),
|
|
]
|
|
|
|
for turn in turns:
|
|
await qdrant_memory.add_turn(test_conversation_id, turn)
|
|
|
|
# Search for AI-related content
|
|
results = await qdrant_memory.similarity_search(
|
|
query="artificial intelligence and AI",
|
|
conversation_id=test_conversation_id,
|
|
limit=3
|
|
)
|
|
|
|
assert len(results) > 0
|
|
# Top results should be about ML/AI, not pizza
|
|
top_result = results[0]
|
|
assert "machine learning" in top_result["content"] or "Deep learning" in top_result["content"]
|
|
assert top_result["score"] > 0.5 # Reasonable similarity score
|
|
print(f"✓ Semantic search works: {len(results)} matches, top score: {results[0]['score']:.3f}")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_conversation_stats(self, qdrant_memory, test_conversation_id):
|
|
"""Test conversation statistics"""
|
|
stats = await qdrant_memory.get_conversation_stats(test_conversation_id)
|
|
|
|
assert stats["conversation_id"] == test_conversation_id
|
|
assert stats["total_turns"] >= 0
|
|
assert "total_tokens" in stats
|
|
print(f"✓ Stats retrieved: {stats['total_turns']} turns, {stats['total_tokens']} tokens")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_clear_conversation(self, qdrant_memory, test_conversation_id):
|
|
"""Test clearing a conversation"""
|
|
# Add a turn
|
|
turn = ConversationTurn(role=MessageRole.USER, content="Test message", turn_number=99)
|
|
await qdrant_memory.add_turn(test_conversation_id, turn)
|
|
|
|
# Clear it
|
|
await qdrant_memory.clear_conversation(test_conversation_id)
|
|
|
|
# Verify it's gone
|
|
exists = await qdrant_memory.conversation_exists(test_conversation_id)
|
|
assert exists is False
|
|
print(f"✓ Conversation cleared: {test_conversation_id}")
|
|
|
|
|
|
class TestIntegration:
|
|
"""Test full integration: Tier 1 + Qdrant + Embeddings"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_full_memory_flow(self):
|
|
"""Test complete memory flow: Buffer → Qdrant"""
|
|
conversation_id = f"integration_test_{int(datetime.utcnow().timestamp())}"
|
|
|
|
# Initialize both tiers
|
|
buffer_memory = get_buffer_memory()
|
|
qdrant_memory = get_qdrant_memory()
|
|
|
|
# 1. Add turns to buffer (Tier 1)
|
|
turns = [
|
|
ConversationTurn(role=MessageRole.USER, content="Hello!", turn_number=1),
|
|
ConversationTurn(role=MessageRole.ASSISTANT, content="Hi there!", turn_number=2),
|
|
ConversationTurn(role=MessageRole.USER, content="How are you?", turn_number=3),
|
|
]
|
|
|
|
for turn in turns:
|
|
await buffer_memory.add_turn(conversation_id, turn)
|
|
|
|
# Verify buffer has them
|
|
buffer = await buffer_memory.get_buffer(conversation_id)
|
|
assert len(buffer.turns) == 3
|
|
print(f"✓ Tier 1 buffer: {len(buffer.turns)} turns")
|
|
|
|
# 2. Move to Qdrant (Tier 2/3)
|
|
for turn in buffer.turns:
|
|
await qdrant_memory.add_turn(conversation_id, turn)
|
|
|
|
# Verify Qdrant has them
|
|
qdrant_turns = await qdrant_memory.get_turns(conversation_id)
|
|
assert len(qdrant_turns) == 3
|
|
print(f"✓ Tier 2/3 Qdrant: {len(qdrant_turns)} turns")
|
|
|
|
# 3. Test semantic search across both
|
|
search_results = await qdrant_memory.similarity_search(
|
|
query="greeting",
|
|
conversation_id=conversation_id,
|
|
limit=2
|
|
)
|
|
assert len(search_results) > 0
|
|
print(f"✓ Semantic search: {len(search_results)} matches")
|
|
|
|
# Cleanup
|
|
await qdrant_memory.clear_conversation(conversation_id)
|
|
await buffer_memory.clear_conversation(conversation_id)
|
|
print(f"✓ Full memory flow complete!")
|
|
|
|
|
|
def run_tests():
|
|
"""Run all tests"""
|
|
print("\n" + "="*60)
|
|
print("Phase 2 Memory System Integration Tests")
|
|
print("="*60 + "\n")
|
|
|
|
# Test 1: Embedding Client
|
|
print("Test 1: Embedding Client")
|
|
print("-" * 40)
|
|
test_embed = TestEmbeddingClient()
|
|
test_embed.test_embedding_client_init()
|
|
test_embed.test_single_embedding()
|
|
test_embed.test_batch_embedding()
|
|
print()
|
|
|
|
# Test 2: Qdrant Memory
|
|
print("Test 2: Qdrant Memory Storage")
|
|
print("-" * 40)
|
|
test_qdrant = TestQdrantMemory()
|
|
qdrant_memory = get_qdrant_memory()
|
|
test_conv_id = f"test_conv_{int(datetime.utcnow().timestamp())}"
|
|
|
|
asyncio.run(test_qdrant.test_qdrant_connection(qdrant_memory))
|
|
asyncio.run(test_qdrant.test_add_turn(qdrant_memory, test_conv_id))
|
|
asyncio.run(test_qdrant.test_chronological_retrieval(qdrant_memory, test_conv_id))
|
|
asyncio.run(test_qdrant.test_semantic_search(qdrant_memory, test_conv_id))
|
|
asyncio.run(test_qdrant.test_conversation_stats(qdrant_memory, test_conv_id))
|
|
asyncio.run(test_qdrant.test_clear_conversation(qdrant_memory, test_conv_id))
|
|
print()
|
|
|
|
# Test 3: Full Integration
|
|
print("Test 3: Full Integration (Tier 1 + Tier 2/3)")
|
|
print("-" * 40)
|
|
test_integration = TestIntegration()
|
|
asyncio.run(test_integration.test_full_memory_flow())
|
|
print()
|
|
|
|
print("="*60)
|
|
print("✅ All Memory System Tests Passed!")
|
|
print("="*60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_tests()
|