diff --git a/.vscode/settings.json b/.vscode/settings.json
index edc76cd0..fe08c123 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -1,4 +1,7 @@
{
"python.analysis.typeCheckingMode": "basic",
- "files.exclude": {}
+ "files.exclude": {},
+ "python.testing.pytestArgs": ["tests"],
+ "python.testing.unittestEnabled": false,
+ "python.testing.pytestEnabled": true
}
diff --git a/src/deriver/queue.py b/src/deriver/queue.py
index c784482f..387d970e 100644
--- a/src/deriver/queue.py
+++ b/src/deriver/queue.py
@@ -1,7 +1,7 @@
import asyncio
import os
import signal
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
from logging import getLogger
import sentry_sdk
@@ -114,7 +114,7 @@ class QueueManager:
async def get_available_sessions(self, db: AsyncSession):
"""Get available sessions that aren't being processed"""
# Clean up stale sessions
- five_minutes_ago = datetime.utcnow() - timedelta(minutes=5)
+ five_minutes_ago = datetime.now(timezone.utc) - timedelta(minutes=5)
await db.execute(
delete(models.ActiveQueueSession).where(
models.ActiveQueueSession.last_updated < five_minutes_ago
diff --git a/tests/conftest.py b/tests/conftest.py
index 0d111596..4cf13be6 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -47,8 +47,11 @@ logging.getLogger("sqlalchemy.engine.Engine").disabled = True
# TODO use environment variable
CONNECTION_URI = make_url(
os.getenv(
- "CONNECTION_URI",
- "postgresql+psycopg://postgres:postgres@localhost:5432/postgres",
+ "TEST_CONNECTION_URI",
+ os.getenv(
+ "CONNECTION_URI",
+ "postgresql+psycopg://postgres:postgres@localhost:5432/postgres",
+ ),
)
)
TEST_DB_URL = CONNECTION_URI.set(database="test_db")
diff --git a/tests/deriver/__init__.py b/tests/deriver/__init__.py
new file mode 100644
index 00000000..74297d6c
--- /dev/null
+++ b/tests/deriver/__init__.py
@@ -0,0 +1 @@
+"""Tests for the Honcho deriver system."""
\ No newline at end of file
diff --git a/tests/deriver/conftest.py b/tests/deriver/conftest.py
new file mode 100644
index 00000000..597e4284
--- /dev/null
+++ b/tests/deriver/conftest.py
@@ -0,0 +1,342 @@
+"""Fixtures and test configuration for deriver tests."""
+
+import json
+from unittest.mock import AsyncMock, MagicMock, patch
+from uuid import uuid4
+
+import pytest
+import pytest_asyncio
+from nanoid import generate as generate_nanoid
+
+from src import models
+from src.deriver.queue import QueueManager
+from src.deriver.tom.embeddings import CollectionEmbeddingStore
+
+
+@pytest.fixture
+def mock_llm_responses():
+ """Provides structured mock responses for different LLM operations."""
+ return {
+ "fact_extraction": json.dumps({
+ "facts": [
+ "User is a software developer",
+ "User works remotely",
+ "User prefers coffee over tea",
+ "User uses Python and JavaScript"
+ ]
+ }),
+ "fact_extraction_empty": json.dumps({"facts": []}),
+ "fact_extraction_malformed": "This is not valid JSON",
+ "tom_single_prompt": json.dumps({
+ "confidence": 0.8,
+ "user_representation": {
+ "personality_traits": ["analytical", "detail-oriented", "collaborative"],
+ "preferences": ["remote work", "technical discussions"],
+ "communication_style": "direct and concise",
+ "expertise_areas": ["software development", "Python programming"]
+ }
+ }),
+ "tom_conversational": "Based on our conversation, I believe the user is a thoughtful software developer who values clear communication and technical excellence.",
+ "summary_short": "User discussed Python development and remote work preferences.",
+ "summary_long": "User is a software developer working remotely who has shown expertise in Python development and expressed preferences for asynchronous communication and detailed technical discussions."
+ }
+
+
+@pytest.fixture
+def mock_embeddings():
+ """Provides mock embedding vectors for testing."""
+ return {
+ "fact_embedding": [0.1, 0.2, 0.3] + [0.0] * 1533, # 1536-dim vector
+ "query_embedding": [0.15, 0.25, 0.35] + [0.0] * 1533,
+ "duplicate_embedding": [0.1, 0.2, 0.3] + [0.0] * 1533, # Identical to fact_embedding
+ "different_embedding": [0.9, 0.8, 0.7] + [0.0] * 1533
+ }
+
+
+@pytest.fixture(autouse=True)
+def mock_model_clients(mock_llm_responses):
+ """Mock ModelClient instances for all TOM methods."""
+ with (
+ patch("src.deriver.tom.single_prompt.ModelClient") as mock_single_prompt_client,
+ patch("src.deriver.tom.long_term.ModelClient") as mock_long_term_client,
+ patch("src.utils.model_client.ModelClient") as mock_utils_client,
+ ):
+ # Utils client for general use
+ mock_utils_instance = AsyncMock()
+ mock_utils_instance.generate.return_value = f"{mock_llm_responses['fact_extraction']}"
+ mock_utils_client.return_value = mock_utils_instance
+
+ # Single prompt TOM client
+ mock_single_prompt_instance = AsyncMock()
+ mock_single_prompt_instance.generate.return_value = mock_llm_responses['tom_single_prompt']
+ mock_single_prompt_client.return_value = mock_single_prompt_instance
+
+ # Long term TOM client
+ mock_long_term_instance = AsyncMock()
+ mock_long_term_instance.generate.return_value = f"{mock_llm_responses['fact_extraction']}"
+ mock_long_term_client.return_value = mock_long_term_instance
+
+ yield {
+ "utils": mock_utils_instance,
+ "single_prompt": mock_single_prompt_instance,
+ "long_term": mock_long_term_instance,
+ }
+
+
+@pytest.fixture(autouse=True)
+def mock_vector_operations(mock_embeddings):
+ """Mock vector similarity operations and database queries."""
+ with (
+ patch("src.crud.get_duplicate_documents") as mock_get_duplicates,
+ patch("src.crud.create_collection") as mock_create_collection,
+ patch("src.crud.get_collection_by_name") as mock_get_collection,
+ patch("src.crud.get_documents") as mock_get_documents,
+ patch("src.crud.query_documents") as mock_query_documents,
+ ):
+ # No duplicates by default
+ mock_get_duplicates.return_value = []
+
+ # Mock collection operations
+ mock_collection = MagicMock()
+ mock_collection.public_id = str(uuid4())
+ mock_create_collection.return_value = mock_collection
+ mock_get_collection.return_value = mock_collection
+
+ # Mock document operations
+ mock_get_documents.return_value = []
+ mock_query_documents.return_value = []
+
+ yield {
+ "get_duplicates": mock_get_duplicates,
+ "create_collection": mock_create_collection,
+ "get_collection": mock_get_collection,
+ "get_documents": mock_get_documents,
+ "query_documents": mock_query_documents
+ }
+
+
+@pytest_asyncio.fixture
+async def queue_manager():
+ """Provides a QueueManager instance for testing."""
+ # Mock the environment variable for test concurrency
+ with patch("src.deriver.queue.os.getenv") as mock_getenv:
+ mock_getenv.return_value = "2" # 2 workers for testing
+ manager = QueueManager()
+ yield manager
+
+
+@pytest_asyncio.fixture
+async def embedding_store(sample_data):
+ """Provides a CollectionEmbeddingStore for testing."""
+ test_app, test_user = sample_data
+ collection_id = str(uuid4())
+ store = CollectionEmbeddingStore(test_app.public_id, test_user.public_id, collection_id)
+ yield store
+
+
+@pytest_asyncio.fixture
+async def sample_messages(db_session, sample_data):
+ """Creates sample messages for testing deriver processing."""
+ test_app, test_user = sample_data
+
+ # Create a test session
+ session = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={}
+ )
+ db_session.add(session)
+ await db_session.flush()
+
+ # Create sample messages
+ messages = []
+
+ # User message
+ user_message = models.Message(
+ session_id=session.public_id,
+ is_user=True,
+ content="I'm a Python developer working on AI projects. I prefer remote work and love debugging complex problems.",
+ metadata={},
+ user_id=test_user.public_id,
+ app_id=test_app.public_id
+ )
+ db_session.add(user_message)
+ messages.append(user_message)
+
+ # AI message
+ ai_message = models.Message(
+ session_id=session.public_id,
+ is_user=False,
+ content="That's great! Python is excellent for AI development. What specific AI frameworks do you work with?",
+ metadata={},
+ user_id=test_user.public_id,
+ app_id=test_app.public_id
+ )
+ db_session.add(ai_message)
+ messages.append(ai_message)
+
+ # Another user message
+ user_message_2 = models.Message(
+ session_id=session.public_id,
+ is_user=True,
+ content="I mainly use PyTorch and transformers. Currently building a chatbot with FastAPI.",
+ metadata={},
+ user_id=test_user.public_id,
+ app_id=test_app.public_id
+ )
+ db_session.add(user_message_2)
+ messages.append(user_message_2)
+
+ await db_session.flush()
+
+ yield session, messages
+
+
+@pytest_asyncio.fixture
+async def sample_queue_items(db_session, sample_messages):
+ """Creates sample queue items for testing queue processing."""
+ session, messages = sample_messages
+
+ queue_items = []
+ for message in messages:
+ if message.is_user: # Only user messages get queued for processing
+ queue_item = models.QueueItem(
+ session_id=session.id, # Use integer ID, not public_id
+ payload={"message_id": message.public_id},
+ processed=False
+ )
+ db_session.add(queue_item)
+ queue_items.append(queue_item)
+
+ await db_session.flush()
+ yield session, messages, queue_items
+
+
+@pytest_asyncio.fixture
+async def sample_facts(db_session, sample_data):
+ """Creates sample user facts stored in collections."""
+ test_app, test_user = sample_data
+
+ # Create user collection
+ collection = models.Collection(
+ app_id=test_app.public_id,
+ user_id=test_user.public_id,
+ name=f"user_{test_user.public_id}",
+ metadata={"type": "user_facts"}
+ )
+ db_session.add(collection)
+ await db_session.flush()
+
+ # Create sample documents (facts)
+ facts = [
+ "User is a Python developer",
+ "User works remotely",
+ "User prefers PyTorch over TensorFlow",
+ "User enjoys debugging complex problems"
+ ]
+
+ documents = []
+ for fact in facts:
+ doc = models.Document(
+ collection_id=collection.public_id,
+ content=fact,
+ metadata={"extracted_at": "2024-01-01T00:00:00Z"},
+ embedding=[0.1] * 1536 # Mock embedding
+ )
+ db_session.add(doc)
+ documents.append(doc)
+
+ await db_session.flush()
+ yield collection, documents
+
+
+@pytest.fixture
+def mock_background_task():
+ """Mock FastAPI BackgroundTasks for testing async task scheduling."""
+ with patch("fastapi.BackgroundTasks") as mock_bg:
+ mock_task = MagicMock()
+ mock_bg.return_value = mock_task
+ yield mock_task
+
+
+@pytest.fixture
+def mock_signal_handling():
+ """Mock signal handling for testing graceful shutdown."""
+ with (
+ patch("signal.signal") as mock_signal,
+ patch("signal.SIGTERM") as mock_sigterm,
+ patch("signal.SIGINT") as mock_sigint,
+ ):
+ yield {
+ "signal": mock_signal,
+ "sigterm": mock_sigterm,
+ "sigint": mock_sigint
+ }
+
+
+@pytest.fixture
+def mock_semaphore():
+ """Mock asyncio.Semaphore for testing concurrency control."""
+ mock_sem = AsyncMock()
+ mock_sem.__aenter__ = AsyncMock(return_value=mock_sem)
+ mock_sem.__aexit__ = AsyncMock(return_value=None)
+
+ with patch("asyncio.Semaphore") as mock_semaphore_class:
+ mock_semaphore_class.return_value = mock_sem
+ yield mock_sem
+
+
+@pytest.fixture
+def tom_method_config():
+ """Fixture for testing different TOM method configurations."""
+ return {
+ "single_prompt": "SINGLE_PROMPT",
+ "conversational": "CONVERSATIONAL",
+ "long_term": "LONG_TERM"
+ }
+
+
+@pytest_asyncio.fixture
+async def mock_session_processing():
+ """Mock session processing for integration tests."""
+ with (
+ patch("src.deriver.consumer.process_user_message") as mock_process_user,
+ patch("src.deriver.consumer.process_ai_message") as mock_process_ai,
+ patch("src.deriver.consumer.maybe_create_summary") as mock_create_summary,
+ ):
+ mock_process_user.return_value = None
+ mock_process_ai.return_value = None
+ mock_create_summary.return_value = None
+
+ yield {
+ "process_user": mock_process_user,
+ "process_ai": mock_process_ai,
+ "create_summary": mock_create_summary
+ }
+
+
+@pytest.fixture
+def performance_config():
+ """Configuration for performance testing."""
+ return {
+ "max_workers": 4,
+ "message_count": 100,
+ "session_count": 10,
+ "timeout_seconds": 30,
+ "fact_extraction_time_limit": 5.0, # seconds
+ "tom_inference_time_limit": 3.0, # seconds
+ "queue_processing_time_limit": 1.0 # seconds per message
+ }
+
+
+@pytest.fixture
+def error_scenarios():
+ """Provides various error scenarios for testing error handling."""
+ return {
+ "llm_timeout": "LLM request timed out",
+ "llm_api_error": "API rate limit exceeded",
+ "database_connection_error": "Database connection failed",
+ "invalid_json_response": "Malformed JSON in LLM response",
+ "embedding_api_error": "Embedding service unavailable",
+ "vector_similarity_error": "Vector similarity calculation failed"
+ }
\ No newline at end of file
diff --git a/tests/deriver/test_background_integration.py b/tests/deriver/test_background_integration.py
new file mode 100644
index 00000000..2a8e32ac
--- /dev/null
+++ b/tests/deriver/test_background_integration.py
@@ -0,0 +1,632 @@
+"""End-to-end integration tests for the complete deriver workflow."""
+
+import json
+from unittest.mock import AsyncMock, MagicMock, patch
+from uuid import uuid4
+
+import pytest
+import pytest_asyncio
+
+from src import models
+from src.deriver import consumer
+from src.deriver.queue import QueueManager
+
+
+class TestMessageToFactsWorkflow:
+ """Test the complete workflow from message creation to fact storage."""
+
+ @pytest_asyncio.fixture
+ async def integration_setup(self, db_session, sample_data):
+ """Setup complete integration test data."""
+ test_app, test_user = sample_data
+
+ # Create a session
+ session = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={}
+ )
+ db_session.add(session)
+ await db_session.flush()
+
+ # Create user collection for fact storage
+ collection = models.Collection(
+ app_id=test_app.public_id,
+ user_id=test_user.public_id,
+ name=f"user_{test_user.public_id}",
+ metadata={"type": "user_facts"}
+ )
+ db_session.add(collection)
+ await db_session.flush()
+
+ # Create user messages
+ messages = []
+ for i, content in enumerate([
+ "Hi, I'm Sarah, a data scientist working remotely from Seattle",
+ "I've been using Python for machine learning for about 3 years",
+ "My current project involves building recommendation systems with PyTorch"
+ ]):
+ message = models.Message(
+ session_id=session.public_id,
+ is_user=True,
+ content=content,
+ metadata={},
+ user_id=test_user.public_id,
+ app_id=test_app.public_id
+ )
+ db_session.add(message)
+ messages.append(message)
+
+ await db_session.flush()
+
+ return test_app, test_user, session, collection, messages
+
+ @pytest.mark.asyncio
+ async def test_complete_message_processing_workflow(self, db_session, integration_setup):
+ """Test complete workflow: message → fact extraction → vector storage."""
+ test_app, test_user, session, collection, messages = integration_setup
+
+ # Mock the fact extraction to return realistic facts
+ extracted_facts = [
+ "User name is Sarah",
+ "User is a data scientist",
+ "User works remotely from Seattle",
+ "User has 3 years of Python experience",
+ "User specializes in machine learning",
+ "User is currently working on recommendation systems",
+ "User uses PyTorch for current project"
+ ]
+
+ with (
+ patch("src.deriver.consumer.extract_facts_long_term") as mock_extract,
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_history,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection,
+ patch("src.deriver.consumer.CollectionEmbeddingStore") as mock_store_class
+ ):
+ # Setup mocks
+ mock_extract.return_value = extracted_facts
+ mock_history.return_value = ("Previous conversation context", [], None)
+ mock_get_collection.return_value = collection
+
+ # Mock embedding store
+ mock_store = AsyncMock()
+ mock_store.remove_duplicates.return_value = extracted_facts # No duplicates
+ mock_store.save_facts.return_value = None
+ mock_store_class.return_value = mock_store
+
+ # Process each message
+ for message in messages:
+ await consumer.process_user_message(
+ message.content,
+ test_app.public_id,
+ test_user.public_id,
+ session.public_id,
+ message.public_id,
+ db_session
+ )
+
+ # Verify fact extraction was called for each message
+ assert mock_extract.call_count == 3
+
+ # Verify facts were saved for each message
+ assert mock_store.save_facts.call_count == 3
+
+ # Verify the facts that would be saved
+ all_saved_facts = []
+ for call in mock_store.save_facts.call_args_list:
+ facts_arg = call[0][0] # First positional argument
+ all_saved_facts.extend(facts_arg)
+
+ # Should have saved all extracted facts
+ assert len(all_saved_facts) == len(extracted_facts) * 3
+
+ @pytest.mark.asyncio
+ async def test_queue_to_consumer_integration(self, db_session, integration_setup):
+ """Test integration between queue management and message processing."""
+ test_app, test_user, session, collection, messages = integration_setup
+
+ # Create queue items for user messages (simulating how they're created in real system)
+ queue_items = []
+ for message in messages:
+ if message.is_user:
+ # Create queue item with the payload structure used in real system
+ payload = {
+ "message_id": message.public_id,
+ "is_user": message.is_user,
+ "content": message.content,
+ "app_id": test_app.public_id,
+ "user_id": test_user.public_id,
+ "session_id": session.public_id
+ }
+
+ queue_item = models.QueueItem(
+ session_id=session.id, # Use integer ID for queue
+ payload=payload,
+ processed=False
+ )
+ db_session.add(queue_item)
+ queue_items.append(queue_item)
+
+ await db_session.flush()
+
+ # Mock the consumer processing functions
+ with (
+ patch("src.deriver.consumer.process_user_message") as mock_process_user,
+ patch("src.deriver.consumer.process_ai_message") as mock_process_ai,
+ patch("src.deriver.consumer.summarize_if_needed") as mock_summarize
+ ):
+ mock_process_user.return_value = None
+ mock_process_ai.return_value = None
+ mock_summarize.return_value = None
+
+ # Process items through the consumer
+ for queue_item in queue_items:
+ await consumer.process_item(db_session, queue_item.payload)
+
+ # Verify all user messages were processed
+ assert mock_process_user.call_count == 3
+ assert mock_process_ai.call_count == 0 # No AI messages
+ assert mock_summarize.call_count == 3 # Summary check for each message
+
+ # Verify the arguments passed to process_user_message
+ for i, call in enumerate(mock_process_user.call_args_list):
+ args = call[0]
+ assert args[0] == messages[i].content # content
+ assert args[1] == test_app.public_id # app_id
+ assert args[2] == test_user.public_id # user_id
+ assert args[3] == session.public_id # session_id
+ assert args[4] == messages[i].public_id # message_id
+ assert args[5] == db_session # db_session
+
+ @pytest.mark.asyncio
+ async def test_tom_inference_integration(self, db_session, integration_setup, mock_llm_responses):
+ """Test integration of TOM inference with fact extraction workflow."""
+ test_app, test_user, session, collection, messages = integration_setup
+
+ # Test facts that would be extracted
+ user_facts = [
+ "User name is Sarah",
+ "User is a data scientist",
+ "User works remotely"
+ ]
+
+ with (
+ patch("src.deriver.tom.get_tom_inference") as mock_tom_inference,
+ patch("src.deriver.tom.get_user_representation") as mock_user_rep,
+ patch("src.deriver.consumer.extract_facts_long_term") as mock_extract_facts,
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_history,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection,
+ patch("src.deriver.consumer.CollectionEmbeddingStore") as mock_store_class
+ ):
+ # Setup TOM mocks
+ mock_tom_inference.return_value = mock_llm_responses['tom_single_prompt']
+ mock_user_rep.return_value = mock_llm_responses['tom_single_prompt']
+
+ # Setup fact extraction mocks
+ mock_extract_facts.return_value = user_facts
+ mock_history.return_value = ("Chat history", [], None)
+ mock_get_collection.return_value = collection
+
+ # Mock embedding store
+ mock_store = AsyncMock()
+ mock_store.remove_duplicates.return_value = user_facts
+ mock_store.save_facts.return_value = None
+ mock_store_class.return_value = mock_store
+
+ # Process a user message
+ await consumer.process_user_message(
+ messages[0].content,
+ test_app.public_id,
+ test_user.public_id,
+ session.public_id,
+ messages[0].public_id,
+ db_session
+ )
+
+ # Verify fact extraction occurred
+ mock_extract_facts.assert_called_once()
+ mock_store.save_facts.assert_called_once_with(user_facts, message_id=messages[0].public_id)
+
+ # The TOM inference methods aren't called directly in consumer,
+ # but we've verified the infrastructure is in place
+
+
+class TestErrorRecoveryIntegration:
+ """Test error recovery and resilience in integrated workflows."""
+
+ @pytest_asyncio.fixture
+ async def error_test_setup(self, db_session, sample_data):
+ """Setup data for error testing."""
+ test_app, test_user = sample_data
+
+ session = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={}
+ )
+ db_session.add(session)
+ await db_session.flush()
+
+ message = models.Message(
+ session_id=session.public_id,
+ is_user=True,
+ content="Test message for error scenarios",
+ metadata={},
+ user_id=test_user.public_id,
+ app_id=test_app.public_id
+ )
+ db_session.add(message)
+ await db_session.flush()
+
+ return test_app, test_user, session, message
+
+ @pytest.mark.asyncio
+ async def test_fact_extraction_error_recovery(self, db_session, error_test_setup):
+ """Test that fact extraction errors don't break the entire workflow."""
+ test_app, test_user, session, message = error_test_setup
+
+ with (
+ patch("src.deriver.consumer.extract_facts_long_term") as mock_extract,
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_history,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection
+ ):
+ # Mock fact extraction to fail
+ mock_extract.side_effect = Exception("LLM API timeout")
+ mock_history.return_value = ("", [], None)
+
+ # Mock collection to avoid that error
+ mock_collection = MagicMock()
+ mock_collection.public_id = str(uuid4())
+ mock_get_collection.return_value = mock_collection
+
+ # Should raise the exception (let caller handle it)
+ with pytest.raises(Exception, match="LLM API timeout"):
+ await consumer.process_user_message(
+ message.content,
+ test_app.public_id,
+ test_user.public_id,
+ session.public_id,
+ message.public_id,
+ db_session
+ )
+
+ @pytest.mark.asyncio
+ async def test_partial_fact_storage_error_recovery(self, db_session, error_test_setup):
+ """Test recovery when some facts fail to store."""
+ test_app, test_user, session, message = error_test_setup
+
+ facts_to_extract = [
+ "User likes programming",
+ "This fact will fail to store",
+ "User works in tech"
+ ]
+
+ with (
+ patch("src.deriver.consumer.extract_facts_long_term") as mock_extract,
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_history,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection,
+ patch("src.deriver.consumer.CollectionEmbeddingStore") as mock_store_class
+ ):
+ mock_extract.return_value = facts_to_extract
+ mock_history.return_value = ("", [], None)
+
+ mock_collection = MagicMock()
+ mock_collection.public_id = str(uuid4())
+ mock_get_collection.return_value = mock_collection
+
+ # Mock embedding store where save_facts has partial failure
+ mock_store = AsyncMock()
+ mock_store.remove_duplicates.return_value = facts_to_extract
+ # save_facts method handles its own errors gracefully
+ mock_store.save_facts.return_value = None
+ mock_store_class.return_value = mock_store
+
+ # Should complete successfully even with partial failures
+ await consumer.process_user_message(
+ message.content,
+ test_app.public_id,
+ test_user.public_id,
+ session.public_id,
+ message.public_id,
+ db_session
+ )
+
+ # Verify the workflow completed
+ mock_extract.assert_called_once()
+ mock_store.save_facts.assert_called_once()
+
+
+class TestSummaryIntegration:
+ """Test summary generation integration with message processing."""
+
+ @pytest_asyncio.fixture
+ async def summary_test_setup(self, db_session, sample_data):
+ """Setup data for summary testing."""
+ test_app, test_user = sample_data
+
+ session = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={}
+ )
+ db_session.add(session)
+ await db_session.flush()
+
+ return test_app, test_user, session
+
+ @pytest.mark.asyncio
+ async def test_summary_generation_integration(self, db_session, summary_test_setup):
+ """Test that summary generation integrates properly with message processing."""
+ test_app, test_user, session = summary_test_setup
+
+ # Create enough messages to trigger summary generation
+ messages = []
+ for i in range(25): # Enough to trigger both short and long summaries
+ message = models.Message(
+ session_id=session.public_id,
+ is_user=True,
+ content=f"Message {i+1}: User discussing various topics",
+ metadata={},
+ user_id=test_user.public_id,
+ app_id=test_app.public_id
+ )
+ db_session.add(message)
+ messages.append(message)
+
+ await db_session.flush()
+
+ with (
+ patch("src.deriver.consumer.extract_facts_long_term") as mock_extract,
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_history,
+ patch("src.deriver.consumer.history.should_create_summary") as mock_should_create,
+ patch("src.deriver.consumer.history.create_summary") as mock_create_summary,
+ patch("src.deriver.consumer.history.save_summary_metamessage") as mock_save_summary,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection,
+ patch("src.deriver.consumer.CollectionEmbeddingStore") as mock_store_class
+ ):
+ # Setup mocks
+ mock_extract.return_value = ["User is engaged in conversation"]
+ mock_history.return_value = ("Previous context", [], None)
+
+ mock_collection = MagicMock()
+ mock_collection.public_id = str(uuid4())
+ mock_get_collection.return_value = mock_collection
+
+ mock_store = AsyncMock()
+ mock_store.remove_duplicates.return_value = ["User is engaged"]
+ mock_store.save_facts.return_value = None
+ mock_store_class.return_value = mock_store
+
+ # Mock summary creation - simulate that summaries are needed
+ mock_should_create.return_value = (True, messages[:10], None)
+ mock_create_summary.return_value = "Summary of recent messages"
+ mock_save_summary.return_value = None
+
+ # Process the last message (which should trigger summary check)
+ await consumer.process_item(db_session, {
+ "message_id": messages[-1].public_id,
+ "is_user": True,
+ "content": messages[-1].content,
+ "app_id": test_app.public_id,
+ "user_id": test_user.public_id,
+ "session_id": session.public_id
+ })
+
+ # Verify fact extraction occurred
+ mock_extract.assert_called_once()
+
+ # Verify summary generation was checked
+ mock_should_create.assert_called()
+
+ # If summaries were triggered, verify they were created
+ if mock_should_create.call_count > 0:
+ # Summary creation logic was invoked
+ assert True # Successfully integrated
+
+
+class TestConcurrentProcessing:
+ """Test concurrent processing scenarios."""
+
+ @pytest.mark.asyncio
+ async def test_concurrent_message_processing(self, db_session, sample_data):
+ """Test that multiple messages can be processed concurrently safely."""
+ test_app, test_user = sample_data
+
+ # Create multiple sessions
+ sessions = []
+ for i in range(3):
+ session = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={"session_num": i}
+ )
+ db_session.add(session)
+ sessions.append(session)
+
+ await db_session.flush()
+
+ # Create messages for each session
+ all_messages = []
+ for i, session in enumerate(sessions):
+ message = models.Message(
+ session_id=session.public_id,
+ is_user=True,
+ content=f"Session {i} message: User sharing information",
+ metadata={},
+ user_id=test_user.public_id,
+ app_id=test_app.public_id
+ )
+ db_session.add(message)
+ all_messages.append(message)
+
+ await db_session.flush()
+
+ processed_count = 0
+
+ async def mock_process_user_message(*args, **kwargs):
+ nonlocal processed_count
+ processed_count += 1
+ # Simulate some processing time
+ import asyncio
+ await asyncio.sleep(0.01)
+
+ with (
+ patch("src.deriver.consumer.process_user_message", side_effect=mock_process_user_message),
+ patch("src.deriver.consumer.process_ai_message") as mock_process_ai,
+ patch("src.deriver.consumer.summarize_if_needed") as mock_summarize
+ ):
+ mock_process_ai.return_value = None
+ mock_summarize.return_value = None
+
+ # Process all messages concurrently
+ import asyncio
+ tasks = []
+ for message in all_messages:
+ payload = {
+ "message_id": message.public_id,
+ "is_user": True,
+ "content": message.content,
+ "app_id": test_app.public_id,
+ "user_id": test_user.public_id,
+ "session_id": session.public_id
+ }
+ task = asyncio.create_task(consumer.process_item(db_session, payload))
+ tasks.append(task)
+
+ # Wait for all processing to complete
+ await asyncio.gather(*tasks)
+
+ # Verify all messages were processed
+ assert processed_count == 3
+
+
+class TestFullSystemIntegration:
+ """Test complete system integration from API to storage."""
+
+ @pytest.mark.asyncio
+ async def test_realistic_user_conversation_workflow(self, db_session, sample_data):
+ """Test a realistic user conversation workflow end-to-end."""
+ test_app, test_user = sample_data
+
+ # Create session
+ session = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={"conversation_type": "onboarding"}
+ )
+ db_session.add(session)
+ await db_session.flush()
+
+ # Realistic conversation messages
+ conversation = [
+ ("user", "Hi! I'm Alex, a software engineer based in San Francisco."),
+ ("ai", "Hello Alex! It's nice to meet you. What kind of software engineering do you focus on?"),
+ ("user", "I mainly work on backend systems using Python and Go. Currently building microservices for a fintech company."),
+ ("ai", "That sounds interesting! Fintech is such a dynamic field. What's the most challenging part of your current project?"),
+ ("user", "The main challenge is handling high-frequency trading data while maintaining low latency. We're processing millions of transactions per second."),
+ ("ai", "That's impressive scale! Are you using any specific technologies for handling that throughput?"),
+ ("user", "Yes, we're using Kafka for streaming, Redis for caching, and PostgreSQL with read replicas. Also experimenting with some Rust components for ultra-low latency parts.")
+ ]
+
+ # Create all messages
+ messages = []
+ for role, content in conversation:
+ message = models.Message(
+ session_id=session.public_id,
+ is_user=(role == "user"),
+ content=content,
+ metadata={},
+ user_id=test_user.public_id,
+ app_id=test_app.public_id
+ )
+ db_session.add(message)
+ messages.append(message)
+
+ await db_session.flush()
+
+ # Expected facts that would be extracted
+ expected_facts = [
+ "User name is Alex",
+ "User is a software engineer",
+ "User is based in San Francisco",
+ "User works on backend systems",
+ "User uses Python and Go",
+ "User works at a fintech company",
+ "User builds microservices",
+ "User handles high-frequency trading data",
+ "User processes millions of transactions per second",
+ "User uses Kafka for streaming",
+ "User uses Redis for caching",
+ "User uses PostgreSQL with read replicas",
+ "User is experimenting with Rust components"
+ ]
+
+ # Track all extracted facts
+ all_extracted_facts = []
+
+ def mock_extract_facts(chat_history):
+ # Simulate realistic fact extraction based on content
+ if "Alex" in chat_history and "software engineer" in chat_history:
+ return ["User name is Alex", "User is a software engineer", "User is based in San Francisco"]
+ elif "Python and Go" in chat_history:
+ return ["User works on backend systems", "User uses Python and Go", "User works at a fintech company"]
+ elif "Kafka" in chat_history:
+ return ["User uses Kafka for streaming", "User uses Redis for caching", "User uses PostgreSQL"]
+ return []
+
+ with (
+ patch("src.deriver.consumer.extract_facts_long_term", side_effect=mock_extract_facts),
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_history,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection,
+ patch("src.deriver.consumer.CollectionEmbeddingStore") as mock_store_class,
+ patch("src.deriver.consumer.summarize_if_needed") as mock_summarize
+ ):
+ mock_history.return_value = ("", [], None)
+
+ mock_collection = MagicMock()
+ mock_collection.public_id = str(uuid4())
+ mock_get_collection.return_value = mock_collection
+
+ # Track saved facts
+ saved_facts = []
+ def track_save_facts(facts, **kwargs):
+ saved_facts.extend(facts)
+
+ mock_store = AsyncMock()
+ mock_store.remove_duplicates.side_effect = lambda facts: facts # No duplicates
+ mock_store.save_facts.side_effect = track_save_facts
+ mock_store_class.return_value = mock_store
+
+ mock_summarize.return_value = None
+
+ # Process only user messages (as would happen in real system)
+ user_messages = [msg for msg in messages if msg.is_user]
+
+ for message in user_messages:
+ await consumer.process_user_message(
+ message.content,
+ test_app.public_id,
+ test_user.public_id,
+ session.public_id,
+ message.public_id,
+ db_session
+ )
+
+ # Verify facts were extracted and saved
+ assert len(saved_facts) > 0
+
+ # Verify user-specific facts were captured
+ saved_facts_str = " ".join(saved_facts)
+ assert "Alex" in saved_facts_str
+ assert "software engineer" in saved_facts_str
+ assert "San Francisco" in saved_facts_str
+
+ # Verify technical details were captured
+ tech_keywords = ["Python", "Go", "fintech", "Kafka", "Redis", "PostgreSQL"]
+ captured_tech = [kw for kw in tech_keywords if kw in saved_facts_str]
+ assert len(captured_tech) > 0
+
+ print(f"✅ Integration test completed successfully!")
+ print(f"📊 Processed {len(user_messages)} user messages")
+ print(f"💾 Saved {len(saved_facts)} facts total")
+ print(f"🔧 Captured {len(captured_tech)} technical details")
\ No newline at end of file
diff --git a/tests/deriver/test_consumer.py b/tests/deriver/test_consumer.py
new file mode 100644
index 00000000..0129ea79
--- /dev/null
+++ b/tests/deriver/test_consumer.py
@@ -0,0 +1,588 @@
+"""Tests for the consumer module and message processing functionality."""
+
+import json
+from unittest.mock import AsyncMock, MagicMock, patch
+from uuid import uuid4
+
+import pytest
+import pytest_asyncio
+
+from src import models
+from src.deriver import consumer
+from src.utils.history import SummaryType
+
+
+class TestProcessItem:
+ """Test the main process_item entry point."""
+
+ @pytest.mark.asyncio
+ async def test_process_item_routes_user_message(self, db_session):
+ """Test that process_item correctly routes user messages."""
+ payload = {
+ "message_id": str(uuid4()),
+ "is_user": True,
+ "content": "Hello, I'm a Python developer",
+ "app_id": str(uuid4()),
+ "user_id": str(uuid4()),
+ "session_id": str(uuid4())
+ }
+
+ with (
+ patch("src.deriver.consumer.process_user_message") as mock_process_user,
+ patch("src.deriver.consumer.process_ai_message") as mock_process_ai,
+ patch("src.deriver.consumer.summarize_if_needed") as mock_summarize
+ ):
+ mock_process_user.return_value = None
+ mock_process_ai.return_value = None
+ mock_summarize.return_value = None
+
+ await consumer.process_item(db_session, payload)
+
+ # Should call user message processing
+ mock_process_user.assert_called_once_with(
+ payload["content"],
+ payload["app_id"],
+ payload["user_id"],
+ payload["session_id"],
+ payload["message_id"],
+ db_session
+ )
+ mock_process_ai.assert_not_called()
+ mock_summarize.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_process_item_routes_ai_message(self, db_session):
+ """Test that process_item correctly routes AI messages."""
+ payload = {
+ "message_id": str(uuid4()),
+ "is_user": False,
+ "content": "I can help you with Python development!",
+ "app_id": str(uuid4()),
+ "user_id": str(uuid4()),
+ "session_id": str(uuid4())
+ }
+
+ with (
+ patch("src.deriver.consumer.process_user_message") as mock_process_user,
+ patch("src.deriver.consumer.process_ai_message") as mock_process_ai,
+ patch("src.deriver.consumer.summarize_if_needed") as mock_summarize
+ ):
+ mock_process_user.return_value = None
+ mock_process_ai.return_value = None
+ mock_summarize.return_value = None
+
+ await consumer.process_item(db_session, payload)
+
+ # Should call AI message processing
+ mock_process_ai.assert_called_once_with(
+ payload["content"],
+ payload["app_id"],
+ payload["user_id"],
+ payload["session_id"],
+ payload["message_id"],
+ db_session
+ )
+ mock_process_user.assert_not_called()
+ mock_summarize.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_process_item_calls_summarize(self, db_session):
+ """Test that process_item always calls summarize_if_needed."""
+ payload = {
+ "message_id": str(uuid4()),
+ "is_user": True,
+ "content": "Test message",
+ "app_id": str(uuid4()),
+ "user_id": str(uuid4()),
+ "session_id": str(uuid4())
+ }
+
+ with (
+ patch("src.deriver.consumer.process_user_message") as mock_process_user,
+ patch("src.deriver.consumer.summarize_if_needed") as mock_summarize
+ ):
+ mock_process_user.return_value = None
+ mock_summarize.return_value = None
+
+ await consumer.process_item(db_session, payload)
+
+ mock_summarize.assert_called_once_with(
+ db_session,
+ payload["app_id"],
+ payload["session_id"],
+ payload["user_id"],
+ payload["message_id"]
+ )
+
+
+class TestProcessUserMessage:
+ """Test user message processing functionality."""
+
+ @pytest_asyncio.fixture
+ async def setup_user_message_test(self, db_session, sample_data):
+ """Setup test data for user message processing."""
+ test_app, test_user = sample_data
+
+ # Create a session
+ session = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={}
+ )
+ db_session.add(session)
+ await db_session.flush()
+
+ # Create a collection for the user
+ collection = models.Collection(
+ app_id=test_app.public_id,
+ user_id=test_user.public_id,
+ name=f"user_{test_user.public_id}",
+ metadata={"type": "user_facts"}
+ )
+ db_session.add(collection)
+ await db_session.flush()
+
+ return test_app, test_user, session, collection
+
+ @pytest.mark.asyncio
+ async def test_process_user_message_extracts_and_saves_facts(self, db_session, setup_user_message_test, mock_llm_responses):
+ """Test that user message processing extracts and saves facts."""
+ test_app, test_user, session, collection = setup_user_message_test
+
+ message_content = "I'm a Python developer who works remotely and loves coffee"
+ message_id = str(uuid4())
+
+ with (
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_get_history,
+ patch("src.deriver.consumer.extract_facts_long_term") as mock_extract_facts,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection,
+ patch("src.deriver.consumer.CollectionEmbeddingStore") as mock_embedding_store_class
+ ):
+ # Mock history retrieval
+ mock_get_history.return_value = ("Previous chat", [], None)
+
+ # Mock fact extraction
+ mock_extract_facts.return_value = [
+ "User is a Python developer",
+ "User works remotely",
+ "User loves coffee"
+ ]
+
+ # Mock collection retrieval
+ mock_get_collection.return_value = collection
+
+ # Mock embedding store
+ mock_embedding_store = AsyncMock()
+ mock_embedding_store.remove_duplicates.return_value = [
+ "User is a Python developer",
+ "User works remotely"
+ ] # Simulate one duplicate removed
+ mock_embedding_store.save_facts.return_value = None
+ mock_embedding_store_class.return_value = mock_embedding_store
+
+ # Process the user message
+ await consumer.process_user_message(
+ message_content,
+ test_app.public_id,
+ test_user.public_id,
+ session.public_id,
+ message_id,
+ db_session
+ )
+
+ # Verify the flow
+ mock_get_history.assert_called_once_with(
+ db_session, session.public_id, summary_type=SummaryType.SHORT
+ )
+ mock_extract_facts.assert_called_once()
+ mock_get_collection.assert_called_once_with(
+ db=db_session, app_id=test_app.public_id, user_id=test_user.public_id
+ )
+ mock_embedding_store.remove_duplicates.assert_called_once_with([
+ "User is a Python developer",
+ "User works remotely",
+ "User loves coffee"
+ ])
+ mock_embedding_store.save_facts.assert_called_once_with(
+ ["User is a Python developer", "User works remotely"],
+ message_id=message_id
+ )
+
+ @pytest.mark.asyncio
+ async def test_process_user_message_no_unique_facts(self, db_session, setup_user_message_test):
+ """Test user message processing when all facts are duplicates."""
+ test_app, test_user, session, collection = setup_user_message_test
+
+ message_content = "I still love Python programming"
+ message_id = str(uuid4())
+
+ with (
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_get_history,
+ patch("src.deriver.consumer.extract_facts_long_term") as mock_extract_facts,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection,
+ patch("src.deriver.consumer.CollectionEmbeddingStore") as mock_embedding_store_class
+ ):
+ mock_get_history.return_value = ("Previous chat", [], None)
+ mock_extract_facts.return_value = ["User loves Python"]
+ mock_get_collection.return_value = collection
+
+ # Mock embedding store to return no unique facts
+ mock_embedding_store = AsyncMock()
+ mock_embedding_store.remove_duplicates.return_value = [] # All duplicates
+ mock_embedding_store_class.return_value = mock_embedding_store
+
+ await consumer.process_user_message(
+ message_content,
+ test_app.public_id,
+ test_user.public_id,
+ session.public_id,
+ message_id,
+ db_session
+ )
+
+ # Should not call save_facts when no unique facts
+ mock_embedding_store.save_facts.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_process_user_message_with_chat_history(self, db_session, setup_user_message_test):
+ """Test that chat history is properly included in fact extraction."""
+ test_app, test_user, session, collection = setup_user_message_test
+
+ message_content = "I prefer PyTorch over TensorFlow"
+ message_id = str(uuid4())
+
+ with (
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_get_history,
+ patch("src.deriver.consumer.extract_facts_long_term") as mock_extract_facts,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection,
+ patch("src.deriver.consumer.CollectionEmbeddingStore") as mock_embedding_store_class
+ ):
+ # Mock history with previous context
+ mock_get_history.return_value = (
+ "AI: Hello! How can I help?\nhuman: I'm a machine learning engineer",
+ [],
+ None
+ )
+ mock_extract_facts.return_value = ["User prefers PyTorch"]
+ mock_get_collection.return_value = collection
+
+ mock_embedding_store = AsyncMock()
+ mock_embedding_store.remove_duplicates.return_value = ["User prefers PyTorch"]
+ mock_embedding_store_class.return_value = mock_embedding_store
+
+ await consumer.process_user_message(
+ message_content,
+ test_app.public_id,
+ test_user.public_id,
+ session.public_id,
+ message_id,
+ db_session
+ )
+
+ # Verify extract_facts was called with combined history + current message
+ expected_history = "AI: Hello! How can I help?\nhuman: I'm a machine learning engineer\nhuman: I prefer PyTorch over TensorFlow"
+ mock_extract_facts.assert_called_once_with(expected_history)
+
+ @pytest.mark.asyncio
+ async def test_process_user_message_handles_extraction_error(self, db_session, setup_user_message_test):
+ """Test that user message processing handles fact extraction errors gracefully."""
+ test_app, test_user, session, collection = setup_user_message_test
+
+ message_content = "Test message"
+ message_id = str(uuid4())
+
+ with (
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_get_history,
+ patch("src.deriver.consumer.extract_facts_long_term") as mock_extract_facts,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection
+ ):
+ mock_get_history.return_value = ("", [], None)
+ mock_extract_facts.side_effect = Exception("LLM API error")
+ mock_get_collection.return_value = collection
+
+ # Should raise the exception (let caller handle it)
+ with pytest.raises(Exception, match="LLM API error"):
+ await consumer.process_user_message(
+ message_content,
+ test_app.public_id,
+ test_user.public_id,
+ session.public_id,
+ message_id,
+ db_session
+ )
+
+
+class TestProcessAIMessage:
+ """Test AI message processing functionality."""
+
+ @pytest.mark.asyncio
+ async def test_process_ai_message_basic_functionality(self, db_session):
+ """Test basic AI message processing (currently just console output)."""
+ content = "I can help you with Python programming!"
+ app_id = str(uuid4())
+ user_id = str(uuid4())
+ session_id = str(uuid4())
+ message_id = str(uuid4())
+
+ # Mock console output
+ with patch("src.deriver.consumer.console.print") as mock_print:
+ await consumer.process_ai_message(
+ content, app_id, user_id, session_id, message_id, db_session
+ )
+
+ # Should print the AI message content
+ mock_print.assert_called_once_with(
+ f"Processing AI message: {content}",
+ style="bright_magenta"
+ )
+
+
+class TestSummarizeIfNeeded:
+ """Test summary generation functionality."""
+
+ @pytest_asyncio.fixture
+ async def setup_summary_test(self, db_session, sample_data):
+ """Setup test data for summary testing."""
+ test_app, test_user = sample_data
+
+ session = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={}
+ )
+ db_session.add(session)
+ await db_session.flush()
+
+ return test_app, test_user, session
+
+ @pytest.mark.asyncio
+ async def test_summarize_if_needed_no_summary_required(self, db_session, setup_summary_test):
+ """Test when no summary is needed."""
+ test_app, test_user, session = setup_summary_test
+ message_id = str(uuid4())
+
+ with patch("src.deriver.consumer.history.should_create_summary") as mock_should_create:
+ # Mock that no summary is needed
+ mock_should_create.return_value = (False, [], None)
+
+ await consumer.summarize_if_needed(
+ db_session,
+ test_app.public_id,
+ session.public_id,
+ test_user.public_id,
+ message_id
+ )
+
+ # Should only check for short summary
+ mock_should_create.assert_called_once_with(
+ db_session, session.public_id, summary_type=SummaryType.SHORT
+ )
+
+ @pytest.mark.asyncio
+ async def test_summarize_if_needed_short_summary_only(self, db_session, setup_summary_test):
+ """Test creating only a short summary."""
+ test_app, test_user, session = setup_summary_test
+ message_id = str(uuid4())
+
+ # Mock messages for short summary
+ mock_messages = [
+ MagicMock(id=1, content="Message 1"),
+ MagicMock(id=2, content="Message 2")
+ ]
+
+ with (
+ patch("src.deriver.consumer.history.should_create_summary") as mock_should_create,
+ patch("src.deriver.consumer.history.create_summary") as mock_create_summary,
+ patch("src.deriver.consumer.history.save_summary_metamessage") as mock_save_summary
+ ):
+ # Mock summary check responses
+ def mock_should_create_side_effect(db, session_id, summary_type):
+ if summary_type == SummaryType.SHORT:
+ return (True, mock_messages, None) # Need short summary
+ else:
+ return (False, [], None) # Don't need long summary
+
+ mock_should_create.side_effect = mock_should_create_side_effect
+ mock_create_summary.return_value = "Short summary of recent messages"
+ mock_save_summary.return_value = None
+
+ await consumer.summarize_if_needed(
+ db_session,
+ test_app.public_id,
+ session.public_id,
+ test_user.public_id,
+ message_id
+ )
+
+ # Should check for both short and long summaries
+ assert mock_should_create.call_count == 2
+
+ # Should create one short summary
+ mock_create_summary.assert_called_once_with(
+ messages=mock_messages,
+ previous_summary=None,
+ summary_type=SummaryType.SHORT
+ )
+
+ # Should save the short summary
+ mock_save_summary.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_summarize_if_needed_both_summaries(self, db_session, setup_summary_test):
+ """Test creating both short and long summaries."""
+ test_app, test_user, session = setup_summary_test
+ message_id = str(uuid4())
+
+ # Mock messages for summaries
+ mock_short_messages = [MagicMock(id=i, content=f"Message {i}") for i in range(1, 11)]
+ mock_long_messages = [MagicMock(id=i, content=f"Message {i}") for i in range(1, 61)]
+
+ with (
+ patch("src.deriver.consumer.history.should_create_summary") as mock_should_create,
+ patch("src.deriver.consumer.history.create_summary") as mock_create_summary,
+ patch("src.deriver.consumer.history.save_summary_metamessage") as mock_save_summary
+ ):
+ # Mock summary check responses
+ def mock_should_create_side_effect(db, session_id, summary_type):
+ if summary_type == SummaryType.SHORT:
+ return (True, mock_short_messages, None)
+ else:
+ return (True, mock_long_messages, None)
+
+ mock_should_create.side_effect = mock_should_create_side_effect
+
+ # Mock summary creation
+ def mock_create_summary_side_effect(messages, previous_summary, summary_type):
+ if summary_type == SummaryType.LONG:
+ return "Long summary of conversation"
+ else:
+ return "Short summary of recent messages"
+
+ mock_create_summary.side_effect = mock_create_summary_side_effect
+
+ # Mock save returning the long summary for short summary context
+ mock_long_summary_obj = MagicMock()
+ mock_long_summary_obj.content = "Long summary of conversation"
+
+ def mock_save_side_effect(db, app_id, user_id, session_id, message_id, summary_content, message_count, summary_type):
+ if summary_type == SummaryType.LONG:
+ return mock_long_summary_obj
+ return None
+
+ mock_save_summary.side_effect = mock_save_side_effect
+
+ await consumer.summarize_if_needed(
+ db_session,
+ test_app.public_id,
+ session.public_id,
+ test_user.public_id,
+ message_id
+ )
+
+ # Should create both summaries
+ assert mock_create_summary.call_count == 2
+
+ # Should save both summaries
+ assert mock_save_summary.call_count == 2
+
+ @pytest.mark.asyncio
+ async def test_summarize_if_needed_handles_summary_creation_error(self, db_session, setup_summary_test):
+ """Test that summary creation errors are handled gracefully."""
+ test_app, test_user, session = setup_summary_test
+ message_id = str(uuid4())
+
+ mock_messages = [MagicMock(id=1, content="Message 1")]
+
+ with (
+ patch("src.deriver.consumer.history.should_create_summary") as mock_should_create,
+ patch("src.deriver.consumer.history.create_summary") as mock_create_summary
+ ):
+ mock_should_create.return_value = (True, mock_messages, None)
+ mock_create_summary.side_effect = Exception("LLM API error")
+
+ # Should not raise exception (errors are logged)
+ await consumer.summarize_if_needed(
+ db_session,
+ test_app.public_id,
+ session.public_id,
+ test_user.public_id,
+ message_id
+ )
+
+ @pytest.mark.asyncio
+ async def test_summarize_if_needed_with_existing_long_summary(self, db_session, setup_summary_test):
+ """Test short summary creation with existing long summary context."""
+ test_app, test_user, session = setup_summary_test
+ message_id = str(uuid4())
+
+ mock_messages = [MagicMock(id=1, content="Message 1")]
+ mock_existing_long_summary = MagicMock()
+ mock_existing_long_summary.content = "Existing long summary"
+
+ with (
+ patch("src.deriver.consumer.history.should_create_summary") as mock_should_create,
+ patch("src.deriver.consumer.history.create_summary") as mock_create_summary,
+ patch("src.deriver.consumer.history.save_summary_metamessage") as mock_save_summary
+ ):
+ # Mock that we need short summary and have existing long summary
+ def mock_should_create_side_effect(db, session_id, summary_type):
+ if summary_type == SummaryType.SHORT:
+ return (True, mock_messages, mock_existing_long_summary)
+ else:
+ return (False, [], mock_existing_long_summary)
+
+ mock_should_create.side_effect = mock_should_create_side_effect
+ mock_create_summary.return_value = "Short summary with context"
+ mock_save_summary.return_value = None
+
+ await consumer.summarize_if_needed(
+ db_session,
+ test_app.public_id,
+ session.public_id,
+ test_user.public_id,
+ message_id
+ )
+
+ # Should create short summary with existing long summary as context
+ mock_create_summary.assert_called_once_with(
+ messages=mock_messages,
+ previous_summary="Existing long summary",
+ summary_type=SummaryType.SHORT
+ )
+
+
+class TestEnvironmentConfiguration:
+ """Test environment variable configuration."""
+
+ def test_tom_method_default(self):
+ """Test TOM_METHOD defaults to single_prompt."""
+ # Test the getenv behavior that the module uses
+ import os
+ default_value = os.getenv("TOM_METHOD", "single_prompt")
+ # If no environment variable is set, should use default
+ if os.getenv("TOM_METHOD") is None:
+ assert default_value == "single_prompt"
+ else:
+ # If environment variable is set, respect it
+ assert default_value == os.getenv("TOM_METHOD")
+
+ def test_tom_method_custom(self):
+ """Test TOM_METHOD can be customized via environment."""
+ # Test that the os.getenv pattern works correctly
+ import os
+ # Simulate the pattern used in the consumer module
+ test_value = os.getenv("TOM_METHOD", "single_prompt")
+ # The result should be either the env var or the default
+ assert test_value in ["single_prompt", "conversational", "long_term"]
+
+ def test_user_representation_method_default(self):
+ """Test USER_REPRESENTATION_METHOD defaults to long_term."""
+ # Test the getenv behavior that the module uses
+ import os
+ default_value = os.getenv("USER_REPRESENTATION_METHOD", "long_term")
+ # If no environment variable is set, should use default
+ if os.getenv("USER_REPRESENTATION_METHOD") is None:
+ assert default_value == "long_term"
+ else:
+ # If environment variable is set, respect it
+ assert default_value == os.getenv("USER_REPRESENTATION_METHOD")
\ No newline at end of file
diff --git a/tests/deriver/test_performance.py b/tests/deriver/test_performance.py
new file mode 100644
index 00000000..30d9e79a
--- /dev/null
+++ b/tests/deriver/test_performance.py
@@ -0,0 +1,263 @@
+"""Performance tests for the deriver system."""
+
+import asyncio
+import time
+from unittest.mock import AsyncMock, patch
+
+import pytest
+import pytest_asyncio
+
+from src.deriver import consumer
+from src.deriver.tom.embeddings import CollectionEmbeddingStore
+
+
+class TestPerformanceValidation:
+ """Test performance characteristics of deriver components."""
+
+ @pytest.mark.asyncio
+ async def test_fact_extraction_performance(self, performance_config):
+ """Test that fact extraction completes within reasonable time limits."""
+ start_time = time.time()
+
+ # Mock a reasonably complex chat history
+ chat_history = """
+ User: Hi, I'm Alex, a senior software engineer working at Google in the machine learning team.
+ AI: Hello Alex! That sounds like an exciting role. What kind of ML projects are you working on?
+ User: I'm primarily focused on building recommendation systems using TensorFlow and PyTorch.
+ We handle millions of user interactions daily and need to provide real-time personalized recommendations.
+ AI: That's impressive scale! How do you handle the computational challenges?
+ User: We use a distributed architecture with Kubernetes, Redis for caching, and BigQuery for data processing.
+ The team also experiments with newer frameworks like JAX for research prototypes.
+ """
+
+ with patch("src.deriver.tom.long_term.ModelClient") as mock_client:
+ mock_instance = AsyncMock()
+ mock_instance.generate.return_value = '{"facts": ["User is Alex", "User works at Google", "User uses TensorFlow"]}'
+ mock_client.return_value = mock_instance
+
+ # Extract facts and measure time
+ from src.deriver.tom.long_term import extract_facts_long_term
+ facts = await extract_facts_long_term(chat_history)
+
+ extraction_time = time.time() - start_time
+
+ # Verify performance meets requirements
+ assert extraction_time < performance_config["fact_extraction_time_limit"]
+ assert len(facts) > 0
+ print(f"✅ Fact extraction completed in {extraction_time:.3f}s")
+
+ @pytest.mark.asyncio
+ async def test_embedding_operations_performance(self, sample_data, performance_config):
+ """Test that embedding operations complete efficiently."""
+ test_app, test_user = sample_data
+
+ # Create embedding store
+ collection_id = "test_collection"
+ store = CollectionEmbeddingStore(test_app.public_id, test_user.public_id, collection_id)
+
+ # Test data
+ facts = [
+ "User is a machine learning engineer",
+ "User works with large-scale systems",
+ "User has expertise in TensorFlow and PyTorch",
+ "User handles millions of daily interactions",
+ "User uses distributed computing"
+ ]
+
+ with (
+ patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db,
+ patch("src.deriver.tom.embeddings.crud.get_duplicate_documents") as mock_get_dupes,
+ patch("src.deriver.tom.embeddings.crud.create_document") as mock_create_doc
+ ):
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ mock_get_dupes.return_value = [] # No duplicates
+ mock_create_doc.return_value = None
+
+ # Test duplicate removal performance
+ start_time = time.time()
+ unique_facts = await store.remove_duplicates(facts)
+ dedup_time = time.time() - start_time
+
+ # Test fact saving performance
+ start_time = time.time()
+ await store.save_facts(unique_facts)
+ save_time = time.time() - start_time
+
+ # Verify performance
+ total_time = dedup_time + save_time
+ assert total_time < 2.0 # Should complete within 2 seconds
+ assert len(unique_facts) == len(facts) # All facts should be unique
+
+ print(f"✅ Embedding operations completed in {total_time:.3f}s")
+ print(f" - Deduplication: {dedup_time:.3f}s")
+ print(f" - Fact saving: {save_time:.3f}s")
+
+ @pytest.mark.asyncio
+ async def test_concurrent_processing_performance(self, sample_data, performance_config):
+ """Test performance under concurrent load."""
+ test_app, test_user = sample_data
+
+ # Create multiple simulated messages
+ messages = [
+ f"Message {i}: User sharing information about their work and interests"
+ for i in range(performance_config["message_count"] // 10) # Smaller load for test
+ ]
+
+ # Mock all the dependencies for speed
+ with (
+ patch("src.deriver.consumer.extract_facts_long_term") as mock_extract,
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_history,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection,
+ patch("src.deriver.consumer.CollectionEmbeddingStore") as mock_store_class,
+ patch("src.deriver.consumer.summarize_if_needed") as mock_summarize
+ ):
+ # Setup fast mocks
+ mock_extract.return_value = ["User fact"]
+ mock_history.return_value = ("", [], None)
+ mock_get_collection.return_value = AsyncMock()
+
+ mock_store = AsyncMock()
+ mock_store.remove_duplicates.return_value = ["User fact"]
+ mock_store.save_facts.return_value = None
+ mock_store_class.return_value = mock_store
+
+ mock_summarize.return_value = None
+
+ # Process messages concurrently
+ start_time = time.time()
+
+ async def process_single_message(content):
+ await consumer.process_user_message(
+ content,
+ test_app.public_id,
+ test_user.public_id,
+ "session_123",
+ f"msg_{hash(content)}",
+ AsyncMock() # Mock DB session
+ )
+
+ # Run concurrent processing
+ tasks = [process_single_message(msg) for msg in messages]
+ await asyncio.gather(*tasks)
+
+ total_time = time.time() - start_time
+
+ # Verify performance
+ messages_per_second = len(messages) / total_time
+ assert messages_per_second > 5 # Should process at least 5 messages per second
+
+ print(f"✅ Processed {len(messages)} messages in {total_time:.3f}s")
+ print(f" - Rate: {messages_per_second:.1f} messages/second")
+
+ @pytest.mark.asyncio
+ async def test_memory_usage_stability(self, sample_data):
+ """Test that memory usage remains stable during processing."""
+ test_app, test_user = sample_data
+
+ # Simulate processing many messages to check for memory leaks
+ message_count = 50
+
+ with (
+ patch("src.deriver.consumer.extract_facts_long_term") as mock_extract,
+ patch("src.deriver.consumer.history.get_summarized_history") as mock_history,
+ patch("src.deriver.consumer.crud.get_or_create_user_protected_collection") as mock_get_collection,
+ patch("src.deriver.consumer.CollectionEmbeddingStore") as mock_store_class,
+ patch("src.deriver.consumer.summarize_if_needed") as mock_summarize
+ ):
+ # Setup mocks
+ mock_extract.return_value = ["Fact"]
+ mock_history.return_value = ("", [], None)
+ mock_get_collection.return_value = AsyncMock()
+
+ mock_store = AsyncMock()
+ mock_store.remove_duplicates.return_value = ["Fact"]
+ mock_store.save_facts.return_value = None
+ mock_store_class.return_value = mock_store
+
+ mock_summarize.return_value = None
+
+ # Process messages in batches to simulate sustained load
+ for batch in range(5): # 5 batches of 10 messages each
+ batch_tasks = []
+ for i in range(10):
+ task = consumer.process_user_message(
+ f"Batch {batch} Message {i}: User information",
+ test_app.public_id,
+ test_user.public_id,
+ f"session_{batch}",
+ f"msg_{batch}_{i}",
+ AsyncMock()
+ )
+ batch_tasks.append(task)
+
+ # Process batch
+ await asyncio.gather(*batch_tasks)
+
+ # Small delay between batches
+ await asyncio.sleep(0.01)
+
+ # Verify all processing completed successfully
+ assert mock_extract.call_count == message_count
+ print(f"✅ Processed {message_count} messages in batches successfully")
+
+ def test_configuration_performance_settings(self, performance_config):
+ """Test that performance configuration is reasonable."""
+ # Verify performance thresholds are achievable
+ assert performance_config["fact_extraction_time_limit"] >= 1.0
+ assert performance_config["tom_inference_time_limit"] >= 1.0
+ assert performance_config["queue_processing_time_limit"] >= 0.1
+ assert performance_config["max_workers"] >= 1
+ assert performance_config["timeout_seconds"] >= 10
+
+ print("✅ Performance configuration validated")
+ print(f" - Fact extraction limit: {performance_config['fact_extraction_time_limit']}s")
+ print(f" - TOM inference limit: {performance_config['tom_inference_time_limit']}s")
+ print(f" - Max workers: {performance_config['max_workers']}")
+
+
+class TestScalabilityValidation:
+ """Test scalability characteristics."""
+
+ @pytest.mark.asyncio
+ async def test_fact_storage_scalability(self, sample_data):
+ """Test that fact storage can handle larger volumes."""
+ test_app, test_user = sample_data
+
+ # Simulate storing many facts
+ large_fact_list = [f"User fact number {i}" for i in range(100)]
+
+ collection_id = "test_scalability"
+ store = CollectionEmbeddingStore(test_app.public_id, test_user.public_id, collection_id)
+
+ with (
+ patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db,
+ patch("src.deriver.tom.embeddings.crud.get_duplicate_documents") as mock_get_dupes,
+ patch("src.deriver.tom.embeddings.crud.create_document") as mock_create_doc
+ ):
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ mock_get_dupes.return_value = []
+ mock_create_doc.return_value = None
+
+ start_time = time.time()
+
+ # Test processing in chunks
+ chunk_size = 20
+ for i in range(0, len(large_fact_list), chunk_size):
+ chunk = large_fact_list[i:i+chunk_size]
+ unique_facts = await store.remove_duplicates(chunk)
+ await store.save_facts(unique_facts)
+
+ total_time = time.time() - start_time
+
+ # Should handle 100 facts efficiently
+ assert total_time < 5.0
+ assert mock_create_doc.call_count == len(large_fact_list)
+
+ print(f"✅ Processed {len(large_fact_list)} facts in {total_time:.3f}s")
+ print(f" - Rate: {len(large_fact_list)/total_time:.1f} facts/second")
\ No newline at end of file
diff --git a/tests/deriver/test_queue_manager.py b/tests/deriver/test_queue_manager.py
new file mode 100644
index 00000000..759e23b6
--- /dev/null
+++ b/tests/deriver/test_queue_manager.py
@@ -0,0 +1,573 @@
+"""Tests for the QueueManager class and queue processing functionality."""
+
+import asyncio
+import signal
+from datetime import datetime, timedelta, timezone
+from unittest.mock import AsyncMock, MagicMock, patch
+from uuid import uuid4
+
+import pytest
+import pytest_asyncio
+from sqlalchemy import select
+
+from src import models
+from src.deriver.queue import QueueManager
+
+
+class TestQueueManagerInitialization:
+ """Test QueueManager initialization and configuration."""
+
+ def test_queue_manager_default_initialization(self):
+ """Test QueueManager initializes with default values."""
+ with patch("src.deriver.queue.os.getenv") as mock_getenv:
+ mock_getenv.return_value = "1" # Default worker count
+
+ manager = QueueManager()
+
+ assert manager.workers == 1
+ assert manager.semaphore._value == 1
+ assert not manager.shutdown_event.is_set()
+ assert len(manager.active_tasks) == 0
+ assert len(manager.owned_sessions) == 0
+
+ def test_queue_manager_custom_workers(self):
+ """Test QueueManager respects DERIVER_WORKERS environment variable."""
+ with patch("src.deriver.queue.os.getenv") as mock_getenv:
+ mock_getenv.return_value = "4"
+
+ manager = QueueManager()
+
+ assert manager.workers == 4
+ assert manager.semaphore._value == 4
+
+ @patch("src.deriver.queue.sentry_sdk")
+ def test_sentry_initialization_enabled(self, mock_sentry):
+ """Test Sentry initialization when enabled."""
+ with patch("src.deriver.queue.os.getenv") as mock_getenv:
+ def getenv_side_effect(key, default=None):
+ if key == "SENTRY_ENABLED":
+ return "True"
+ elif key == "SENTRY_DSN":
+ return "https://test@sentry.io/123"
+ elif key == "DERIVER_WORKERS":
+ return "1"
+ return default
+
+ mock_getenv.side_effect = getenv_side_effect
+
+ QueueManager()
+
+ mock_sentry.init.assert_called_once()
+
+ @patch("src.deriver.queue.sentry_sdk")
+ def test_sentry_initialization_disabled(self, mock_sentry):
+ """Test Sentry is not initialized when disabled."""
+ with patch("src.deriver.queue.os.getenv") as mock_getenv:
+ def getenv_side_effect(key, default=None):
+ if key == "SENTRY_ENABLED":
+ return "False"
+ elif key == "DERIVER_WORKERS":
+ return "1"
+ return default
+
+ mock_getenv.side_effect = getenv_side_effect
+
+ QueueManager()
+
+ mock_sentry.init.assert_not_called()
+
+
+class TestTaskAndSessionTracking:
+ """Test task and session tracking functionality."""
+
+ def test_add_task_tracking(self):
+ """Test adding tasks to tracking set."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+
+ # Create a mock task
+ task = MagicMock()
+ task.add_done_callback = MagicMock()
+
+ manager.add_task(task)
+
+ assert task in manager.active_tasks
+ task.add_done_callback.assert_called_once()
+
+ def test_session_tracking(self):
+ """Test session tracking and untracking."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+ session_id = 123
+
+ # Track session
+ manager.track_session(session_id)
+ assert session_id in manager.owned_sessions
+
+ # Untrack session
+ manager.untrack_session(session_id)
+ assert session_id not in manager.owned_sessions
+
+ def test_track_session_multiple(self):
+ """Test tracking multiple sessions."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+ session_ids = [123, 456, 789]
+
+ for session_id in session_ids:
+ manager.track_session(session_id)
+
+ assert all(sid in manager.owned_sessions for sid in session_ids)
+ assert len(manager.owned_sessions) == 3
+
+
+class TestDatabaseOperations:
+ """Test database operations for queue management."""
+
+ @pytest_asyncio.fixture
+ async def setup_queue_data(self, db_session, sample_data):
+ """Setup test data for queue operations."""
+ test_app, test_user = sample_data
+
+ # Create sessions
+ session1 = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={}
+ )
+ session2 = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={}
+ )
+ db_session.add_all([session1, session2])
+ await db_session.flush()
+
+ # Create queue items (use integer session.id, not public_id)
+ queue_item1 = models.QueueItem(
+ session_id=session1.id,
+ payload={"message_id": str(uuid4())},
+ processed=False
+ )
+ queue_item2 = models.QueueItem(
+ session_id=session2.id,
+ payload={"message_id": str(uuid4())},
+ processed=False
+ )
+ queue_item3 = models.QueueItem(
+ session_id=session1.id,
+ payload={"message_id": str(uuid4())},
+ processed=True # Already processed
+ )
+
+ db_session.add_all([queue_item1, queue_item2, queue_item3])
+ await db_session.flush()
+
+ return session1, session2, [queue_item1, queue_item2, queue_item3]
+
+ @pytest.mark.asyncio
+ async def test_get_available_sessions(self, db_session, setup_queue_data):
+ """Test getting available sessions for processing."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+ session1, session2, queue_items = setup_queue_data
+
+ # Get available sessions
+ available_sessions = await manager.get_available_sessions(db_session)
+
+ # Should return sessions with unprocessed items
+ assert len(available_sessions) == 1 # Limited to 1 by the query
+ assert available_sessions[0] in [session1.id, session2.id]
+
+ @pytest.mark.asyncio
+ async def test_get_available_sessions_with_active_session(self, db_session, setup_queue_data):
+ """Test that active sessions are excluded from available sessions."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+ session1, session2, queue_items = setup_queue_data
+
+ # Mark session1 as active
+ active_session = models.ActiveQueueSession(session_id=session1.id)
+ db_session.add(active_session)
+ await db_session.flush()
+
+ # Get available sessions
+ available_sessions = await manager.get_available_sessions(db_session)
+
+ # Should only return session2
+ assert len(available_sessions) == 1
+ assert available_sessions[0] == session2.id
+
+ @pytest.mark.asyncio
+ async def test_stale_session_cleanup(self, db_session, setup_queue_data):
+ """Test cleanup of stale active sessions."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+ session1, session2, queue_items = setup_queue_data
+
+ # Create a stale active session (older than 5 minutes)
+ stale_time = datetime.now(timezone.utc) - timedelta(minutes=10)
+ stale_session = models.ActiveQueueSession(
+ session_id=session1.id,
+ last_updated=stale_time
+ )
+ db_session.add(stale_session)
+ await db_session.flush()
+
+ # Get available sessions (this should trigger cleanup)
+ available_sessions = await manager.get_available_sessions(db_session)
+
+ # Stale session should be cleaned up, making session1 available
+ result = await db_session.execute(
+ select(models.ActiveQueueSession).where(
+ models.ActiveQueueSession.session_id == session1.id
+ )
+ )
+ assert result.scalar_one_or_none() is None
+
+ @pytest.mark.asyncio
+ async def test_get_next_message(self, db_session, setup_queue_data):
+ """Test getting the next unprocessed message for a session."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+ session1, session2, queue_items = setup_queue_data
+
+ # Get next message for session1
+ next_message = await manager.get_next_message(db_session, session1.id)
+
+ # Should return the unprocessed message
+ assert next_message is not None
+ assert next_message.session_id == session1.id
+ assert not next_message.processed
+
+ @pytest.mark.asyncio
+ async def test_get_next_message_no_unprocessed(self, db_session, setup_queue_data):
+ """Test getting next message when all are processed."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+ session1, session2, queue_items = setup_queue_data
+
+ # Mark all messages as processed
+ for item in queue_items:
+ item.processed = True
+ await db_session.flush()
+
+ # Get next message
+ next_message = await manager.get_next_message(db_session, session1.id)
+
+ # Should return None
+ assert next_message is None
+
+
+class TestConcurrencyControl:
+ """Test concurrency control and semaphore behavior."""
+
+ @pytest.mark.asyncio
+ async def test_semaphore_limits_concurrent_processing(self, mock_semaphore):
+ """Test that semaphore properly limits concurrent session processing."""
+ with patch("src.deriver.queue.os.getenv", return_value="2"):
+ with patch("asyncio.Semaphore") as mock_semaphore_class:
+ mock_semaphore_class.return_value = mock_semaphore
+ manager = QueueManager()
+
+ # Mock the process_session method to return actual async function
+ async def mock_process_session(session_id):
+ async with manager.semaphore:
+ await asyncio.sleep(0.01) # Simulate work
+
+ with patch.object(manager, 'process_session', side_effect=mock_process_session):
+ # Try to process multiple sessions
+ tasks = []
+ for i in range(5):
+ task = asyncio.create_task(manager.process_session(i))
+ tasks.append(task)
+ manager.add_task(task)
+
+ # Wait for tasks to complete
+ await asyncio.gather(*tasks, return_exceptions=True)
+
+ # Verify semaphore was used
+ assert mock_semaphore.__aenter__.call_count == 5
+
+ @pytest.mark.asyncio
+ async def test_polling_loop_respects_semaphore_capacity(self):
+ """Test that polling loop waits when all workers are busy."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+
+ # Mock semaphore as locked (no capacity)
+ manager.semaphore.locked = MagicMock(return_value=True)
+
+ # Mock tracked_db to avoid database operations
+ with patch("src.deriver.queue.tracked_db"):
+ # Set shutdown event after a short delay to exit the loop
+ async def set_shutdown():
+ await asyncio.sleep(0.1)
+ manager.shutdown_event.set()
+
+ asyncio.create_task(set_shutdown())
+
+ # Run polling loop
+ await manager.polling_loop()
+
+ # Should have checked semaphore status
+ manager.semaphore.locked.assert_called()
+
+
+class TestSignalHandling:
+ """Test signal handling and graceful shutdown."""
+
+ @pytest.mark.asyncio
+ async def test_shutdown_signal_handling(self, mock_signal_handling):
+ """Test that shutdown properly handles signals."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+
+ # Create actual async tasks instead of AsyncMock
+ async def dummy_task():
+ await asyncio.sleep(0.01)
+
+ task1 = asyncio.create_task(dummy_task())
+ task2 = asyncio.create_task(dummy_task())
+ manager.active_tasks = {task1, task2}
+
+ # Call shutdown
+ await manager.shutdown(signal.SIGTERM)
+
+ # Shutdown event should be set
+ assert manager.shutdown_event.is_set()
+
+ @pytest.mark.asyncio
+ async def test_cleanup_owned_sessions(self, db_session):
+ """Test cleanup of owned sessions during shutdown."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+
+ # Add owned sessions
+ session_ids = [123, 456, 789]
+ for session_id in session_ids:
+ manager.track_session(session_id)
+ # Create corresponding active session records
+ active_session = models.ActiveQueueSession(session_id=session_id)
+ db_session.add(active_session)
+
+ await db_session.flush()
+
+ # Mock tracked_db to use our test session
+ with patch("src.deriver.queue.tracked_db") as mock_tracked_db:
+ mock_tracked_db.return_value.__aenter__.return_value = db_session
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ # Run cleanup
+ await manager.cleanup()
+
+ # Verify sessions were removed from database
+ result = await db_session.execute(
+ select(models.ActiveQueueSession).where(
+ models.ActiveQueueSession.session_id.in_(session_ids)
+ )
+ )
+ remaining_sessions = result.scalars().all()
+ assert len(remaining_sessions) == 0
+
+ @pytest.mark.asyncio
+ async def test_cleanup_with_database_error(self, db_session):
+ """Test cleanup handles database errors gracefully."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+ manager.track_session(123)
+
+ # Mock tracked_db to raise an exception
+ with patch("src.deriver.queue.tracked_db") as mock_tracked_db:
+ mock_tracked_db.side_effect = Exception("Database connection failed")
+
+ # Cleanup should not raise exception
+ await manager.cleanup()
+
+ # Session should still be tracked (cleanup failed)
+ assert 123 in manager.owned_sessions
+
+
+class TestErrorHandling:
+ """Test error handling in various scenarios."""
+
+ @pytest.mark.asyncio
+ async def test_polling_loop_handles_database_errors(self):
+ """Test polling loop handles database errors gracefully."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+
+ # Mock tracked_db as proper async context manager that fails inside the context
+ call_count = 0
+
+ class MockTrackedDBContext:
+ def __init__(self, *args, **kwargs):
+ nonlocal call_count
+ call_count += 1
+
+ async def __aenter__(self):
+ mock_db = MagicMock()
+ # Make get_available_sessions fail on first call
+ if call_count == 1:
+ mock_db.execute.side_effect = Exception("Database connection failed")
+ else:
+ # Set shutdown on second call to exit loop
+ manager.shutdown_event.set()
+ mock_db.execute.return_value = MagicMock()
+ return mock_db
+
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
+ return None
+
+ with patch("src.deriver.queue.tracked_db", MockTrackedDBContext):
+ # Should not raise exception and should retry
+ await manager.polling_loop()
+
+ # Should have attempted multiple calls
+ assert call_count >= 2
+
+ @pytest.mark.asyncio
+ async def test_process_session_marks_failed_messages_as_processed(self, db_session, sample_queue_items):
+ """Test that failed message processing still marks messages as processed."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+ session, messages, queue_items = sample_queue_items
+
+ # Mock process_item to raise an exception
+ with patch("src.deriver.queue.process_item", side_effect=Exception("Processing failed")):
+ with patch("src.deriver.queue.tracked_db") as mock_tracked_db:
+ mock_tracked_db.return_value.__aenter__.return_value = db_session
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ # Process the session
+ await manager.process_session(session.id)
+
+ # All messages should be marked as processed despite the error
+ result = await db_session.execute(
+ select(models.QueueItem).where(
+ models.QueueItem.session_id == session.id
+ )
+ )
+ queue_items_after = result.scalars().all()
+ assert all(item.processed for item in queue_items_after)
+
+ @pytest.mark.asyncio
+ async def test_session_claiming_handles_integrity_error(self, db_session, sample_data):
+ """Test that session claiming handles race conditions gracefully."""
+ test_app, test_user = sample_data
+
+ # Create sessions
+ session1 = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={}
+ )
+ session2 = models.Session(
+ user_id=test_user.public_id,
+ app_id=test_app.public_id,
+ metadata={}
+ )
+ db_session.add_all([session1, session2])
+ await db_session.flush()
+
+ # Create queue items
+ queue_item1 = models.QueueItem(
+ session_id=session1.id,
+ payload={"message_id": str(uuid4())},
+ processed=False
+ )
+ queue_item2 = models.QueueItem(
+ session_id=session2.id,
+ payload={"message_id": str(uuid4())},
+ processed=False
+ )
+ db_session.add_all([queue_item1, queue_item2])
+ await db_session.flush()
+
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+
+ # Create an active session to cause IntegrityError
+ active_session = models.ActiveQueueSession(session_id=session1.id)
+ db_session.add(active_session)
+ await db_session.flush()
+
+ # Try to get available sessions and claim them
+ available_sessions = await manager.get_available_sessions(db_session)
+
+ # Should get session2 (session1 is active)
+ assert len(available_sessions) == 1
+ assert available_sessions[0] == session2.id
+
+
+class TestIntegrationScenarios:
+ """Test integration scenarios and real-world usage patterns."""
+
+ @pytest.mark.asyncio
+ async def test_full_session_processing_cycle(self, db_session, sample_queue_items):
+ """Test complete processing cycle for a session."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+ session, messages, queue_items = sample_queue_items
+
+ # Mock process_item to simulate successful processing
+ with patch("src.deriver.queue.process_item") as mock_process:
+ mock_process.return_value = None
+
+ with patch("src.deriver.queue.tracked_db") as mock_tracked_db:
+ mock_tracked_db.return_value.__aenter__.return_value = db_session
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ # Process the session
+ await manager.process_session(session.id)
+
+ # Verify all user messages were processed
+ user_message_count = len([item for item in queue_items])
+ assert mock_process.call_count == user_message_count
+
+ # Verify session is not in active sessions
+ result = await db_session.execute(
+ select(models.ActiveQueueSession).where(
+ models.ActiveQueueSession.session_id == session.id
+ )
+ )
+ assert result.scalar_one_or_none() is None
+
+ # Verify session is untracked
+ assert session.id not in manager.owned_sessions
+
+ @pytest.mark.asyncio
+ async def test_shutdown_during_processing(self, db_session, sample_queue_items):
+ """Test graceful shutdown while processing messages."""
+ with patch("src.deriver.queue.os.getenv", return_value="1"):
+ manager = QueueManager()
+ session, messages, queue_items = sample_queue_items
+
+ # Mock process_item to be slow and check shutdown event
+ async def slow_process_item(db, payload):
+ await asyncio.sleep(0.1)
+ if manager.shutdown_event.is_set():
+ return
+ # Continue processing
+
+ with patch("src.deriver.queue.process_item", side_effect=slow_process_item):
+ with patch("src.deriver.queue.tracked_db") as mock_tracked_db:
+ mock_tracked_db.return_value.__aenter__.return_value = db_session
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ # Start processing
+ process_task = asyncio.create_task(manager.process_session(session.id))
+
+ # Trigger shutdown after a short delay
+ async def trigger_shutdown():
+ await asyncio.sleep(0.05)
+ manager.shutdown_event.set()
+
+ shutdown_task = asyncio.create_task(trigger_shutdown())
+
+ # Wait for both tasks
+ await asyncio.gather(process_task, shutdown_task, return_exceptions=True)
+
+ # Session should be cleaned up even with shutdown
+ assert session.id not in manager.owned_sessions
+
diff --git a/tests/deriver/test_tom_embeddings.py b/tests/deriver/test_tom_embeddings.py
new file mode 100644
index 00000000..d2e2e4b5
--- /dev/null
+++ b/tests/deriver/test_tom_embeddings.py
@@ -0,0 +1,594 @@
+"""Tests for the TOM embeddings module and CollectionEmbeddingStore."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+from uuid import uuid4
+
+import pytest
+import pytest_asyncio
+
+from src import schemas
+from src.deriver.tom.embeddings import CollectionEmbeddingStore
+
+
+class TestCollectionEmbeddingStoreInitialization:
+ """Test CollectionEmbeddingStore initialization."""
+
+ def test_initialization(self):
+ """Test basic initialization of CollectionEmbeddingStore."""
+ app_id = str(uuid4())
+ user_id = str(uuid4())
+ collection_id = str(uuid4())
+
+ store = CollectionEmbeddingStore(app_id, user_id, collection_id)
+
+ assert store.app_id == app_id
+ assert store.user_id == user_id
+ assert store.collection_id == collection_id
+
+
+class TestSaveFacts:
+ """Test fact saving functionality."""
+
+ @pytest_asyncio.fixture
+ async def embedding_store(self, sample_data):
+ """Create an embedding store for testing."""
+ test_app, test_user = sample_data
+ collection_id = str(uuid4())
+ return CollectionEmbeddingStore(test_app.public_id, test_user.public_id, collection_id)
+
+ @pytest.mark.asyncio
+ async def test_save_facts_basic(self, embedding_store):
+ """Test basic fact saving functionality."""
+ facts = [
+ "User is a Python developer",
+ "User works remotely",
+ "User loves coffee"
+ ]
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.create_document") as mock_create_doc:
+ mock_create_doc.return_value = None
+
+ await embedding_store.save_facts(facts)
+
+ # Should create a document for each fact
+ assert mock_create_doc.call_count == 3
+
+ # Verify each call
+ for i, fact in enumerate(facts):
+ call_args = mock_create_doc.call_args_list[i]
+ assert call_args[1]["app_id"] == embedding_store.app_id
+ assert call_args[1]["user_id"] == embedding_store.user_id
+ assert call_args[1]["collection_id"] == embedding_store.collection_id
+ assert call_args[1]["document"].content == fact
+ assert abs(call_args[1]["duplicate_threshold"] - 0.15) < 1e-10 # 1 - 0.85
+
+ @pytest.mark.asyncio
+ async def test_save_facts_with_message_id(self, embedding_store):
+ """Test saving facts with message ID metadata."""
+ facts = ["User prefers PyTorch over TensorFlow"]
+ message_id = str(uuid4())
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.create_document") as mock_create_doc:
+ mock_create_doc.return_value = None
+
+ await embedding_store.save_facts(facts, message_id=message_id)
+
+ # Verify message_id is included in metadata
+ call_args = mock_create_doc.call_args_list[0]
+ document = call_args[1]["document"]
+ assert document.metadata == {"message_id": message_id}
+
+ @pytest.mark.asyncio
+ async def test_save_facts_custom_similarity_threshold(self, embedding_store):
+ """Test saving facts with custom similarity threshold."""
+ facts = ["User enjoys debugging"]
+ similarity_threshold = 0.9
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.create_document") as mock_create_doc:
+ mock_create_doc.return_value = None
+
+ await embedding_store.save_facts(facts, similarity_threshold=similarity_threshold)
+
+ # Verify duplicate threshold is calculated correctly
+ call_args = mock_create_doc.call_args_list[0]
+ assert abs(call_args[1]["duplicate_threshold"] - 0.1) < 1e-10 # 1 - 0.9
+
+ @pytest.mark.asyncio
+ async def test_save_facts_handles_document_creation_error(self, embedding_store):
+ """Test that fact saving handles document creation errors gracefully."""
+ facts = [
+ "User is a Python developer",
+ "This fact will fail",
+ "User works remotely"
+ ]
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.create_document") as mock_create_doc:
+ # Mock second call to raise an exception
+ def mock_create_side_effect(*args, **kwargs):
+ if "This fact will fail" in str(kwargs.get("document", "")):
+ raise Exception("Database error")
+ return None
+
+ mock_create_doc.side_effect = mock_create_side_effect
+
+ # Should not raise exception (errors are handled gracefully)
+ await embedding_store.save_facts(facts)
+
+ # Should still attempt to create all documents
+ assert mock_create_doc.call_count == 3
+
+ @pytest.mark.asyncio
+ async def test_save_facts_empty_list(self, embedding_store):
+ """Test saving empty fact list."""
+ facts = []
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.create_document") as mock_create_doc:
+ await embedding_store.save_facts(facts)
+
+ # Should not call create_document
+ mock_create_doc.assert_not_called()
+
+
+class TestGetRelevantFacts:
+ """Test fact retrieval functionality."""
+
+ @pytest_asyncio.fixture
+ async def embedding_store(self, sample_data):
+ """Create an embedding store for testing."""
+ test_app, test_user = sample_data
+ collection_id = str(uuid4())
+ return CollectionEmbeddingStore(test_app.public_id, test_user.public_id, collection_id)
+
+ @pytest.mark.asyncio
+ async def test_get_relevant_facts_basic(self, embedding_store):
+ """Test basic fact retrieval functionality."""
+ query = "What programming languages does the user know?"
+
+ # Mock documents returned from query
+ mock_documents = [
+ MagicMock(content="User is proficient in Python"),
+ MagicMock(content="User has experience with JavaScript"),
+ MagicMock(content="User knows SQL")
+ ]
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.query_documents") as mock_query_docs:
+ mock_query_docs.return_value = mock_documents
+
+ facts = await embedding_store.get_relevant_facts(query)
+
+ # Verify query parameters
+ mock_query_docs.assert_called_once_with(
+ mock_db,
+ app_id=embedding_store.app_id,
+ user_id=embedding_store.user_id,
+ collection_id=embedding_store.collection_id,
+ query=query,
+ max_distance=0.3,
+ top_k=5
+ )
+
+ # Verify returned facts
+ expected_facts = [
+ "User is proficient in Python",
+ "User has experience with JavaScript",
+ "User knows SQL"
+ ]
+ assert facts == expected_facts
+
+ @pytest.mark.asyncio
+ async def test_get_relevant_facts_custom_parameters(self, embedding_store):
+ """Test fact retrieval with custom parameters."""
+ query = "What does the user do for work?"
+ top_k = 10
+ max_distance = 0.2
+
+ mock_documents = [MagicMock(content="User is a software engineer")]
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.query_documents") as mock_query_docs:
+ mock_query_docs.return_value = mock_documents
+
+ facts = await embedding_store.get_relevant_facts(
+ query, top_k=top_k, max_distance=max_distance
+ )
+
+ # Verify custom parameters were used
+ mock_query_docs.assert_called_once_with(
+ mock_db,
+ app_id=embedding_store.app_id,
+ user_id=embedding_store.user_id,
+ collection_id=embedding_store.collection_id,
+ query=query,
+ max_distance=max_distance,
+ top_k=top_k
+ )
+
+ @pytest.mark.asyncio
+ async def test_get_relevant_facts_no_results(self, embedding_store):
+ """Test fact retrieval when no relevant facts are found."""
+ query = "What is the user's favorite food?"
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.query_documents") as mock_query_docs:
+ mock_query_docs.return_value = [] # No results
+
+ facts = await embedding_store.get_relevant_facts(query)
+
+ assert facts == []
+
+ @pytest.mark.asyncio
+ async def test_get_relevant_facts_empty_query(self, embedding_store):
+ """Test fact retrieval with empty query."""
+ query = ""
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.query_documents") as mock_query_docs:
+ mock_query_docs.return_value = []
+
+ facts = await embedding_store.get_relevant_facts(query)
+
+ # Should still call query_documents with empty query
+ mock_query_docs.assert_called_once()
+ assert facts == []
+
+
+class TestRemoveDuplicates:
+ """Test duplicate detection and removal functionality."""
+
+ @pytest_asyncio.fixture
+ async def embedding_store(self, sample_data):
+ """Create an embedding store for testing."""
+ test_app, test_user = sample_data
+ collection_id = str(uuid4())
+ return CollectionEmbeddingStore(test_app.public_id, test_user.public_id, collection_id)
+
+ @pytest.mark.asyncio
+ async def test_remove_duplicates_no_duplicates(self, embedding_store):
+ """Test duplicate removal when no duplicates exist."""
+ facts = [
+ "User is a Python developer",
+ "User works remotely",
+ "User loves coffee"
+ ]
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.get_duplicate_documents") as mock_get_dupes:
+ mock_get_dupes.return_value = [] # No duplicates
+
+ unique_facts = await embedding_store.remove_duplicates(facts)
+
+ # All facts should be considered unique
+ assert unique_facts == facts
+
+ # Should check each fact for duplicates
+ assert mock_get_dupes.call_count == 3
+
+ @pytest.mark.asyncio
+ async def test_remove_duplicates_with_duplicates(self, embedding_store):
+ """Test duplicate removal when duplicates exist."""
+ facts = [
+ "User is a Python developer",
+ "User codes in Python", # Similar to first fact
+ "User works remotely"
+ ]
+
+ # Mock duplicate document for second fact
+ mock_duplicate_doc = MagicMock()
+ mock_duplicate_doc.content = "User is a Python programmer"
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.get_duplicate_documents") as mock_get_dupes:
+ def mock_get_dupes_side_effect(db, app_id, user_id, collection_id, content, similarity_threshold):
+ if "codes in Python" in content:
+ return [mock_duplicate_doc] # Duplicate found
+ return [] # No duplicates
+
+ mock_get_dupes.side_effect = mock_get_dupes_side_effect
+
+ unique_facts = await embedding_store.remove_duplicates(facts)
+
+ # Should remove the duplicate fact
+ expected_unique = [
+ "User is a Python developer",
+ "User works remotely"
+ ]
+ assert unique_facts == expected_unique
+
+ @pytest.mark.asyncio
+ async def test_remove_duplicates_custom_threshold(self, embedding_store):
+ """Test duplicate removal with custom similarity threshold."""
+ facts = ["User enjoys programming"]
+ similarity_threshold = 0.9
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.get_duplicate_documents") as mock_get_dupes:
+ mock_get_dupes.return_value = []
+
+ await embedding_store.remove_duplicates(facts, similarity_threshold=similarity_threshold)
+
+ # Verify custom threshold was passed
+ mock_get_dupes.assert_called_once_with(
+ mock_db,
+ app_id=embedding_store.app_id,
+ user_id=embedding_store.user_id,
+ collection_id=embedding_store.collection_id,
+ content=facts[0],
+ similarity_threshold=similarity_threshold
+ )
+
+ @pytest.mark.asyncio
+ async def test_remove_duplicates_handles_errors(self, embedding_store):
+ """Test that duplicate checking handles errors gracefully."""
+ facts = [
+ "User is a Python developer",
+ "This fact will cause an error",
+ "User works remotely"
+ ]
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.get_duplicate_documents") as mock_get_dupes:
+ def mock_get_dupes_side_effect(*args, **kwargs):
+ if "cause an error" in kwargs.get("content", ""):
+ raise Exception("Database connection error")
+ return []
+
+ mock_get_dupes.side_effect = mock_get_dupes_side_effect
+
+ unique_facts = await embedding_store.remove_duplicates(facts)
+
+ # Should include all facts (error results in keeping the fact)
+ assert unique_facts == facts
+ assert mock_get_dupes.call_count == 3
+
+ @pytest.mark.asyncio
+ async def test_remove_duplicates_empty_list(self, embedding_store):
+ """Test duplicate removal with empty fact list."""
+ facts = []
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.get_duplicate_documents") as mock_get_dupes:
+ unique_facts = await embedding_store.remove_duplicates(facts)
+
+ assert unique_facts == []
+ mock_get_dupes.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_remove_duplicates_logs_duplicate_found(self, embedding_store):
+ """Test that duplicate detection logs when duplicates are found."""
+ facts = ["User loves Python programming"]
+
+ mock_duplicate_doc = MagicMock()
+ mock_duplicate_doc.content = "User enjoys Python development"
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.get_duplicate_documents") as mock_get_dupes:
+ mock_get_dupes.return_value = [mock_duplicate_doc]
+
+ with patch("src.deriver.tom.embeddings.logger.debug") as mock_log:
+ unique_facts = await embedding_store.remove_duplicates(facts)
+
+ # Should log the duplicate detection
+ mock_log.assert_called_once()
+ log_message = mock_log.call_args[0][0]
+ assert "Duplicate found" in log_message
+ assert mock_duplicate_doc.content in log_message
+ assert facts[0] in log_message
+
+ # Should not include the duplicate fact
+ assert unique_facts == []
+
+
+class TestIntegrationScenarios:
+ """Test integration scenarios combining multiple operations."""
+
+ @pytest_asyncio.fixture
+ async def embedding_store(self, sample_data):
+ """Create an embedding store for testing."""
+ test_app, test_user = sample_data
+ collection_id = str(uuid4())
+ return CollectionEmbeddingStore(test_app.public_id, test_user.public_id, collection_id)
+
+ @pytest.mark.asyncio
+ async def test_full_workflow_save_and_retrieve(self, embedding_store):
+ """Test complete workflow of saving facts and retrieving them."""
+ # First save some facts
+ facts_to_save = [
+ "User is a senior Python developer",
+ "User has 5 years of experience with FastAPI",
+ "User prefers async programming"
+ ]
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.create_document") as mock_create_doc:
+ mock_create_doc.return_value = None
+
+ await embedding_store.save_facts(facts_to_save)
+
+ # Verify all facts were saved
+ assert mock_create_doc.call_count == 3
+
+ # Then retrieve relevant facts
+ query = "What is the user's programming experience?"
+ mock_documents = [
+ MagicMock(content="User is a senior Python developer"),
+ MagicMock(content="User has 5 years of experience with FastAPI")
+ ]
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.query_documents") as mock_query_docs:
+ mock_query_docs.return_value = mock_documents
+
+ retrieved_facts = await embedding_store.get_relevant_facts(query)
+
+ expected_facts = [
+ "User is a senior Python developer",
+ "User has 5 years of experience with FastAPI"
+ ]
+ assert retrieved_facts == expected_facts
+
+ @pytest.mark.asyncio
+ async def test_duplicate_removal_before_saving(self, embedding_store):
+ """Test the typical workflow of removing duplicates before saving."""
+ facts_to_check = [
+ "User is a Python developer",
+ "User writes code in Python", # Potential duplicate
+ "User works from home"
+ ]
+
+ # Mock existing duplicate
+ mock_duplicate_doc = MagicMock()
+ mock_duplicate_doc.content = "User is proficient in Python"
+
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.get_duplicate_documents") as mock_get_dupes:
+ def mock_get_dupes_side_effect(*args, **kwargs):
+ if "writes code in Python" in kwargs.get("content", ""):
+ return [mock_duplicate_doc]
+ return []
+
+ mock_get_dupes.side_effect = mock_get_dupes_side_effect
+
+ # Remove duplicates
+ unique_facts = await embedding_store.remove_duplicates(facts_to_check)
+
+ expected_unique = [
+ "User is a Python developer",
+ "User works from home"
+ ]
+ assert unique_facts == expected_unique
+
+ # Now save the unique facts
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.create_document") as mock_create_doc:
+ mock_create_doc.return_value = None
+
+ await embedding_store.save_facts(unique_facts)
+
+ # Should only save the unique facts
+ assert mock_create_doc.call_count == 2
+
+ @pytest.mark.asyncio
+ async def test_error_recovery_in_workflow(self, embedding_store):
+ """Test error recovery across multiple operations."""
+ facts = [
+ "User is experienced with machine learning",
+ "User uses scikit-learn and pandas"
+ ]
+
+ # Test save_facts with partial failure
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.create_document") as mock_create_doc:
+ def mock_create_side_effect(*args, **kwargs):
+ if "scikit-learn" in str(kwargs.get("document", "")):
+ raise Exception("Database error")
+ return None
+
+ mock_create_doc.side_effect = mock_create_side_effect
+
+ # Should handle the error gracefully
+ await embedding_store.save_facts(facts)
+
+ # Should attempt to save both facts
+ assert mock_create_doc.call_count == 2
+
+ # Test get_relevant_facts after partial save
+ with patch("src.deriver.tom.embeddings.tracked_db") as mock_tracked_db:
+ mock_db = AsyncMock()
+ mock_tracked_db.return_value.__aenter__.return_value = mock_db
+ mock_tracked_db.return_value.__aexit__.return_value = None
+
+ with patch("src.deriver.tom.embeddings.crud.query_documents") as mock_query_docs:
+ # Only return the successfully saved fact
+ mock_query_docs.return_value = [
+ MagicMock(content="User is experienced with machine learning")
+ ]
+
+ retrieved_facts = await embedding_store.get_relevant_facts("machine learning")
+
+ assert retrieved_facts == ["User is experienced with machine learning"]
\ No newline at end of file
diff --git a/tests/deriver/test_tom_modules.py b/tests/deriver/test_tom_modules.py
new file mode 100644
index 00000000..f07fe156
--- /dev/null
+++ b/tests/deriver/test_tom_modules.py
@@ -0,0 +1,701 @@
+"""Tests for TOM (Theory of Mind) inference modules."""
+
+import json
+from unittest.mock import AsyncMock, MagicMock, patch
+from uuid import uuid4
+
+import pytest
+import pytest_asyncio
+
+from src.deriver.tom import (
+ get_tom_inference,
+ get_user_representation
+)
+from src.deriver.tom.single_prompt import (
+ get_tom_inference_single_prompt,
+ get_user_representation_single_prompt
+)
+from src.deriver.tom.conversational import (
+ get_tom_inference_conversational,
+ get_user_representation_conversational
+)
+from src.deriver.tom.long_term import (
+ get_user_representation_long_term,
+ extract_facts_long_term
+)
+
+
+class TestTOMRouter:
+ """Test the main TOM routing functions in __init__.py."""
+
+ @pytest.mark.asyncio
+ async def test_get_tom_inference_routes_to_conversational(self):
+ """Test routing to conversational TOM inference method."""
+ chat_history = "User: I'm a Python developer\nAI: How long have you been coding?"
+ session_id = str(uuid4())
+ user_representation = "User is technical"
+
+ with patch("src.deriver.tom.get_tom_inference_conversational") as mock_conversational:
+ mock_conversational.return_value = "Conversational TOM response"
+
+ result = await get_tom_inference(
+ chat_history, session_id, user_representation, method="conversational"
+ )
+
+ mock_conversational.assert_called_once_with(
+ chat_history, session_id, user_representation
+ )
+ assert result == "Conversational TOM response"
+
+ @pytest.mark.asyncio
+ async def test_get_tom_inference_routes_to_single_prompt(self):
+ """Test routing to single prompt TOM inference method."""
+ chat_history = "User: I love machine learning\nAI: What frameworks do you use?"
+ session_id = str(uuid4())
+
+ with patch("src.deriver.tom.get_tom_inference_single_prompt") as mock_single_prompt:
+ mock_single_prompt.return_value = "Single prompt TOM response"
+
+ result = await get_tom_inference(
+ chat_history, session_id, method="single_prompt"
+ )
+
+ mock_single_prompt.assert_called_once_with(
+ chat_history, session_id, "None"
+ )
+ assert result == "Single prompt TOM response"
+
+ @pytest.mark.asyncio
+ async def test_get_tom_inference_invalid_method_raises_error(self):
+ """Test that invalid TOM inference method raises ValueError."""
+ with pytest.raises(ValueError, match="Invalid method: invalid_method"):
+ await get_tom_inference(
+ "chat history", "session_id", method="invalid_method"
+ )
+
+ @pytest.mark.asyncio
+ async def test_get_user_representation_routes_to_conversational(self):
+ """Test routing to conversational user representation method."""
+ chat_history = "User: I work in AI research"
+ session_id = str(uuid4())
+ tom_inference = "User is excited about AI"
+
+ with patch("src.deriver.tom.get_user_representation_conversational") as mock_conversational:
+ mock_conversational.return_value = "Conversational representation"
+
+ result = await get_user_representation(
+ chat_history, session_id, tom_inference=tom_inference, method="conversational"
+ )
+
+ mock_conversational.assert_called_once_with(
+ chat_history, session_id, "None", tom_inference
+ )
+ assert result == "Conversational representation"
+
+ @pytest.mark.asyncio
+ async def test_get_user_representation_routes_to_long_term(self):
+ """Test routing to long term user representation method."""
+ chat_history = "User: I've been programming for 5 years"
+ session_id = str(uuid4())
+
+ with patch("src.deriver.tom.get_user_representation_long_term") as mock_long_term:
+ mock_long_term.return_value = "Long term representation"
+
+ result = await get_user_representation(
+ chat_history, session_id, method="long_term"
+ )
+
+ mock_long_term.assert_called_once_with(
+ chat_history, session_id, "None", "None"
+ )
+ assert result == "Long term representation"
+
+ @pytest.mark.asyncio
+ async def test_get_user_representation_invalid_method_raises_error(self):
+ """Test that invalid user representation method raises ValueError."""
+ with pytest.raises(ValueError, match="Invalid method: unknown_method"):
+ await get_user_representation(
+ "chat history", "session_id", method="unknown_method"
+ )
+
+ @pytest.mark.asyncio
+ async def test_tom_inference_with_kwargs(self):
+ """Test that kwargs are properly passed through to TOM methods."""
+ chat_history = "User: Test message"
+ session_id = str(uuid4())
+ extra_param = "test_value"
+
+ with patch("src.deriver.tom.get_tom_inference_single_prompt") as mock_single_prompt:
+ mock_single_prompt.return_value = "Response with kwargs"
+
+ await get_tom_inference(
+ chat_history, session_id, method="single_prompt", extra_param=extra_param
+ )
+
+ # Verify kwargs were passed through
+ mock_single_prompt.assert_called_once_with(
+ chat_history, session_id, "None", extra_param=extra_param
+ )
+
+
+class TestSinglePromptMethods:
+ """Test the single prompt TOM inference methods."""
+
+ @pytest.mark.asyncio
+ async def test_get_tom_inference_single_prompt_basic(self, mock_model_clients):
+ """Test basic single prompt TOM inference."""
+ chat_history = "User: I'm feeling stressed about work\nAI: What's causing the stress?"
+ session_id = str(uuid4())
+
+ with patch("src.deriver.tom.single_prompt.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ result = await get_tom_inference_single_prompt(chat_history, session_id)
+
+ # Verify model client was called with correct parameters
+ mock_client = mock_model_clients["single_prompt"]
+ mock_client.generate.assert_called_once()
+ call_kwargs = mock_client.generate.call_args[1]
+
+ assert call_kwargs["max_tokens"] == 1000
+ assert call_kwargs["temperature"] == 0
+ assert call_kwargs["use_caching"] is True
+ # The system prompt should contain key TOM instruction phrases
+ assert "system" in call_kwargs
+ system_prompt = call_kwargs["system"]
+ assert "theory of mind" in system_prompt.lower() or "prediction" in system_prompt.lower()
+
+ @pytest.mark.asyncio
+ async def test_get_tom_inference_single_prompt_with_user_representation(self, mock_model_clients):
+ """Test single prompt TOM inference with existing user representation."""
+ chat_history = "User: I changed my mind about the project"
+ session_id = str(uuid4())
+ user_representation = "User is decisive and goal-oriented"
+
+ with patch("src.deriver.tom.single_prompt.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ await get_tom_inference_single_prompt(
+ chat_history, session_id, user_representation
+ )
+
+ # Verify user representation was included in messages
+ mock_client = mock_model_clients["single_prompt"]
+ call_args = mock_client.generate.call_args[1]
+ messages = call_args["messages"]
+
+ # Should have two messages: main analysis + user representation context
+ assert len(messages) == 2
+ assert user_representation in str(messages)
+
+ @pytest.mark.asyncio
+ async def test_get_tom_inference_single_prompt_handles_error(self, mock_model_clients):
+ """Test that single prompt TOM inference handles LLM errors."""
+ chat_history = "User: Test message"
+ session_id = str(uuid4())
+
+ # Mock the model client to raise an exception
+ mock_model_clients["single_prompt"].generate.side_effect = Exception("LLM API Error")
+
+ with patch("src.deriver.tom.single_prompt.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ with patch("src.deriver.tom.single_prompt.sentry_sdk.capture_exception") as mock_capture:
+ with pytest.raises(Exception, match="LLM API Error"):
+ await get_tom_inference_single_prompt(chat_history, session_id)
+
+ # Verify error was captured
+ mock_capture.assert_called_once()
+
+ @pytest.mark.asyncio
+ async def test_get_user_representation_single_prompt_basic(self, mock_model_clients):
+ """Test basic single prompt user representation."""
+ chat_history = "User: I'm a data scientist\nAI: What tools do you use?"
+ session_id = str(uuid4())
+ tom_inference = "User is passionate about data science"
+
+ with patch("src.deriver.tom.single_prompt.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ result = await get_user_representation_single_prompt(
+ chat_history, session_id, tom_inference=tom_inference
+ )
+
+ # Verify correct system prompt was used
+ mock_client = mock_model_clients["single_prompt"]
+ call_kwargs = mock_client.generate.call_args[1]
+ # The system prompt should contain user representation instructions
+ assert "system" in call_kwargs
+ system_prompt = call_kwargs["system"]
+ assert "user representation" in system_prompt.lower() or "factual" in system_prompt.lower()
+
+ # Verify TOM inference was included in context
+ messages = call_kwargs["messages"]
+ assert tom_inference in str(messages)
+
+ @pytest.mark.asyncio
+ async def test_get_user_representation_single_prompt_all_inputs(self, mock_model_clients):
+ """Test single prompt user representation with all optional inputs."""
+ chat_history = "User: I've been learning React lately"
+ session_id = str(uuid4())
+ user_representation = "User is a full-stack developer"
+ tom_inference = "User is eager to learn new technologies"
+
+ with patch("src.deriver.tom.single_prompt.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ await get_user_representation_single_prompt(
+ chat_history, session_id, user_representation, tom_inference
+ )
+
+ # Verify all inputs were included in the context
+ mock_client = mock_model_clients["single_prompt"]
+ call_kwargs = mock_client.generate.call_args[1]
+ messages = call_kwargs["messages"]
+
+ message_content = str(messages)
+ assert chat_history in message_content
+ assert user_representation in message_content
+ assert tom_inference in message_content
+
+
+class TestConversationalMethods:
+ """Test the conversational TOM inference methods."""
+
+ @pytest.mark.asyncio
+ async def test_get_tom_inference_conversational_basic(self):
+ """Test basic conversational TOM inference."""
+ chat_history = "User: I'm learning to cook\nAI: That's exciting! What dishes interest you?"
+ session_id = str(uuid4())
+ user_representation = "User enjoys trying new things"
+
+ # Mock the Anthropic client
+ mock_message = MagicMock()
+ mock_message.content = [MagicMock(text="User is enthusiastic about cooking")]
+
+ with patch("src.deriver.tom.conversational.anthropic.messages.create") as mock_create:
+ mock_create.return_value = mock_message
+
+ with patch("src.deriver.tom.conversational.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ result = await get_tom_inference_conversational(
+ chat_history, session_id, user_representation
+ )
+
+ # Verify Anthropic client was called
+ mock_create.assert_called_once()
+ call_kwargs = mock_create.call_args[1]
+
+ assert call_kwargs["model"] == "claude-3-5-sonnet-20240620"
+ assert call_kwargs["max_tokens"] == 1000
+ assert call_kwargs["temperature"] == 0
+
+ # Verify chat history and user representation were included
+ messages = call_kwargs["messages"]
+ message_content = str(messages)
+ # Check for key parts of the chat history and user representation
+ assert "learning to cook" in message_content.lower()
+ assert "enjoys trying new things" in message_content.lower() or user_representation in message_content
+
+ assert result == "User is enthusiastic about cooking"
+
+ @pytest.mark.asyncio
+ async def test_get_tom_inference_conversational_complex_prompting(self):
+ """Test that conversational method uses complex metanarrative prompting."""
+ chat_history = "User: I'm having trouble with my team\nAI: What kind of challenges are you facing?"
+ session_id = str(uuid4())
+
+ mock_message = MagicMock()
+ mock_message.content = [MagicMock(text="User seems frustrated with team dynamics")]
+
+ with patch("src.deriver.tom.conversational.anthropic.messages.create") as mock_create:
+ mock_create.return_value = mock_message
+
+ with patch("src.deriver.tom.conversational.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ await get_tom_inference_conversational(chat_history, session_id)
+
+ # Verify complex prompting structure
+ call_kwargs = mock_create.call_args[1]
+ messages = call_kwargs["messages"]
+
+ # Should have multiple role-playing messages
+ assert len(messages) >= 5
+
+ # Verify OOC (out of character) setup is included
+ message_content = str(messages)
+ assert "OOC" in message_content
+ assert "experiment" in message_content.lower()
+
+ @pytest.mark.asyncio
+ async def test_get_user_representation_conversational_basic(self):
+ """Test basic conversational user representation."""
+ chat_history = "User: I work in finance but I'm passionate about art"
+ session_id = str(uuid4())
+ tom_inference = "User has diverse interests spanning analytical and creative domains"
+
+ mock_message = MagicMock()
+ mock_message.content = [MagicMock(text="User balances analytical work with creative pursuits")]
+
+ with patch("src.deriver.tom.conversational.anthropic.messages.create") as mock_create:
+ mock_create.return_value = mock_message
+
+ with patch("src.deriver.tom.conversational.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ result = await get_user_representation_conversational(
+ chat_history, session_id, tom_inference=tom_inference
+ )
+
+ # Verify TOM inference was included in the prompt
+ call_kwargs = mock_create.call_args[1]
+ messages = call_kwargs["messages"]
+ assert tom_inference in str(messages)
+
+ assert result == "User balances analytical work with creative pursuits"
+
+ @pytest.mark.asyncio
+ async def test_get_user_representation_conversational_with_existing_representation(self):
+ """Test conversational user representation with existing representation."""
+ chat_history = "User: I've started learning piano"
+ session_id = str(uuid4())
+ user_representation = "User enjoys creative hobbies"
+ tom_inference = "User is expanding creative skills"
+
+ mock_message = MagicMock()
+ mock_message.content = [MagicMock(text="Updated representation with piano learning")]
+
+ with patch("src.deriver.tom.conversational.anthropic.messages.create") as mock_create:
+ mock_create.return_value = mock_message
+
+ with patch("src.deriver.tom.conversational.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ await get_user_representation_conversational(
+ chat_history, session_id, user_representation, tom_inference
+ )
+
+ # Verify all inputs were included
+ call_kwargs = mock_create.call_args[1]
+ messages = call_kwargs["messages"]
+ message_content = str(messages)
+
+ assert chat_history in message_content
+ assert user_representation in message_content
+ assert tom_inference in message_content
+
+
+class TestLongTermMethods:
+ """Test the long term TOM methods."""
+
+ @pytest.mark.asyncio
+ async def test_extract_facts_long_term_basic(self, mock_model_clients, mock_llm_responses):
+ """Test basic fact extraction from chat history."""
+ chat_history = "User: I'm a software engineer at Google and I love hiking on weekends"
+
+ # Mock the response with proper XML format
+ mock_response = f'{mock_llm_responses["fact_extraction"]}'
+ mock_model_clients["long_term"].generate.return_value = mock_response
+
+ with patch("src.deriver.tom.long_term.parse_xml_content") as mock_parse_xml:
+ mock_parse_xml.return_value = mock_llm_responses["fact_extraction"]
+
+ facts = await extract_facts_long_term(chat_history)
+
+ # Verify model client was called
+ mock_client = mock_model_clients["long_term"]
+ mock_client.generate.assert_called_once()
+
+ call_kwargs = mock_client.generate.call_args[1]
+ assert call_kwargs["temperature"] == 0.0
+ assert call_kwargs["use_caching"] is True
+
+ # Verify the system prompt includes the chat history
+ messages = call_kwargs["messages"]
+ message_content = str(messages)
+ # Check for key parts of the chat history
+ assert "software engineer" in message_content.lower()
+ assert "google" in message_content.lower() or "hiking" in message_content.lower()
+
+ # Verify facts were extracted correctly
+ expected_facts = [
+ "User is a software developer",
+ "User works remotely",
+ "User prefers coffee over tea",
+ "User uses Python and JavaScript"
+ ]
+ assert facts == expected_facts
+
+ @pytest.mark.asyncio
+ async def test_extract_facts_long_term_handles_json_error(self, mock_model_clients):
+ """Test that fact extraction handles JSON parsing errors gracefully."""
+ chat_history = "User: I like programming"
+
+ # Mock malformed response
+ mock_model_clients["long_term"].generate.return_value = "Invalid JSON response"
+
+ with patch("src.deriver.tom.long_term.parse_xml_content") as mock_parse_xml:
+ mock_parse_xml.return_value = "Not valid JSON"
+
+ facts = await extract_facts_long_term(chat_history)
+
+ # Should return empty list on error
+ assert facts == []
+
+ @pytest.mark.asyncio
+ async def test_extract_facts_long_term_handles_missing_facts_key(self, mock_model_clients):
+ """Test that fact extraction handles missing 'facts' key in response."""
+ chat_history = "User: Test message"
+
+ # Mock response with missing facts key
+ invalid_response = json.dumps({"other_key": "some_value"})
+ mock_model_clients["long_term"].generate.return_value = f"{invalid_response}"
+
+ with patch("src.deriver.tom.long_term.parse_xml_content") as mock_parse_xml:
+ mock_parse_xml.return_value = invalid_response
+
+ facts = await extract_facts_long_term(chat_history)
+
+ # Should return empty list on KeyError
+ assert facts == []
+
+ @pytest.mark.asyncio
+ async def test_get_user_representation_long_term_basic(self, mock_model_clients):
+ """Test basic long term user representation."""
+ chat_history = "User: I'm starting a new job next week"
+ session_id = str(uuid4())
+ facts = ["User is a software engineer", "User is changing jobs"]
+
+ mock_response = "CURRENT STATE:\n- Active Context: Starting new job\n\nTENTATIVE PATTERNS:\n- High confidence: Career-focused"
+ mock_model_clients["long_term"].generate.return_value = mock_response
+
+ result = await get_user_representation_long_term(
+ chat_history, session_id, facts=facts
+ )
+
+ # Verify model client was called
+ mock_client = mock_model_clients["long_term"]
+ mock_client.generate.assert_called_once()
+
+ call_kwargs = mock_client.generate.call_args[1]
+ assert call_kwargs["temperature"] == 0
+ assert call_kwargs["use_caching"] is True
+
+ # Verify chat history was included
+ messages = call_kwargs["messages"]
+ assert chat_history in str(messages)
+
+ # Verify facts were injected into the response
+ assert "User is a software engineer" in result
+ assert "User is changing jobs" in result
+ assert "" not in result # Should be replaced
+
+ @pytest.mark.asyncio
+ async def test_get_user_representation_long_term_with_all_inputs(self, mock_model_clients):
+ """Test long term user representation with all optional inputs."""
+ chat_history = "User: I'm excited about the new project"
+ session_id = str(uuid4())
+ user_representation = "User is enthusiastic about work"
+ tom_inference = "User is feeling motivated"
+ facts = ["User works in tech", "User enjoys new challenges"]
+
+ mock_response = "Updated representation with placeholder"
+ mock_model_clients["long_term"].generate.return_value = mock_response
+
+ result = await get_user_representation_long_term(
+ chat_history, session_id, user_representation, tom_inference, facts
+ )
+
+ # Verify all inputs were included in the context
+ mock_client = mock_model_clients["long_term"]
+ call_kwargs = mock_client.generate.call_args[1]
+ messages = call_kwargs["messages"]
+ message_content = str(messages)
+
+ assert chat_history in message_content
+ assert user_representation in message_content
+ assert tom_inference in message_content
+
+ # Verify facts injection worked
+ assert "User works in tech" in result
+ assert "User enjoys new challenges" in result
+
+ @pytest.mark.asyncio
+ async def test_get_user_representation_long_term_empty_facts(self, mock_model_clients):
+ """Test long term user representation with empty facts list."""
+ chat_history = "User: Hello there"
+ session_id = str(uuid4())
+
+ mock_response = "Basic representation with placeholder"
+ mock_model_clients["long_term"].generate.return_value = mock_response
+
+ result = await get_user_representation_long_term(chat_history, session_id)
+
+ # Verify empty facts are handled gracefully
+ assert "PERSISTENT INFORMATION:\n" in result
+ assert "" not in result
+
+ @pytest.mark.asyncio
+ async def test_get_user_representation_long_term_none_inputs(self, mock_model_clients):
+ """Test long term user representation with None inputs."""
+ chat_history = "User: Test message"
+ session_id = str(uuid4())
+
+ mock_response = "Representation with "
+ mock_model_clients["long_term"].generate.return_value = mock_response
+
+ result = await get_user_representation_long_term(
+ chat_history, session_id, user_representation="None", tom_inference="None"
+ )
+
+ # Verify None inputs are handled (not included in context)
+ mock_client = mock_model_clients["long_term"]
+ call_kwargs = mock_client.generate.call_args[1]
+ messages = call_kwargs["messages"]
+ message_content = str(messages)
+
+ # "None" values should not be included in context strings
+ assert "EXISTING USER REPRESENTATION - INCOMPLETE" not in message_content
+ assert "PREDICTION OF USER MENTAL STATE" not in message_content
+
+
+class TestTOMIntegration:
+ """Test integration scenarios across TOM methods."""
+
+ @pytest.mark.asyncio
+ async def test_method_configuration_via_environment(self, mock_model_clients):
+ """Test that TOM methods can be configured via environment variables."""
+ chat_history = "User: I'm learning data science"
+ session_id = str(uuid4())
+
+ # Test single_prompt method
+ with patch("src.deriver.tom.single_prompt.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ await get_tom_inference(chat_history, session_id, method="single_prompt")
+
+ # Verify single prompt was called
+ mock_model_clients["single_prompt"].generate.assert_called()
+
+ @pytest.mark.asyncio
+ async def test_error_handling_across_methods(self, mock_model_clients):
+ """Test error handling consistency across different TOM methods."""
+ chat_history = "User: Error test"
+ session_id = str(uuid4())
+
+ # Test single_prompt error handling
+ mock_model_clients["single_prompt"].generate.side_effect = Exception("API Error")
+
+ with patch("src.deriver.tom.single_prompt.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ with patch("src.deriver.tom.single_prompt.sentry_sdk.capture_exception"):
+ with pytest.raises(Exception):
+ await get_tom_inference(chat_history, session_id, method="single_prompt")
+
+ @pytest.mark.asyncio
+ async def test_response_format_consistency(self, mock_model_clients, mock_llm_responses):
+ """Test that different methods return appropriately formatted responses."""
+ chat_history = "User: I'm a product manager"
+ session_id = str(uuid4())
+
+ # Test single prompt response format
+ with patch("src.deriver.tom.single_prompt.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ single_prompt_result = await get_tom_inference(
+ chat_history, session_id, method="single_prompt"
+ )
+
+ # Verify response is returned correctly
+ assert single_prompt_result == mock_llm_responses["tom_single_prompt"]
+
+ # Test long term fact extraction format
+ with patch("src.deriver.tom.long_term.parse_xml_content") as mock_parse_xml:
+ mock_parse_xml.return_value = mock_llm_responses["fact_extraction"]
+
+ facts = await extract_facts_long_term(chat_history)
+
+ # Verify facts are returned as list of strings
+ assert isinstance(facts, list)
+ assert all(isinstance(fact, str) for fact in facts)
+
+ @pytest.mark.asyncio
+ async def test_caching_behavior_across_methods(self, mock_model_clients):
+ """Test that caching is properly enabled across different methods."""
+ chat_history = "User: Testing caching"
+ session_id = str(uuid4())
+
+ # Test single prompt caching
+ with patch("src.deriver.tom.single_prompt.sentry_sdk.start_transaction") as mock_transaction:
+ mock_transaction.return_value.__enter__.return_value = MagicMock()
+ mock_transaction.return_value.__exit__.return_value = None
+
+ await get_tom_inference(chat_history, session_id, method="single_prompt")
+
+ call_kwargs = mock_model_clients["single_prompt"].generate.call_args[1]
+ assert call_kwargs["use_caching"] is True
+
+ # Test long term fact extraction caching
+ with patch("src.deriver.tom.long_term.parse_xml_content") as mock_parse_xml:
+ mock_parse_xml.return_value = '{"facts": []}'
+
+ await extract_facts_long_term(chat_history)
+
+ call_kwargs = mock_model_clients["long_term"].generate.call_args[1]
+ assert call_kwargs["use_caching"] is True
+
+ @pytest.mark.asyncio
+ async def test_observability_integration(self, mock_model_clients):
+ """Test that observability tools (Sentry, Langfuse) are properly integrated."""
+ chat_history = "User: Testing observability"
+ session_id = str(uuid4())
+
+ # Test single prompt observability
+ with (
+ patch("src.deriver.tom.single_prompt.sentry_sdk.start_transaction") as mock_sentry,
+ patch("src.deriver.tom.single_prompt.langfuse_context.update_current_observation") as mock_langfuse
+ ):
+ mock_sentry.return_value.__enter__.return_value = MagicMock()
+ mock_sentry.return_value.__exit__.return_value = None
+
+ await get_tom_inference(chat_history, session_id, method="single_prompt")
+
+ # Verify Sentry transaction was started
+ mock_sentry.assert_called_once_with(op="tom-inference", name="ToM Inference")
+
+ # Verify Langfuse observation was updated
+ mock_langfuse.assert_called_once()
+
+ # Test conversational method observability
+ mock_message = MagicMock()
+ mock_message.content = [MagicMock(text="Test response")]
+
+ with (
+ patch("src.deriver.tom.conversational.anthropic.messages.create") as mock_create,
+ patch("src.deriver.tom.conversational.sentry_sdk.start_transaction") as mock_sentry,
+ patch("src.deriver.tom.conversational.langfuse_context.update_current_observation") as mock_langfuse
+ ):
+ mock_create.return_value = mock_message
+ mock_sentry.return_value.__enter__.return_value = MagicMock()
+ mock_sentry.return_value.__exit__.return_value = None
+
+ await get_tom_inference(chat_history, session_id, method="conversational")
+
+ # Verify observability integration
+ mock_sentry.assert_called_once()
+ mock_langfuse.assert_called_once()
\ No newline at end of file