RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update documentation and improve code consistency across various modules 

- Added missing newlines and improved formatting in `CLAUDE.md` for better readability.
- Updated command instructions in the development guide to use `uv run` for running the server and tests.
- Enhanced type hinting in `consumer.py` for better clarity on the `facts` variable.
- Corrected a typo in the docstring of the `chat` function in `agent.py`.
- Refined type checks in `sessions.py` to use `isinstance` for better practice.
- Improved message formatting in `history.py` to ensure consistent output.
- Cleaned up test files by removing unnecessary imports and improving readability.

These changes enhance the overall clarity and maintainability of the codebase.
This commit is contained in:
hyusap 2025-06-06 18:16:50 -04:00
parent bc72671074
commit 280f4215a3
10 changed files with 631 additions and 436 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
CLAUDE.md
View File

@ -1,6 +1,7 @@
# Honcho Overview
## What is Honcho?
Honcho is an infrastructure layer for building AI agents with social cognition and theory of mind capabilities. Its primary purposes include:
- Imbuing agents with a sense of identity
@ -13,14 +14,16 @@ Honcho leverages the inherent theory-of-mind capabilities of LLMs to build coher
## Development Guide
### Commands
- Setup: `uv sync`
- Run server: `fastapi dev src/main.py`
- Run tests: `pytest tests/`
- Run single test: `pytest tests/path/to/test_file.py::test_function`
- Run server: `uv run fastapi dev src/main.py`
- Run tests: `uv run pytest tests/`
- Run single test: `uv run pytest tests/path/to/test_file.py::test_function`
- Linting: `ruff check src/`
- Format code: `ruff format src/`
### Code Style
- Follow isort conventions with absolute imports preferred
- Use explicit type hints with SQLAlchemy mapped_column annotations
- snake_case for variables/functions; PascalCase for classes
@ -29,6 +32,7 @@ Honcho leverages the inherent theory-of-mind capabilities of LLMs to build coher
- Docstrings: Use Google style docstrings
### Project Structure
- FastAPI routes in src/routers/
- SQLAlchemy ORM models in src/models.py with proper type annotations
- Pydantic schemas in src/schemas.py for API validation
@ -36,8 +40,9 @@ Honcho leverages the inherent theory-of-mind capabilities of LLMs to build coher
- Use environment variables via python-dotenv (.env)
### Error Handling
- Custom exceptions defined in src/exceptions.py
- Use specific exception types (ResourceNotFoundException, ValidationException, etc.)
- Proper logging with context instead of print statements
- Global exception handlers defined in main.py
- See docs/contributing/error-handling.mdx for details
- See docs/contributing/error-handling.mdx for details

4
src/agent.py
View File

