feat(context): implement main ContextEngine with full integration (#85)
Phase 7 of Context Management Engine - Main Engine: - Add ContextEngine as main orchestration class - Integrate all components: calculator, scorer, ranker, compressor, cache - Add high-level assemble_context() API with: - System prompt support - Task description support - Knowledge Base integration via MCP - Conversation history conversion - Tool results conversion - Custom contexts support - Add helper methods: - get_budget_for_model() - count_tokens() with caching - invalidate_cache() - get_stats() - Add create_context_engine() factory function Tests: 26 new tests, 311 total context tests passing 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
458
backend/tests/services/context/test_engine.py
Normal file
458
backend/tests/services/context/test_engine.py
Normal file
@@ -0,0 +1,458 @@
|
||||
"""Tests for ContextEngine."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from app.services.context.config import ContextSettings
|
||||
from app.services.context.engine import ContextEngine, create_context_engine
|
||||
from app.services.context.types import (
|
||||
AssembledContext,
|
||||
ConversationContext,
|
||||
KnowledgeContext,
|
||||
MessageRole,
|
||||
SystemContext,
|
||||
TaskContext,
|
||||
ToolContext,
|
||||
)
|
||||
|
||||
|
||||
class TestContextEngineCreation:
|
||||
"""Tests for ContextEngine creation."""
|
||||
|
||||
def test_creation_minimal(self) -> None:
|
||||
"""Test creating engine with minimal config."""
|
||||
engine = ContextEngine()
|
||||
|
||||
assert engine._mcp is None
|
||||
assert engine._settings is not None
|
||||
assert engine._calculator is not None
|
||||
assert engine._scorer is not None
|
||||
assert engine._ranker is not None
|
||||
assert engine._compressor is not None
|
||||
assert engine._cache is not None
|
||||
assert engine._pipeline is not None
|
||||
|
||||
def test_creation_with_settings(self) -> None:
|
||||
"""Test creating engine with custom settings."""
|
||||
settings = ContextSettings(
|
||||
compression_threshold=0.7,
|
||||
cache_enabled=False,
|
||||
)
|
||||
engine = ContextEngine(settings=settings)
|
||||
|
||||
assert engine._settings.compression_threshold == 0.7
|
||||
assert engine._settings.cache_enabled is False
|
||||
|
||||
def test_creation_with_redis(self) -> None:
|
||||
"""Test creating engine with Redis."""
|
||||
mock_redis = MagicMock()
|
||||
settings = ContextSettings(cache_enabled=True)
|
||||
engine = ContextEngine(redis=mock_redis, settings=settings)
|
||||
|
||||
assert engine._cache.is_enabled
|
||||
|
||||
def test_set_mcp_manager(self) -> None:
|
||||
"""Test setting MCP manager."""
|
||||
engine = ContextEngine()
|
||||
mock_mcp = MagicMock()
|
||||
|
||||
engine.set_mcp_manager(mock_mcp)
|
||||
|
||||
assert engine._mcp is mock_mcp
|
||||
|
||||
def test_set_redis(self) -> None:
|
||||
"""Test setting Redis connection."""
|
||||
engine = ContextEngine()
|
||||
mock_redis = MagicMock()
|
||||
|
||||
engine.set_redis(mock_redis)
|
||||
|
||||
assert engine._cache._redis is mock_redis
|
||||
|
||||
|
||||
class TestContextEngineHelpers:
|
||||
"""Tests for ContextEngine helper methods."""
|
||||
|
||||
def test_convert_conversation(self) -> None:
|
||||
"""Test converting conversation history."""
|
||||
engine = ContextEngine()
|
||||
|
||||
history = [
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
{"role": "user", "content": "How are you?"},
|
||||
]
|
||||
|
||||
contexts = engine._convert_conversation(history)
|
||||
|
||||
assert len(contexts) == 3
|
||||
assert all(isinstance(c, ConversationContext) for c in contexts)
|
||||
assert contexts[0].role == MessageRole.USER
|
||||
assert contexts[1].role == MessageRole.ASSISTANT
|
||||
assert contexts[0].content == "Hello!"
|
||||
assert contexts[0].metadata["turn"] == 0
|
||||
|
||||
def test_convert_tool_results(self) -> None:
|
||||
"""Test converting tool results."""
|
||||
engine = ContextEngine()
|
||||
|
||||
results = [
|
||||
{"tool_name": "search", "content": "Result 1", "status": "success"},
|
||||
{"tool_name": "read", "result": {"file": "test.txt"}, "status": "success"},
|
||||
]
|
||||
|
||||
contexts = engine._convert_tool_results(results)
|
||||
|
||||
assert len(contexts) == 2
|
||||
assert all(isinstance(c, ToolContext) for c in contexts)
|
||||
assert contexts[0].content == "Result 1"
|
||||
assert contexts[0].metadata["tool_name"] == "search"
|
||||
# Dict content should be JSON serialized
|
||||
assert "file" in contexts[1].content
|
||||
assert "test.txt" in contexts[1].content
|
||||
|
||||
|
||||
class TestContextEngineAssembly:
|
||||
"""Tests for context assembly."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assemble_minimal(self) -> None:
|
||||
"""Test assembling with minimal inputs."""
|
||||
engine = ContextEngine()
|
||||
|
||||
result = await engine.assemble_context(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="test query",
|
||||
model="claude-3-sonnet",
|
||||
use_cache=False, # Disable cache for test
|
||||
)
|
||||
|
||||
assert isinstance(result, AssembledContext)
|
||||
assert result.context_count == 0 # No contexts provided
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assemble_with_system_prompt(self) -> None:
|
||||
"""Test assembling with system prompt."""
|
||||
engine = ContextEngine()
|
||||
|
||||
result = await engine.assemble_context(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="test query",
|
||||
model="claude-3-sonnet",
|
||||
system_prompt="You are a helpful assistant.",
|
||||
use_cache=False,
|
||||
)
|
||||
|
||||
assert result.context_count == 1
|
||||
assert "helpful assistant" in result.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assemble_with_task(self) -> None:
|
||||
"""Test assembling with task description."""
|
||||
engine = ContextEngine()
|
||||
|
||||
result = await engine.assemble_context(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="implement feature",
|
||||
model="claude-3-sonnet",
|
||||
task_description="Implement user authentication",
|
||||
use_cache=False,
|
||||
)
|
||||
|
||||
assert result.context_count == 1
|
||||
assert "authentication" in result.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assemble_with_conversation(self) -> None:
|
||||
"""Test assembling with conversation history."""
|
||||
engine = ContextEngine()
|
||||
|
||||
result = await engine.assemble_context(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="continue",
|
||||
model="claude-3-sonnet",
|
||||
conversation_history=[
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{"role": "assistant", "content": "Hi!"},
|
||||
],
|
||||
use_cache=False,
|
||||
)
|
||||
|
||||
assert result.context_count == 2
|
||||
assert "Hello" in result.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assemble_with_tool_results(self) -> None:
|
||||
"""Test assembling with tool results."""
|
||||
engine = ContextEngine()
|
||||
|
||||
result = await engine.assemble_context(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="continue",
|
||||
model="claude-3-sonnet",
|
||||
tool_results=[
|
||||
{"tool_name": "search", "content": "Found 5 results"},
|
||||
],
|
||||
use_cache=False,
|
||||
)
|
||||
|
||||
assert result.context_count == 1
|
||||
assert "Found 5 results" in result.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assemble_with_custom_contexts(self) -> None:
|
||||
"""Test assembling with custom contexts."""
|
||||
engine = ContextEngine()
|
||||
|
||||
custom = [
|
||||
KnowledgeContext(
|
||||
content="Custom knowledge.",
|
||||
source="custom",
|
||||
relevance_score=0.9,
|
||||
)
|
||||
]
|
||||
|
||||
result = await engine.assemble_context(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="test",
|
||||
model="claude-3-sonnet",
|
||||
custom_contexts=custom,
|
||||
use_cache=False,
|
||||
)
|
||||
|
||||
assert result.context_count == 1
|
||||
assert "Custom knowledge" in result.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_assemble_full_workflow(self) -> None:
|
||||
"""Test full assembly workflow."""
|
||||
engine = ContextEngine()
|
||||
|
||||
result = await engine.assemble_context(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="implement login",
|
||||
model="claude-3-sonnet",
|
||||
system_prompt="You are an expert Python developer.",
|
||||
task_description="Implement user authentication.",
|
||||
conversation_history=[
|
||||
{"role": "user", "content": "Can you help me implement JWT auth?"},
|
||||
],
|
||||
tool_results=[
|
||||
{"tool_name": "file_create", "content": "Created auth.py"},
|
||||
],
|
||||
use_cache=False,
|
||||
)
|
||||
|
||||
assert result.context_count >= 4
|
||||
assert result.total_tokens > 0
|
||||
assert result.model == "claude-3-sonnet"
|
||||
|
||||
# Check for expected content
|
||||
assert "expert Python developer" in result.content
|
||||
assert "authentication" in result.content
|
||||
|
||||
|
||||
class TestContextEngineKnowledge:
|
||||
"""Tests for knowledge fetching."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_knowledge_no_mcp(self) -> None:
|
||||
"""Test fetching knowledge without MCP returns empty."""
|
||||
engine = ContextEngine()
|
||||
|
||||
result = await engine._fetch_knowledge(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="test",
|
||||
)
|
||||
|
||||
assert result == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_knowledge_with_mcp(self) -> None:
|
||||
"""Test fetching knowledge with MCP."""
|
||||
mock_mcp = AsyncMock()
|
||||
mock_mcp.call_tool.return_value.data = {
|
||||
"results": [
|
||||
{
|
||||
"content": "Document content",
|
||||
"source_path": "docs/api.md",
|
||||
"score": 0.9,
|
||||
"chunk_id": "chunk-1",
|
||||
},
|
||||
{
|
||||
"content": "Another document",
|
||||
"source_path": "docs/auth.md",
|
||||
"score": 0.8,
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
engine = ContextEngine(mcp_manager=mock_mcp)
|
||||
|
||||
result = await engine._fetch_knowledge(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="authentication",
|
||||
)
|
||||
|
||||
assert len(result) == 2
|
||||
assert all(isinstance(c, KnowledgeContext) for c in result)
|
||||
assert result[0].content == "Document content"
|
||||
assert result[0].source == "docs/api.md"
|
||||
assert result[0].relevance_score == 0.9
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_knowledge_error_handling(self) -> None:
|
||||
"""Test knowledge fetch error handling."""
|
||||
mock_mcp = AsyncMock()
|
||||
mock_mcp.call_tool.side_effect = Exception("MCP error")
|
||||
|
||||
engine = ContextEngine(mcp_manager=mock_mcp)
|
||||
|
||||
# Should not raise, returns empty
|
||||
result = await engine._fetch_knowledge(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="test",
|
||||
)
|
||||
|
||||
assert result == []
|
||||
|
||||
|
||||
class TestContextEngineCaching:
|
||||
"""Tests for caching behavior."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_disabled(self) -> None:
|
||||
"""Test assembly with cache disabled."""
|
||||
engine = ContextEngine()
|
||||
|
||||
result = await engine.assemble_context(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="test",
|
||||
model="claude-3-sonnet",
|
||||
system_prompt="Test prompt",
|
||||
use_cache=False,
|
||||
)
|
||||
|
||||
assert not result.cache_hit
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_hit(self) -> None:
|
||||
"""Test cache hit."""
|
||||
mock_redis = AsyncMock()
|
||||
settings = ContextSettings(cache_enabled=True)
|
||||
engine = ContextEngine(redis=mock_redis, settings=settings)
|
||||
|
||||
# First call - cache miss
|
||||
mock_redis.get.return_value = None
|
||||
|
||||
result1 = await engine.assemble_context(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="test",
|
||||
model="claude-3-sonnet",
|
||||
system_prompt="Test prompt",
|
||||
)
|
||||
|
||||
# Second call - mock cache hit
|
||||
mock_redis.get.return_value = result1.to_json()
|
||||
|
||||
result2 = await engine.assemble_context(
|
||||
project_id="proj-123",
|
||||
agent_id="agent-456",
|
||||
query="test",
|
||||
model="claude-3-sonnet",
|
||||
system_prompt="Test prompt",
|
||||
)
|
||||
|
||||
assert result2.cache_hit
|
||||
|
||||
|
||||
class TestContextEngineUtilities:
|
||||
"""Tests for utility methods."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_budget_for_model(self) -> None:
|
||||
"""Test getting budget for model."""
|
||||
engine = ContextEngine()
|
||||
|
||||
budget = await engine.get_budget_for_model("claude-3-sonnet")
|
||||
|
||||
assert budget.total > 0
|
||||
assert budget.system > 0
|
||||
assert budget.knowledge > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_budget_with_max_tokens(self) -> None:
|
||||
"""Test getting budget with max tokens."""
|
||||
engine = ContextEngine()
|
||||
|
||||
budget = await engine.get_budget_for_model("claude-3-sonnet", max_tokens=5000)
|
||||
|
||||
assert budget.total == 5000
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_tokens(self) -> None:
|
||||
"""Test token counting."""
|
||||
engine = ContextEngine()
|
||||
|
||||
count = await engine.count_tokens("Hello world")
|
||||
|
||||
assert count > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalidate_cache(self) -> None:
|
||||
"""Test cache invalidation."""
|
||||
mock_redis = AsyncMock()
|
||||
|
||||
async def mock_scan_iter(match=None):
|
||||
for key in ["ctx:1", "ctx:2"]:
|
||||
yield key
|
||||
|
||||
mock_redis.scan_iter = mock_scan_iter
|
||||
|
||||
settings = ContextSettings(cache_enabled=True)
|
||||
engine = ContextEngine(redis=mock_redis, settings=settings)
|
||||
|
||||
deleted = await engine.invalidate_cache(pattern="*test*")
|
||||
|
||||
assert deleted >= 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_stats(self) -> None:
|
||||
"""Test getting engine stats."""
|
||||
engine = ContextEngine()
|
||||
|
||||
stats = await engine.get_stats()
|
||||
|
||||
assert "cache" in stats
|
||||
assert "settings" in stats
|
||||
assert "compression_threshold" in stats["settings"]
|
||||
|
||||
|
||||
class TestCreateContextEngine:
|
||||
"""Tests for factory function."""
|
||||
|
||||
def test_create_context_engine(self) -> None:
|
||||
"""Test factory function."""
|
||||
engine = create_context_engine()
|
||||
|
||||
assert isinstance(engine, ContextEngine)
|
||||
|
||||
def test_create_context_engine_with_settings(self) -> None:
|
||||
"""Test factory with settings."""
|
||||
settings = ContextSettings(cache_enabled=False)
|
||||
engine = create_context_engine(settings=settings)
|
||||
|
||||
assert engine._settings.cache_enabled is False
|
||||
Reference in New Issue
Block a user