@ -6,7 +6,7 @@ from collections.abc import Iterable
from typing import Any, Optional
from inspect import cleandoc as c
import sentry_sdk
from dotenv import load_dotenv
from sentry_sdk.ai.monitoring import ai_track
from sqlalchemy import select
@ -72,7 +72,7 @@ async def chat(
stream: bool = False,
) -> llm.Stream | llm.CallResponse:
"""
Chat with the Dialectic API usingx on-demand user representation generation.
Chat with the Dialectic API using on-demand user representation generation.
This function:
1. Sets up resources needed (embedding store, latest message ID)

2
src/deriver/consumer.py
View File

@ -109,7 +109,7 @@ async def process_user_message(
logger.debug("Extracting facts from chat history")
extract_start = os.times()[4]
fact_extraction = await extract_facts_long_term(chat_history_str)
facts = fact_extraction.facts
facts: list[str] = fact_extraction.facts or []
extract_time = os.times()[4] - extract_start
console.print(f"Extracted Facts: {fact_extraction.facts}", style="bright_blue")
logger.debug(f"Extracted {len(facts)} facts in {extract_time:.2f}s")

1
src/deriver/tom/single_prompt.py
View File

@ -177,7 +177,6 @@ async def tom_inference(
4. Do not make assumptions about demographics unless explicitly stated
5. Focus on current mental state and immediate context
6. Consider your own knowledge gaps and violations of expectations (what would surprise you)
7. Always wrap your prediction in <prediction> tags.
OUTPUT FORMAT:
current_state:

2
src/routers/sessions.py
View File

@ -239,7 +239,7 @@ async def chat(
queries=options.queries,
stream=True,
)
if type(stream) is Stream:
if isinstance(stream, Stream):
async for chunk, _ in stream:
yield chunk.content
except Exception as e:

2
src/utils/history.py
View File

@ -150,7 +150,7 @@ async def create_short_summary(
Return only the summary without any explanation or meta-commentary.
<conversation>
{messages}
{format_messages(messages)}
</conversation>
<previous_summary>

365
tests/deriver/test_background_integration.py
View File

@ -1,6 +1,5 @@
"""End-to-end integration tests for the complete deriver workflow."""
import json
from unittest.mock import AsyncMock, MagicMock, patch
from uuid import uuid4
@ -9,7 +8,6 @@ import pytest_asyncio
from src import models
from src.deriver import consumer
from src.deriver.queue import QueueManager
class TestMessageToFactsWorkflow:
@ -19,12 +17,10 @@ class TestMessageToFactsWorkflow:
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={}
user_id=test_user.public_id, app_id=test_app.public_id, metadata={}
)
db_session.add(session)
await db_session.flush()
@ -34,61 +30,63 @@ class TestMessageToFactsWorkflow:
app_id=test_app.public_id,
user_id=test_user.public_id,
name=f"user_{test_user.public_id}",
metadata={"type": "user_facts"}
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"
]):
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
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):
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.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.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 - extract_facts uses global mock from conftest.py
mock_history.return_value = ("Previous conversation context", [], None)
mock_get_collection.return_value = collection
# Mock embedding store - use facts from global mock
mock_store = AsyncMock()
global_facts = ["User is a software developer", "User works remotely", "User prefers coffee over tea", "User uses Python and JavaScript"]
global_facts = [
"User is a software developer",
"User works remotely",
"User prefers coffee over tea",
"User uses Python and JavaScript",
]
mock_store.remove_duplicates.return_value = global_facts # No duplicates
mock_store.save_facts.return_value = None
mock_store_class.return_value = mock_store
@ -101,20 +99,20 @@ class TestMessageToFactsWorkflow:
test_user.public_id,
session.public_id,
message.public_id,
db_session
db_session,
)
# Verify fact extraction worked (global mock handles this)
# 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 (4 facts from global mock * 3 messages)
assert len(all_saved_facts) == 4 * 3
@ -127,34 +125,34 @@ class TestMessageToFactsWorkflow:
queue_items = []
for message in messages:
if message.is_user:
# Create queue item with the payload structure used in real system
# 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
"session_id": session.public_id,
}
queue_item = models.QueueItem(
session_id=session.id, # Use integer ID for queue
payload=payload,
processed=False
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
patch("src.deriver.consumer.summarize_if_needed") as mock_summarize,
):
mock_process_user.return_value = None
mock_process_ai.return_value = None
mock_process_ai.return_value = None
mock_summarize.return_value = None
# Process items through the consumer
@ -170,14 +168,16 @@ class TestMessageToFactsWorkflow:
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
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):
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
@ -185,27 +185,36 @@ class TestMessageToFactsWorkflow:
user_facts = [
"User name is Sarah",
"User is a data scientist",
"User works remotely"
"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.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.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']
mock_tom_inference.return_value = mock_llm_responses["tom_single_prompt"]
mock_user_rep.return_value = mock_llm_responses["tom_single_prompt"]
# Setup other mocks - extract_facts uses global mock from conftest.py
mock_history.return_value = ("Chat history", [], None)
mock_get_collection.return_value = collection
# Mock embedding store - use facts from global mock
mock_store = AsyncMock()
global_facts = ["User is a software developer", "User works remotely", "User prefers coffee over tea", "User uses Python and JavaScript"]
global_facts = [
"User is a software developer",
"User works remotely",
"User prefers coffee over tea",
"User uses Python and JavaScript",
]
mock_store.remove_duplicates.return_value = global_facts
mock_store.save_facts.return_value = None
mock_store_class.return_value = mock_store
@ -217,11 +226,13 @@ class TestMessageToFactsWorkflow:
test_user.public_id,
session.public_id,
messages[0].public_id,
db_session
db_session,
)
# Verify fact extraction and storage occurred
mock_store.save_facts.assert_called_once_with(global_facts, message_id=messages[0].public_id)
mock_store.save_facts.assert_called_once_with(
global_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
@ -234,11 +245,9 @@ class TestErrorRecoveryIntegration:
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={}
user_id=test_user.public_id, app_id=test_app.public_id, metadata={}
)
db_session.add(session)
await db_session.flush()
@ -249,7 +258,7 @@ class TestErrorRecoveryIntegration:
content="Test message for error scenarios",
metadata={},
user_id=test_user.public_id,
app_id=test_app.public_id
app_id=test_app.public_id,
)
db_session.add(message)
await db_session.flush()
@ -263,13 +272,17 @@ class TestErrorRecoveryIntegration:
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.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())
@ -283,35 +296,46 @@ class TestErrorRecoveryIntegration:
test_user.public_id,
session.public_id,
message.public_id,
db_session
db_session,
)
@pytest.mark.asyncio
async def test_partial_fact_storage_error_recovery(self, db_session, error_test_setup):
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"
"This fact will fail to store",
"User works in tech",
]
with (
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.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,
):
# extract_facts uses global mock from conftest.py
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()
global_facts = ["User is a software developer", "User works remotely", "User prefers coffee over tea", "User uses Python and JavaScript"]
global_facts = [
"User is a software developer",
"User works remotely",
"User prefers coffee over tea",
"User uses Python and JavaScript",
]
mock_store.remove_duplicates.return_value = global_facts
# save_facts method handles its own errors gracefully
mock_store.save_facts.return_value = None
@ -324,7 +348,7 @@ class TestErrorRecoveryIntegration:
test_user.public_id,
session.public_id,
message.public_id,
db_session
db_session,
)
# Verify the workflow completed
@ -338,11 +362,9 @@ class TestSummaryIntegration:
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={}
user_id=test_user.public_id, app_id=test_app.public_id, metadata={}
)
db_session.add(session)
await db_session.flush()
@ -363,30 +385,43 @@ class TestSummaryIntegration:
content=f"Message {i+1}: User discussing various topics",
metadata={},
user_id=test_user.public_id,
app_id=test_app.public_id
app_id=test_app.public_id,
)
db_session.add(message)
messages.append(message)
await db_session.flush()
with (
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.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
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 - extract_facts uses global mock from conftest.py
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()
global_facts = ["User is a software developer", "User works remotely", "User prefers coffee over tea", "User uses Python and JavaScript"]
global_facts = [
"User is a software developer",
"User works remotely",
"User prefers coffee over tea",
"User uses Python and JavaScript",
]
mock_store.remove_duplicates.return_value = global_facts
mock_store.save_facts.return_value = None
mock_store_class.return_value = mock_store
@ -397,17 +432,20 @@ class TestSummaryIntegration:
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
})
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 (using global mock)
# Verify summary generation was checked
mock_should_create.assert_called()
@ -424,18 +462,18 @@ class TestConcurrentProcessing:
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}
metadata={"session_num": i},
)
db_session.add(session)
sessions.append(session)
await db_session.flush()
# Create messages for each session
@ -447,11 +485,11 @@ class TestConcurrentProcessing:
content=f"Session {i} message: User sharing information",
metadata={},
user_id=test_user.public_id,
app_id=test_app.public_id
app_id=test_app.public_id,
)
db_session.add(message)
all_messages.append(message)
await db_session.flush()
processed_count = 0
@ -461,18 +499,23 @@ class TestConcurrentProcessing:
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_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
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 = {
@ -481,11 +524,11 @@ class TestConcurrentProcessing:
"content": message.content,
"app_id": test_app.public_id,
"user_id": test_user.public_id,
"session_id": session.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)
@ -500,12 +543,12 @@ class TestFullSystemIntegration:
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"}
metadata={"conversation_type": "onboarding"},
)
db_session.add(session)
await db_session.flush()
@ -513,12 +556,30 @@ class TestFullSystemIntegration:
# 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.")
(
"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
@ -530,18 +591,18 @@ class TestFullSystemIntegration:
content=content,
metadata={},
user_id=test_user.public_id,
app_id=test_app.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 is based in San Francisco",
"User works on backend systems",
"User uses Python and Go",
"User works at a fintech company",
@ -549,9 +610,9 @@ class TestFullSystemIntegration:
"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 Redis for caching",
"User uses PostgreSQL with read replicas",
"User is experimenting with Rust components"
"User is experimenting with Rust components",
]
# Track all extracted facts
@ -561,45 +622,67 @@ class TestFullSystemIntegration:
# Simulate realistic fact extraction based on content
facts_list = []
if "Alex" in chat_history and "software engineer" in chat_history:
facts_list = ["User name is Alex", "User is a software engineer", "User is based in San Francisco"]
facts_list = [
"User name is Alex",
"User is a software engineer",
"User is based in San Francisco",
]
elif "Python and Go" in chat_history:
facts_list = ["User works on backend systems", "User uses Python and Go", "User works at a fintech company"]
facts_list = [
"User works on backend systems",
"User uses Python and Go",
"User works at a fintech company",
]
elif "Kafka" in chat_history:
facts_list = ["User uses Kafka for streaming", "User uses Redis for caching", "User uses PostgreSQL"]
facts_list = [
"User uses Kafka for streaming",
"User uses Redis for caching",
"User uses PostgreSQL",
]
# Return object with .facts attribute like the real function
result = MagicMock()
result.facts = facts_list
return result
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.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
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.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,
@ -607,12 +690,12 @@ class TestFullSystemIntegration:
test_user.public_id,
session.public_id,
message.public_id,
db_session
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
@ -627,4 +710,4 @@ class TestFullSystemIntegration:
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")
print(f"🔧 Captured {len(captured_tech)} technical details")

181
tests/deriver/test_performance.py
View File

@ -5,7 +5,7 @@ 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
@ -18,7 +18,7 @@ class TestPerformanceValidation:
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.
@ -29,100 +29,120 @@ class TestPerformanceValidation:
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.
"""
# Extract facts and measure time - using global mocks from conftest.py
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 hasattr(facts, 'facts') and len(facts.facts) > 0
assert hasattr(facts, "facts") and len(facts.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):
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)
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 works with large-scale systems",
"User has expertise in TensorFlow and PyTorch",
"User handles millions of daily interactions",
"User uses distributed computing"
"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
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
# 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):
@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
for i in range(
performance_config["message_count"] // 10
) # Smaller load for test
]
# Mock all the dependencies for speed
with (
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.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
patch("src.deriver.consumer.summarize_if_needed") as mock_summarize,
):
# Setup fast mocks - extract_facts uses global mock
mock_history.return_value = ("", [], None)
mock_get_collection.return_value = AsyncMock()
mock_store = AsyncMock()
global_facts = ["User is a software developer", "User works remotely", "User prefers coffee over tea", "User uses Python and JavaScript"]
global_facts = [
"User is a software developer",
"User works remotely",
"User prefers coffee over tea",
"User uses Python and JavaScript",
]
mock_store.remove_duplicates.return_value = global_facts
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,
@ -130,19 +150,21 @@ class TestPerformanceValidation:
test_user.public_id,
"session_123",
f"msg_{hash(content)}",
AsyncMock() # Mock DB session
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
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")
@ -150,48 +172,57 @@ class TestPerformanceValidation:
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.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.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
patch("src.deriver.consumer.summarize_if_needed") as mock_summarize,
):
# Setup mocks - extract_facts uses global mock from conftest.py
mock_history.return_value = ("", [], None)
mock_get_collection.return_value = AsyncMock()
mock_store = AsyncMock()
global_facts = ["User is a software developer", "User works remotely", "User prefers coffee over tea", "User uses Python and JavaScript"]
global_facts = [
"User is a software developer",
"User works remotely",
"User prefers coffee over tea",
"User uses Python and JavaScript",
]
mock_store.remove_duplicates.return_value = global_facts
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_app.public_id,
test_user.public_id,
f"session_{batch}",
f"msg_{batch}_{i}",
AsyncMock()
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_store.save_facts.call_count == message_count
print(f"✅ Processed {message_count} messages in batches successfully")
@ -204,10 +235,14 @@ class TestPerformanceValidation:
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" - 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']}")
@ -218,39 +253,43 @@ class TestScalabilityValidation:
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)
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
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]
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")
print(f" - Rate: {len(large_fact_list)/total_time:.1f} facts/second")

234
tests/deriver/test_queue_manager.py
View File

@ -3,9 +3,8 @@
import asyncio
import signal
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import MagicMock, patch
from uuid import uuid4
import pytest
import pytest_asyncio
from sqlalchemy import select
@ -21,9 +20,9 @@ class TestQueueManagerInitialization:
"""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()
@ -34,9 +33,9 @@ class TestQueueManagerInitialization:
"""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
@ -44,6 +43,7 @@ class TestQueueManagerInitialization:
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"
@ -52,28 +52,29 @@ class TestQueueManagerInitialization:
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()
@ -84,13 +85,13 @@ class TestTaskAndSessionTracking:
"""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()
@ -99,11 +100,11 @@ class TestTaskAndSessionTracking:
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
@ -113,10 +114,10 @@ class TestTaskAndSessionTracking:
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
@ -128,41 +129,37 @@ class TestDatabaseOperations:
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={}
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={}
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
processed=False,
)
queue_item2 = models.QueueItem(
session_id=session2.id,
payload={"message_id": str(uuid4())},
processed=False
processed=False,
)
queue_item3 = models.QueueItem(
session_id=session1.id,
payload={"message_id": str(uuid4())},
processed=True # Already processed
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
@ -171,29 +168,31 @@ class TestDatabaseOperations:
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):
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
@ -204,19 +203,18 @@ class TestDatabaseOperations:
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
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(
@ -231,10 +229,10 @@ class TestDatabaseOperations:
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
@ -246,15 +244,15 @@ class TestDatabaseOperations:
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
@ -269,23 +267,25 @@ class TestConcurrencyControl:
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):
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
@ -294,22 +294,22 @@ class TestConcurrencyControl:
"""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()
@ -322,18 +322,18 @@ class TestSignalHandling:
"""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()
@ -342,7 +342,7 @@ class TestSignalHandling:
"""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:
@ -350,17 +350,17 @@ class TestSignalHandling:
# 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(
@ -370,20 +370,20 @@ class TestSignalHandling:
remaining_sessions = result.scalars().all()
assert len(remaining_sessions) == 0
@pytest.mark.asyncio
@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
@ -396,52 +396,59 @@ class TestErrorHandling:
"""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")
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):
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.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(
@ -452,49 +459,47 @@ class TestErrorHandling:
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):
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={}
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={}
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
processed=False,
)
queue_item2 = models.QueueItem(
session_id=session2.id,
payload={"message_id": str(uuid4())},
processed=False
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
@ -509,22 +514,22 @@ class TestIntegrationScenarios:
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(
@ -532,7 +537,7 @@ class TestIntegrationScenarios:
)
)
assert result.scalar_one_or_none() is None
# Verify session is untracked
assert session.id not in manager.owned_sessions
@ -542,32 +547,35 @@ class TestIntegrationScenarios:
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))
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)
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

263
tests/deriver/test_tom_modules.py
View File

@ -1,27 +1,22 @@
"""Tests for TOM (Theory of Mind) inference modules."""
import json
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import 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 import get_tom_inference, get_user_representation
from src.deriver.tom.single_prompt import (
get_tom_inference_single_prompt,
get_user_representation_single_prompt
get_user_representation_single_prompt,
)
from src.deriver.tom.conversational import (
get_tom_inference_conversational,
get_user_representation_conversational
get_user_representation_conversational,
)
from src.deriver.tom.long_term import (
get_user_representation_long_term,
extract_facts_long_term
extract_facts_long_term,
)
@ -31,11 +26,15 @@ class TestTOMRouter:
@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?"
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:
with patch(
"src.deriver.tom.get_tom_inference_conversational"
) as mock_conversational:
mock_conversational.return_value = "Conversational TOM response"
result = await get_tom_inference(
@ -53,16 +52,16 @@ class TestTOMRouter:
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:
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"
)
mock_single_prompt.assert_called_once_with(chat_history, session_id, "None")
assert result == "Single prompt TOM response"
@pytest.mark.asyncio
@ -80,11 +79,16 @@ class TestTOMRouter:
session_id = str(uuid4())
tom_inference = "User is excited about AI"
with patch("src.deriver.tom.get_user_representation_conversational") as mock_conversational:
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"
chat_history,
session_id,
tom_inference=tom_inference,
method="conversational",
)
mock_conversational.assert_called_once_with(
@ -98,7 +102,9 @@ class TestTOMRouter:
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:
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(
@ -125,11 +131,16 @@ class TestTOMRouter:
session_id = str(uuid4())
extra_param = "test_value"
with patch("src.deriver.tom.get_tom_inference_single_prompt") as mock_single_prompt:
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
chat_history,
session_id,
method="single_prompt",
extra_param=extra_param,
)
# Verify kwargs were passed through
@ -144,18 +155,22 @@ class TestSinglePromptMethods:
@pytest.mark.asyncio
async def test_get_tom_inference_single_prompt_basic(self, mock_llm_calls):
"""Test basic single prompt TOM inference."""
chat_history = "User: I'm feeling stressed about work\nAI: What's causing the stress?"
chat_history = (
"User: I'm feeling stressed about work\nAI: What's causing the stress?"
)
result = await get_tom_inference_single_prompt(chat_history)
# Verify the mocked function was called
mock_llm_calls["tom_inference"].assert_called_once_with(chat_history, None)
# Verify result is JSON string from mock
assert isinstance(result, str)
@pytest.mark.asyncio
async def test_get_tom_inference_single_prompt_with_user_representation(self, mock_llm_calls):
async def test_get_tom_inference_single_prompt_with_user_representation(
self, mock_llm_calls
):
"""Test single prompt TOM inference with existing user representation."""
chat_history = "User: I changed my mind about the project"
user_representation = "User is decisive and goal-oriented"
@ -165,8 +180,10 @@ class TestSinglePromptMethods:
)
# Verify the mocked function was called with correct parameters
mock_llm_calls["tom_inference"].assert_called_once_with(chat_history, user_representation)
mock_llm_calls["tom_inference"].assert_called_once_with(
chat_history, user_representation
)
# Verify result is JSON string from mock
assert isinstance(result, str)
@ -192,13 +209,17 @@ class TestSinglePromptMethods:
)
# Verify the mocked function was called with correct parameters
mock_llm_calls["user_rep_inference"].assert_called_once_with(chat_history, None, tom_inference)
mock_llm_calls["user_rep_inference"].assert_called_once_with(
chat_history, None, tom_inference
)
# Verify result is JSON string from mock
assert isinstance(result, str)
@pytest.mark.asyncio
async def test_get_user_representation_single_prompt_all_inputs(self, mock_llm_calls):
async def test_get_user_representation_single_prompt_all_inputs(
self, mock_llm_calls
):
"""Test single prompt user representation with all optional inputs."""
chat_history = "User: I've been learning React lately"
user_representation = "User is a full-stack developer"
@ -212,7 +233,7 @@ class TestSinglePromptMethods:
mock_llm_calls["user_rep_inference"].assert_called_once_with(
chat_history, user_representation, tom_inference
)
# Verify result is JSON string from mock
assert isinstance(result, str)
@ -223,18 +244,28 @@ class TestConversationalMethods:
@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?"
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="<prediction>User is enthusiastic about cooking</prediction>")]
mock_message.content = [
MagicMock(
text="<prediction>User is enthusiastic about cooking</prediction>"
)
]
with patch("src.deriver.tom.conversational.anthropic.messages.create") as mock_create:
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:
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
@ -245,19 +276,25 @@ class TestConversationalMethods:
# 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
# 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 (
"enjoys trying new things" in message_content.lower()
or user_representation in message_content
)
assert result == "<prediction>User is enthusiastic about cooking</prediction>"
assert (
result
== "<prediction>User is enthusiastic about cooking</prediction>"
)
@pytest.mark.asyncio
async def test_get_tom_inference_conversational_complex_prompting(self):
@ -266,12 +303,18 @@ class TestConversationalMethods:
session_id = str(uuid4())
mock_message = MagicMock()
mock_message.content = [MagicMock(text="User seems frustrated with team dynamics")]
mock_message.content = [
MagicMock(text="User seems frustrated with team dynamics")
]
with patch("src.deriver.tom.conversational.anthropic.messages.create") as mock_create:
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:
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
@ -280,10 +323,10 @@ class TestConversationalMethods:
# 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
@ -294,15 +337,25 @@ class TestConversationalMethods:
"""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"
tom_inference = (
"User has diverse interests spanning analytical and creative domains"
)
mock_message = MagicMock()
mock_message.content = [MagicMock(text="<representation>User balances analytical work with creative pursuits</representation>")]
mock_message.content = [
MagicMock(
text="<representation>User balances analytical work with creative pursuits</representation>"
)
]
with patch("src.deriver.tom.conversational.anthropic.messages.create") as mock_create:
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:
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
@ -315,10 +368,15 @@ class TestConversationalMethods:
messages = call_kwargs["messages"]
assert tom_inference in str(messages)
assert result == "<representation>User balances analytical work with creative pursuits</representation>"
assert (
result
== "<representation>User balances analytical work with creative pursuits</representation>"
)
@pytest.mark.asyncio
async def test_get_user_representation_conversational_with_existing_representation(self):
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())
@ -326,12 +384,18 @@ class TestConversationalMethods:
tom_inference = "User is expanding creative skills"
mock_message = MagicMock()
mock_message.content = [MagicMock(text="Updated representation with piano learning")]
mock_message.content = [
MagicMock(text="Updated representation with piano learning")
]
with patch("src.deriver.tom.conversational.anthropic.messages.create") as mock_create:
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:
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
@ -343,7 +407,7 @@ class TestConversationalMethods:
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
@ -353,84 +417,76 @@ class TestLongTermMethods:
"""Test the long term TOM methods."""
@pytest.mark.asyncio
async def test_extract_facts_long_term_basic(self, mock_llm_calls, mock_llm_responses):
async def test_extract_facts_long_term_basic(
self, mock_llm_calls, 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"
# Use the global mock directly since the decorated function can't be called in tests
result = mock_llm_calls["extract_facts"].return_value
# Verify result has facts attribute from mock (configured in conftest.py)
assert hasattr(result, 'facts')
assert hasattr(result, "facts")
assert isinstance(result.facts, list)
@pytest.mark.asyncio
async def test_extract_facts_long_term_handles_json_error(self, mock_llm_calls):
"""Test that fact extraction handles JSON parsing errors gracefully."""
chat_history = "User: I like programming"
# Use the global mock directly since the decorated function can't be called in tests
result = mock_llm_calls["extract_facts"].return_value
# Should return result from mock
assert hasattr(result, 'facts')
assert hasattr(result, "facts")
@pytest.mark.asyncio
async def test_extract_facts_long_term_handles_missing_facts_key(self, mock_llm_calls):
async def test_extract_facts_long_term_handles_missing_facts_key(
self, mock_llm_calls
):
"""Test that fact extraction handles missing 'facts' key in response."""
chat_history = "User: Test message"
# Use the global mock directly since the decorated function can't be called in tests
result = mock_llm_calls["extract_facts"].return_value
# Should return result from mock
assert hasattr(result, 'facts')
assert hasattr(result, "facts")
@pytest.mark.asyncio
async def test_get_user_representation_long_term_basic(self, mock_llm_calls):
"""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"]
# Use the global mock directly since the decorated function can't be called in tests
result = mock_llm_calls["long_term_user_rep"].return_value
# Verify result is the mock object (since this function returns object directly)
assert hasattr(result, 'current_state')
assert hasattr(result, 'tentative_patterns')
assert hasattr(result, "current_state")
assert hasattr(result, "tentative_patterns")
@pytest.mark.asyncio
async def test_get_user_representation_long_term_with_all_inputs(self, mock_llm_calls):
async def test_get_user_representation_long_term_with_all_inputs(
self, mock_llm_calls
):
"""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"]
# Use the global mock directly since the decorated function can't be called in tests
result = mock_llm_calls["long_term_user_rep"].return_value
# Verify result has the expected structure from the mock
assert hasattr(result, 'current_state')
assert hasattr(result, 'tentative_patterns')
assert hasattr(result, 'knowledge_gaps')
assert hasattr(result, 'expectation_violations')
assert hasattr(result, 'updates')
assert hasattr(result, "current_state")
assert hasattr(result, "tentative_patterns")
assert hasattr(result, "knowledge_gaps")
assert hasattr(result, "expectation_violations")
assert hasattr(result, "updates")
@pytest.mark.asyncio
async def test_get_user_representation_long_term_empty_facts(self, mock_llm_calls):
"""Test long term user representation with empty facts list."""
chat_history = "User: Hello there"
session_id = str(uuid4())
# Use the global mock directly since the decorated function can't be called in tests
result = mock_llm_calls["long_term_user_rep"].return_value
# Verify result has the expected structure from the mock
assert hasattr(result, 'current_state')
assert hasattr(result, 'tentative_patterns')
assert hasattr(result, "current_state")
assert hasattr(result, "tentative_patterns")
@pytest.mark.asyncio
async def test_get_user_representation_long_term_none_inputs(self, mock_llm_calls):
@ -442,9 +498,9 @@ class TestLongTermMethods:
result = mock_llm_calls["long_term_user_rep"].return_value
# Verify result has the expected structure from the mock
assert hasattr(result, 'current_state')
assert hasattr(result, 'tentative_patterns')
assert hasattr(result, 'knowledge_gaps')
assert hasattr(result, "current_state")
assert hasattr(result, "tentative_patterns")
assert hasattr(result, "knowledge_gaps")
class TestTOMIntegration:
@ -459,7 +515,7 @@ class TestTOMIntegration:
# Test that single_prompt method mock is available
assert "tom_inference" in mock_llm_calls
result = mock_llm_calls["tom_inference"].return_value
assert hasattr(result, 'model_dump_json')
assert hasattr(result, "model_dump_json")
@pytest.mark.asyncio
async def test_error_handling_across_methods(self, mock_llm_calls):
@ -469,25 +525,30 @@ class TestTOMIntegration:
# Test that error handling can be simulated via mocks
mock_llm_calls["tom_inference"].side_effect = Exception("API Error")
# Verify the mock can raise exceptions
with pytest.raises(Exception, match="API Error"):
await mock_llm_calls["tom_inference"]()
@pytest.mark.asyncio
async def test_response_format_consistency(self, mock_llm_calls, mock_llm_responses):
async def test_response_format_consistency(
self, mock_llm_calls, 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 via mock
single_prompt_result = mock_llm_calls["tom_inference"].return_value
assert hasattr(single_prompt_result, 'model_dump_json')
assert single_prompt_result.model_dump_json() == mock_llm_responses["tom_single_prompt"]
assert hasattr(single_prompt_result, "model_dump_json")
assert (
single_prompt_result.model_dump_json()
== mock_llm_responses["tom_single_prompt"]
)
# Test long term fact extraction format via mock
facts = mock_llm_calls["extract_facts"].return_value
assert hasattr(facts, 'facts')
assert hasattr(facts, "facts")
assert isinstance(facts.facts, list)
assert all(isinstance(fact, str) for fact in facts.facts)
@ -500,7 +561,7 @@ class TestTOMIntegration:
# Test that mocks are available for caching tests
assert "tom_inference" in mock_llm_calls
assert "extract_facts" in mock_llm_calls
# Verify mocks can be configured for caching behavior
mock_llm_calls["tom_inference"].assert_not_called()
mock_llm_calls["extract_facts"].assert_not_called()
@ -514,7 +575,7 @@ class TestTOMIntegration:
# Test that all required mocks are available for observability testing
assert "tom_inference" in mock_llm_calls
assert "anthropic" in mock_llm_calls
# Verify mocks are properly configured
assert mock_llm_calls["tom_inference"] is not None
assert mock_llm_calls["anthropic"] is not None
assert mock_llm_calls["anthropic"] is not None