diff --git a/headroom/memory/backends/local.py b/headroom/memory/backends/local.py index 9b31e407b..e13c2d2cb 100644 --- a/headroom/memory/backends/local.py +++ b/headroom/memory/backends/local.py @@ -316,13 +316,11 @@ class LocalBackend: entities: Optional filter by related entities. include_related: If True, expand results via knowledge graph. min_similarity: Minimum cosine similarity threshold. - session_id: Optional session filter (not yet implemented). + session_id: Optional session filter to isolate memories by session. Returns: List of MemorySearchResult objects with scores and related entities. """ - # Note: session_id filtering is not yet implemented in LocalBackend - _ = session_id # Acknowledge parameter for protocol compliance await self._ensure_initialized() assert self._hierarchical_memory is not None assert self._graph is not None @@ -331,6 +329,7 @@ class LocalBackend: vector_results = await self._hierarchical_memory.search( query=query, user_id=user_id, + session_id=session_id, top_k=top_k * 2 if include_related else top_k, # Over-fetch for deduplication min_similarity=min_similarity, ) @@ -386,6 +385,9 @@ class LocalBackend: for mem_id in new_memory_ids: memory = await self._hierarchical_memory.get(mem_id) if memory and memory.user_id == user_id: + # Filter by session_id if specified (security: prevent session leakage) + if session_id is not None and memory.session_id != session_id: + continue # Add with lower score since it's from graph expansion results.append( MemorySearchResult( diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 9e6ff7db6..252cdd68f 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -1608,8 +1608,10 @@ class HeadroomProxy: resp_json = None try: resp_json = response.json() - except Exception: - pass + except (json.JSONDecodeError, ValueError) as e: + logger.debug( + f"[{request_id}] Failed to parse response JSON for CCR handling: {e}" + ) # CCR Response Handling: Handle headroom_retrieve tool calls automatically if ( @@ -2961,9 +2963,9 @@ class HeadroomProxy: if usage: return usage - except Exception: + except (UnicodeDecodeError, KeyError, TypeError) as e: # Don't fail streaming on parse errors - pass + logger.debug(f"SSE usage parsing error for {provider}: {e}") return None @@ -3687,8 +3689,10 @@ class HeadroomProxy: # These are charged at 50% of the input price prompt_details = usage.get("prompt_tokens_details", {}) cache_read_tokens = prompt_details.get("cached_tokens", 0) - except Exception: - pass + except (KeyError, TypeError, AttributeError) as e: + logger.debug( + f"[{request_id}] Failed to extract cached tokens from OpenAI response: {e}" + ) # For OpenAI, prompt_tokens is TOTAL (includes cached) # Normalize to non-cached input for consistent cost calculation @@ -4423,8 +4427,10 @@ class HeadroomProxy: "prompt_tokens_details", usage.get("input_tokens_details", {}) ) cache_read_tokens = prompt_details.get("cached_tokens", 0) - except Exception: - pass + except (KeyError, TypeError, AttributeError) as e: + logger.debug( + f"[{request_id}] Failed to extract cached tokens from OpenAI passthrough response: {e}" + ) # For OpenAI, input_tokens is TOTAL (includes cached) # Normalize to non-cached input for consistent cost calculation @@ -4687,8 +4693,10 @@ class HeadroomProxy: # Gemini returns cachedContentTokenCount for context-cached tokens # These are charged at 10-25% of the input price depending on model cache_read_tokens = usage.get("cachedContentTokenCount", 0) - except Exception: - pass + except (KeyError, TypeError, AttributeError) as e: + logger.debug( + f"[{request_id}] Failed to extract cached tokens from Gemini response: {e}" + ) # For Gemini, promptTokenCount is TOTAL (includes cached) # Normalize to non-cached input for consistent cost calculation @@ -4939,8 +4947,8 @@ class HeadroomProxy: try: resp_json = response.json() compressed_tokens = resp_json.get("totalTokens", 0) - except Exception: - pass + except (json.JSONDecodeError, ValueError) as e: + logger.debug(f"[{request_id}] Failed to parse Gemini token count response: {e}") # Track stats tokens_saved = original_tokens - compressed_tokens if compressed_tokens > 0 else 0 diff --git a/tests/test_ccr_batch_processor.py b/tests/test_ccr_batch_processor.py new file mode 100644 index 000000000..f63fcc546 --- /dev/null +++ b/tests/test_ccr_batch_processor.py @@ -0,0 +1,1194 @@ +"""Tests for CCR batch result processor. + +These tests verify that: +1. BatchResultProcessor class initialization works correctly +2. Result parsing for Anthropic, OpenAI, and Google batch formats +3. CCR tool call detection in batch results +4. Continuation call handling works for all providers +5. Error cases and edge cases are handled gracefully +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from headroom.ccr.batch_processor import ( + BatchResultProcessor, + BatchResultProcessorConfig, + ProcessedBatchResult, + process_batch_results, +) +from headroom.ccr.batch_store import ( + BatchContext, + BatchContextStore, + BatchRequestContext, + reset_batch_context_store, +) +from headroom.ccr.tool_injection import CCR_TOOL_NAME + + +class TestBatchResultProcessorConfig: + """Test BatchResultProcessorConfig dataclass.""" + + def test_default_config(self): + """Default config values.""" + config = BatchResultProcessorConfig() + + assert config.enabled is True + assert config.continuation_timeout == 120 + assert config.max_continuation_rounds == 3 + + def test_custom_config(self): + """Custom config values.""" + config = BatchResultProcessorConfig( + enabled=False, + continuation_timeout=60, + max_continuation_rounds=5, + ) + + assert config.enabled is False + assert config.continuation_timeout == 60 + assert config.max_continuation_rounds == 5 + + +class TestProcessedBatchResult: + """Test ProcessedBatchResult dataclass.""" + + def test_default_values(self): + """Default values for ProcessedBatchResult.""" + result = ProcessedBatchResult( + custom_id="req_123", + result={"content": "test"}, + ) + + assert result.custom_id == "req_123" + assert result.result == {"content": "test"} + assert result.was_processed is False + assert result.continuation_rounds == 0 + assert result.error is None + + def test_processed_result(self): + """ProcessedBatchResult with CCR processing.""" + result = ProcessedBatchResult( + custom_id="req_456", + result={"content": "processed"}, + was_processed=True, + continuation_rounds=2, + ) + + assert result.was_processed is True + assert result.continuation_rounds == 2 + + def test_error_result(self): + """ProcessedBatchResult with error.""" + result = ProcessedBatchResult( + custom_id="req_789", + result={"content": "partial"}, + error="Retrieval failed", + ) + + assert result.error == "Retrieval failed" + + +class TestBatchResultProcessorInit: + """Test BatchResultProcessor initialization.""" + + def test_default_initialization(self): + """Initialize with default config.""" + http_client = MagicMock(spec=httpx.AsyncClient) + processor = BatchResultProcessor(http_client) + + assert processor.http_client == http_client + assert processor.config.enabled is True + assert processor.config.continuation_timeout == 120 + assert processor.ccr_handler is not None + + def test_custom_config_initialization(self): + """Initialize with custom config.""" + http_client = MagicMock(spec=httpx.AsyncClient) + config = BatchResultProcessorConfig( + enabled=False, + continuation_timeout=60, + ) + processor = BatchResultProcessor(http_client, config) + + assert processor.config.enabled is False + assert processor.config.continuation_timeout == 60 + + def test_api_urls_set(self): + """API URLs are set for all providers.""" + http_client = MagicMock(spec=httpx.AsyncClient) + processor = BatchResultProcessor(http_client) + + assert "anthropic" in processor.api_urls + assert "openai" in processor.api_urls + assert "google" in processor.api_urls + assert processor.api_urls["anthropic"] == "https://api.anthropic.com" + assert processor.api_urls["openai"] == "https://api.openai.com" + assert processor.api_urls["google"] == "https://generativelanguage.googleapis.com" + + +class TestCustomIdExtraction: + """Test _get_custom_id method for different providers.""" + + @pytest.fixture + def processor(self): + """Create a processor instance.""" + http_client = MagicMock(spec=httpx.AsyncClient) + return BatchResultProcessor(http_client) + + def test_anthropic_custom_id(self, processor): + """Extract custom_id from Anthropic batch result.""" + result = {"custom_id": "anthropic_req_123", "result": {"message": {}}} + custom_id = processor._get_custom_id(result, "anthropic") + assert custom_id == "anthropic_req_123" + + def test_openai_custom_id(self, processor): + """Extract custom_id from OpenAI batch result.""" + result = {"custom_id": "openai_req_456", "response": {"body": {}}} + custom_id = processor._get_custom_id(result, "openai") + assert custom_id == "openai_req_456" + + def test_google_custom_id(self, processor): + """Extract custom_id from Google batch result (metadata.key).""" + result = {"metadata": {"key": "google_req_789"}, "response": {}} + custom_id = processor._get_custom_id(result, "google") + assert custom_id == "google_req_789" + + def test_google_missing_metadata(self, processor): + """Handle missing metadata in Google result.""" + result = {"response": {}} + custom_id = processor._get_custom_id(result, "google") + assert custom_id == "" + + def test_unknown_provider_fallback(self, processor): + """Fallback extraction for unknown provider.""" + result = {"custom_id": "unknown_req", "id": "backup_id"} + custom_id = processor._get_custom_id(result, "unknown") + assert custom_id == "unknown_req" + + def test_unknown_provider_uses_id_fallback(self, processor): + """Unknown provider falls back to 'id' field.""" + result = {"id": "id_field_value"} + custom_id = processor._get_custom_id(result, "unknown") + assert custom_id == "id_field_value" + + +class TestResponseExtraction: + """Test _extract_response method for different providers.""" + + @pytest.fixture + def processor(self): + """Create a processor instance.""" + http_client = MagicMock(spec=httpx.AsyncClient) + return BatchResultProcessor(http_client) + + def test_anthropic_response_extraction(self, processor): + """Extract response from Anthropic batch result.""" + result = { + "custom_id": "req_1", + "result": { + "type": "message", + "message": { + "content": [{"type": "text", "text": "Hello"}], + "stop_reason": "end_turn", + }, + }, + } + response = processor._extract_response(result, "anthropic") + assert response is not None + assert response["content"] == [{"type": "text", "text": "Hello"}] + + def test_openai_response_extraction(self, processor): + """Extract response from OpenAI batch result.""" + result = { + "custom_id": "req_2", + "response": { + "status_code": 200, + "body": { + "choices": [{"message": {"content": "Hello"}}], + }, + }, + } + response = processor._extract_response(result, "openai") + assert response is not None + assert response["choices"][0]["message"]["content"] == "Hello" + + def test_google_response_extraction(self, processor): + """Extract response from Google batch result.""" + result = { + "metadata": {"key": "req_3"}, + "response": { + "candidates": [{"content": {"parts": [{"text": "Hello"}]}}], + }, + } + response = processor._extract_response(result, "google") + assert response is not None + assert response["candidates"][0]["content"]["parts"][0]["text"] == "Hello" + + def test_anthropic_missing_result(self, processor): + """Handle missing result in Anthropic format.""" + result = {"custom_id": "req_4"} + response = processor._extract_response(result, "anthropic") + assert response is None + + def test_openai_missing_body(self, processor): + """Handle missing body in OpenAI format.""" + result = {"custom_id": "req_5", "response": {"status_code": 500}} + response = processor._extract_response(result, "openai") + assert response is None + + def test_invalid_response_type(self, processor): + """Handle non-dict response.""" + result = {"custom_id": "req_6", "result": {"message": "not_a_dict"}} + response = processor._extract_response(result, "anthropic") + assert response is None + + +class TestCCRToolCallDetectionInBatch: + """Test CCR tool call detection within batch results.""" + + @pytest.fixture + def processor(self): + """Create a processor instance.""" + http_client = MagicMock(spec=httpx.AsyncClient) + return BatchResultProcessor(http_client) + + def test_detect_anthropic_ccr_in_batch(self, processor): + """Detect CCR tool call in Anthropic batch result.""" + response = { + "content": [ + {"type": "text", "text": "Let me retrieve that data."}, + { + "type": "tool_use", + "id": "tool_123", + "name": CCR_TOOL_NAME, + "input": {"hash": "abc123"}, + }, + ] + } + assert processor.ccr_handler.has_ccr_tool_calls(response, "anthropic") + + def test_detect_openai_ccr_in_batch(self, processor): + """Detect CCR tool call in OpenAI batch result.""" + response = { + "choices": [ + { + "message": { + "role": "assistant", + "content": "Retrieving data...", + "tool_calls": [ + { + "id": "call_abc", + "type": "function", + "function": { + "name": CCR_TOOL_NAME, + "arguments": '{"hash": "def456"}', + }, + } + ], + } + } + ] + } + assert processor.ccr_handler.has_ccr_tool_calls(response, "openai") + + def test_detect_google_ccr_in_batch(self, processor): + """Detect CCR tool call in Google batch result.""" + response = { + "candidates": [ + { + "content": { + "parts": [ + {"text": "Retrieving data..."}, + { + "functionCall": { + "name": CCR_TOOL_NAME, + "args": {"hash": "ghi789"}, + } + }, + ] + } + } + ] + } + assert processor.ccr_handler.has_ccr_tool_calls(response, "google") + + def test_no_ccr_in_text_only_response(self, processor): + """No false positive for text-only response.""" + response = {"content": [{"type": "text", "text": "Just a text response."}]} + assert not processor.ccr_handler.has_ccr_tool_calls(response, "anthropic") + + def test_no_ccr_for_other_tools(self, processor): + """No false positive for other tool calls.""" + response = { + "content": [ + { + "type": "tool_use", + "id": "tool_xyz", + "name": "read_file", + "input": {"path": "/etc/config"}, + } + ] + } + assert not processor.ccr_handler.has_ccr_tool_calls(response, "anthropic") + + +class TestResultUpdate: + """Test _update_result method for different providers.""" + + @pytest.fixture + def processor(self): + """Create a processor instance.""" + http_client = MagicMock(spec=httpx.AsyncClient) + return BatchResultProcessor(http_client) + + def test_update_anthropic_result(self, processor): + """Update Anthropic batch result with final response.""" + original = { + "custom_id": "req_1", + "result": { + "type": "tool_use", + "message": {"content": [{"type": "tool_use", "name": CCR_TOOL_NAME}]}, + }, + } + final_response = { + "content": [{"type": "text", "text": "Final answer"}], + "stop_reason": "end_turn", + } + + updated = processor._update_result(original, final_response, "anthropic") + + assert updated["result"]["message"] == final_response + assert updated["result"]["type"] == "succeeded" + assert updated["custom_id"] == "req_1" + + def test_update_openai_result(self, processor): + """Update OpenAI batch result with final response.""" + original = { + "custom_id": "req_2", + "response": { + "status_code": 200, + "body": { + "choices": [ + {"message": {"tool_calls": [{"function": {"name": CCR_TOOL_NAME}}]}} + ] + }, + }, + } + final_response = {"choices": [{"message": {"content": "Final answer"}}]} + + updated = processor._update_result(original, final_response, "openai") + + assert updated["response"]["body"] == final_response + assert updated["custom_id"] == "req_2" + + def test_update_google_result(self, processor): + """Update Google batch result with final response.""" + original = { + "metadata": {"key": "req_3"}, + "response": { + "candidates": [{"content": {"parts": [{"functionCall": {"name": CCR_TOOL_NAME}}]}}] + }, + } + final_response = {"candidates": [{"content": {"parts": [{"text": "Final answer"}]}}]} + + updated = processor._update_result(original, final_response, "google") + + assert updated["response"] == final_response + + def test_update_creates_missing_containers(self, processor): + """Update creates missing result/response containers.""" + original_anthropic = {"custom_id": "req_1"} + original_openai = {"custom_id": "req_2"} + final = {"content": "test"} + + updated_anthropic = processor._update_result(original_anthropic, final, "anthropic") + updated_openai = processor._update_result(original_openai, final, "openai") + + assert "result" in updated_anthropic + assert "response" in updated_openai + + +class TestMessagesToGoogleContents: + """Test _messages_to_google_contents conversion.""" + + @pytest.fixture + def processor(self): + """Create a processor instance.""" + http_client = MagicMock(spec=httpx.AsyncClient) + return BatchResultProcessor(http_client) + + def test_convert_simple_text_message(self, processor): + """Convert simple text messages.""" + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there"}, + ] + + contents = processor._messages_to_google_contents(messages) + + assert len(contents) == 2 + assert contents[0]["role"] == "user" + assert contents[0]["parts"] == [{"text": "Hello"}] + assert contents[1]["role"] == "model" + assert contents[1]["parts"] == [{"text": "Hi there"}] + + def test_skip_system_messages(self, processor): + """System messages are skipped (handled separately in Google).""" + messages = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + ] + + contents = processor._messages_to_google_contents(messages) + + assert len(contents) == 1 + assert contents[0]["role"] == "user" + + def test_convert_tool_result_content(self, processor): + """Convert structured content with tool results.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "tool_123", "content": "Result data"} + ], + } + ] + + contents = processor._messages_to_google_contents(messages) + + assert len(contents) == 1 + assert contents[0]["parts"][0]["functionResponse"]["response"]["content"] == "Result data" + + def test_convert_tool_use_content(self, processor): + """Convert content with tool_use blocks.""" + messages = [ + { + "role": "assistant", + "content": [{"type": "tool_use", "name": "read_file", "input": {"path": "/test"}}], + } + ] + + contents = processor._messages_to_google_contents(messages) + + assert len(contents) == 1 + assert contents[0]["role"] == "model" + assert contents[0]["parts"][0]["functionCall"]["name"] == "read_file" + + def test_preserve_google_format_messages(self, processor): + """Messages already in Google format are preserved.""" + messages = [{"role": "model", "parts": [{"text": "Already Google format"}]}] + + contents = processor._messages_to_google_contents(messages) + + assert len(contents) == 1 + assert contents[0]["parts"] == [{"text": "Already Google format"}] + + +class TestProcessResults: + """Test process_results method.""" + + @pytest.fixture(autouse=True) + def reset_store(self): + """Reset global store before each test.""" + reset_batch_context_store() + yield + reset_batch_context_store() + + @pytest.fixture + def processor(self): + """Create a processor instance.""" + http_client = AsyncMock(spec=httpx.AsyncClient) + return BatchResultProcessor(http_client) + + @pytest.mark.asyncio + async def test_disabled_processor_passthrough(self, processor): + """Disabled processor passes through results unchanged.""" + processor.config.enabled = False + + results = [{"custom_id": "req_1", "result": {"message": {"content": "test"}}}] + + processed = await processor.process_results("batch_123", results, "anthropic") + + assert len(processed) == 1 + assert processed[0].custom_id == "req_1" + assert processed[0].was_processed is False + + @pytest.mark.asyncio + async def test_missing_batch_context_passthrough(self, processor): + """Missing batch context passes through results unchanged.""" + results = [{"custom_id": "req_1", "result": {"message": {"content": "test"}}}] + + processed = await processor.process_results("nonexistent_batch", results, "anthropic") + + assert len(processed) == 1 + assert processed[0].was_processed is False + + @pytest.mark.asyncio + async def test_no_ccr_tool_calls_passthrough(self): + """Results without CCR tool calls pass through.""" + http_client = AsyncMock(spec=httpx.AsyncClient) + processor = BatchResultProcessor(http_client) + + # Set up batch context + store = BatchContextStore() + context = BatchContext(batch_id="batch_123", provider="anthropic") + context.add_request( + BatchRequestContext( + custom_id="req_1", + messages=[{"role": "user", "content": "Hi"}], + model="claude-3-opus", + ) + ) + await store.store(context) + + with patch( + "headroom.ccr.batch_processor.get_batch_context_store", + return_value=store, + ): + results = [ + { + "custom_id": "req_1", + "result": { + "message": { + "content": [{"type": "text", "text": "Hello!"}], + "stop_reason": "end_turn", + } + }, + } + ] + + processed = await processor.process_results("batch_123", results, "anthropic") + + assert len(processed) == 1 + assert processed[0].was_processed is False + + @pytest.mark.asyncio + async def test_missing_request_context_passthrough(self): + """Missing request context for custom_id passes through.""" + http_client = AsyncMock(spec=httpx.AsyncClient) + processor = BatchResultProcessor(http_client) + + # Set up batch context without matching request + store = BatchContextStore() + context = BatchContext(batch_id="batch_123", provider="anthropic") + # Don't add any requests + await store.store(context) + + with patch( + "headroom.ccr.batch_processor.get_batch_context_store", + return_value=store, + ): + results = [ + { + "custom_id": "unknown_req", + "result": {"message": {"content": [{"type": "text", "text": "Test"}]}}, + } + ] + + processed = await processor.process_results("batch_123", results, "anthropic") + + assert len(processed) == 1 + assert processed[0].custom_id == "unknown_req" + assert processed[0].was_processed is False + + +class TestContinuationCalls: + """Test continuation API calls for different providers.""" + + @pytest.fixture(autouse=True) + def reset_store(self): + """Reset global store before each test.""" + reset_batch_context_store() + yield + reset_batch_context_store() + + @pytest.mark.asyncio + async def test_anthropic_continuation_call(self): + """Test Anthropic continuation call format.""" + mock_response = MagicMock() + mock_response.json.return_value = {"content": [{"type": "text", "text": "Final answer"}]} + mock_response.raise_for_status = MagicMock() + + http_client = AsyncMock(spec=httpx.AsyncClient) + http_client.post.return_value = mock_response + + processor = BatchResultProcessor(http_client) + + request_context = BatchRequestContext( + custom_id="req_1", + messages=[{"role": "user", "content": "Hello"}], + model="claude-3-opus-20240229", + extras={"max_tokens": 1000}, + ) + batch_context = BatchContext( + batch_id="batch_123", + provider="anthropic", + api_key="test_api_key", + ) + + messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Tool result here"}, + ] + + await processor._anthropic_continuation(messages, None, request_context, batch_context) + + # Verify API call + http_client.post.assert_called_once() + call_args = http_client.post.call_args + + assert "api.anthropic.com" in call_args.args[0] + assert call_args.kwargs["headers"]["x-api-key"] == "test_api_key" + assert call_args.kwargs["json"]["model"] == "claude-3-opus-20240229" + + @pytest.mark.asyncio + async def test_openai_continuation_call(self): + """Test OpenAI continuation call format.""" + mock_response = MagicMock() + mock_response.json.return_value = {"choices": [{"message": {"content": "Final answer"}}]} + mock_response.raise_for_status = MagicMock() + + http_client = AsyncMock(spec=httpx.AsyncClient) + http_client.post.return_value = mock_response + + processor = BatchResultProcessor(http_client) + + request_context = BatchRequestContext( + custom_id="req_1", + messages=[{"role": "user", "content": "Hello"}], + model="gpt-4", + ) + batch_context = BatchContext( + batch_id="batch_123", + provider="openai", + api_key="sk-test123", + ) + + await processor._openai_continuation( + [{"role": "user", "content": "Hello"}], + None, + request_context, + batch_context, + ) + + # Verify API call + http_client.post.assert_called_once() + call_args = http_client.post.call_args + + assert "api.openai.com" in call_args.args[0] + assert "Bearer sk-test123" in call_args.kwargs["headers"]["Authorization"] + assert call_args.kwargs["json"]["model"] == "gpt-4" + + @pytest.mark.asyncio + async def test_google_continuation_call(self): + """Test Google continuation call format.""" + mock_response = MagicMock() + mock_response.json.return_value = { + "candidates": [{"content": {"parts": [{"text": "Final answer"}]}}] + } + mock_response.raise_for_status = MagicMock() + + http_client = AsyncMock(spec=httpx.AsyncClient) + http_client.post.return_value = mock_response + + processor = BatchResultProcessor(http_client) + + request_context = BatchRequestContext( + custom_id="req_1", + messages=[{"role": "user", "content": "Hello"}], + model="gemini-pro", + system_instruction="Be helpful.", + ) + batch_context = BatchContext( + batch_id="batch_123", + provider="google", + api_key="google_api_key", + ) + + await processor._google_continuation( + [{"role": "user", "content": "Hello"}], + [{"name": "test_tool", "parameters": {}}], + request_context, + batch_context, + ) + + # Verify API call + http_client.post.assert_called_once() + call_args = http_client.post.call_args + + assert "generativelanguage.googleapis.com" in call_args.args[0] + assert "gemini-pro" in call_args.args[0] + assert "key=google_api_key" in call_args.args[0] + assert "contents" in call_args.kwargs["json"] + assert "systemInstruction" in call_args.kwargs["json"] + + @pytest.mark.asyncio + async def test_continuation_with_tools(self): + """Test continuation includes tools when present.""" + mock_response = MagicMock() + mock_response.json.return_value = {"content": [{"type": "text", "text": "Done"}]} + mock_response.raise_for_status = MagicMock() + + http_client = AsyncMock(spec=httpx.AsyncClient) + http_client.post.return_value = mock_response + + processor = BatchResultProcessor(http_client) + + request_context = BatchRequestContext( + custom_id="req_1", + messages=[], + model="claude-3-opus", + extras={"max_tokens": 1000}, + ) + batch_context = BatchContext( + batch_id="batch_123", + provider="anthropic", + api_key="test_key", + ) + + tools = [ + {"name": "read_file", "input_schema": {"type": "object"}}, + {"name": CCR_TOOL_NAME, "input_schema": {"type": "object"}}, + ] + + await processor._anthropic_continuation( + [{"role": "user", "content": "test"}], + tools, + request_context, + batch_context, + ) + + call_args = http_client.post.call_args + assert "tools" in call_args.kwargs["json"] + assert len(call_args.kwargs["json"]["tools"]) == 2 + + @pytest.mark.asyncio + async def test_make_continuation_call_unknown_provider(self): + """Test continuation call raises for unknown provider.""" + http_client = AsyncMock(spec=httpx.AsyncClient) + processor = BatchResultProcessor(http_client) + + request_context = BatchRequestContext( + custom_id="req_1", + messages=[], + model="unknown-model", + ) + batch_context = BatchContext( + batch_id="batch_123", + provider="unknown", + ) + + with pytest.raises(ValueError, match="Unknown provider"): + await processor._make_continuation_call( + [], + None, + request_context, + batch_context, + "unknown", + ) + + +class TestProcessSingleResult: + """Test _process_single_result method.""" + + @pytest.fixture(autouse=True) + def reset_store(self): + """Reset global store before each test.""" + reset_batch_context_store() + yield + reset_batch_context_store() + + @pytest.mark.asyncio + async def test_process_single_result_with_ccr(self): + """Process a single result containing CCR tool call.""" + # Mock the CCR handler to simulate CCR processing + mock_response = MagicMock() + mock_response.json.return_value = { + "content": [{"type": "text", "text": "Final processed answer"}] + } + mock_response.raise_for_status = MagicMock() + + http_client = AsyncMock(spec=httpx.AsyncClient) + http_client.post.return_value = mock_response + + processor = BatchResultProcessor(http_client) + + # Mock the CCR handler's handle_response + processor.ccr_handler.handle_response = AsyncMock( + return_value={"content": [{"type": "text", "text": "Final processed answer"}]} + ) + + original_result = { + "custom_id": "req_1", + "result": { + "message": { + "content": [ + { + "type": "tool_use", + "id": "tool_123", + "name": CCR_TOOL_NAME, + "input": {"hash": "abc123"}, + } + ] + } + }, + } + response = original_result["result"]["message"] + + request_context = BatchRequestContext( + custom_id="req_1", + messages=[{"role": "user", "content": "Get data"}], + tools=[{"name": CCR_TOOL_NAME}], + model="claude-3-opus", + ) + batch_context = BatchContext( + batch_id="batch_123", + provider="anthropic", + api_key="test_key", + ) + + processed = await processor._process_single_result( + original_result, + response, + request_context, + batch_context, + "anthropic", + ) + + assert processed.custom_id == "req_1" + assert processed.was_processed is True + + +class TestConvenienceFunction: + """Test process_batch_results convenience function.""" + + @pytest.fixture(autouse=True) + def reset_store(self): + """Reset global store before each test.""" + reset_batch_context_store() + yield + reset_batch_context_store() + + @pytest.mark.asyncio + async def test_process_batch_results_function(self): + """Test the convenience function.""" + http_client = AsyncMock(spec=httpx.AsyncClient) + + results = [ + { + "custom_id": "req_1", + "result": {"message": {"content": [{"type": "text", "text": "Hello"}]}}, + } + ] + + processed = await process_batch_results( + "batch_123", + results, + "anthropic", + http_client, + ) + + assert len(processed) == 1 + assert processed[0].custom_id == "req_1" + + +class TestErrorHandling: + """Test error handling scenarios.""" + + @pytest.fixture(autouse=True) + def reset_store(self): + """Reset global store before each test.""" + reset_batch_context_store() + yield + reset_batch_context_store() + + @pytest.mark.asyncio + async def test_continuation_api_error(self): + """Handle API errors during continuation.""" + http_client = AsyncMock(spec=httpx.AsyncClient) + http_client.post.side_effect = httpx.HTTPStatusError( + "Internal Server Error", + request=MagicMock(), + response=MagicMock(status_code=500), + ) + + processor = BatchResultProcessor(http_client) + + request_context = BatchRequestContext( + custom_id="req_1", + messages=[], + model="claude-3-opus", + extras={"max_tokens": 1000}, + ) + batch_context = BatchContext( + batch_id="batch_123", + provider="anthropic", + api_key="test_key", + ) + + with pytest.raises(httpx.HTTPStatusError): + await processor._anthropic_continuation( + [{"role": "user", "content": "test"}], + None, + request_context, + batch_context, + ) + + @pytest.mark.asyncio + async def test_processing_error_captured(self): + """Processing errors are captured in result.""" + http_client = AsyncMock(spec=httpx.AsyncClient) + processor = BatchResultProcessor(http_client) + + # Set up batch context + store = BatchContextStore() + context = BatchContext(batch_id="batch_123", provider="anthropic") + context.add_request( + BatchRequestContext( + custom_id="req_1", + messages=[{"role": "user", "content": "Hi"}], + model="claude-3-opus", + ) + ) + await store.store(context) + + # Mock the handler to raise an error + processor.ccr_handler.has_ccr_tool_calls = MagicMock(return_value=True) + processor._process_single_result = AsyncMock(side_effect=Exception("Processing failed")) + + with patch( + "headroom.ccr.batch_processor.get_batch_context_store", + return_value=store, + ): + results = [ + { + "custom_id": "req_1", + "result": { + "message": { + "content": [ + { + "type": "tool_use", + "id": "tool_1", + "name": CCR_TOOL_NAME, + "input": {"hash": "abc"}, + } + ] + } + }, + } + ] + + processed = await processor.process_results("batch_123", results, "anthropic") + + assert len(processed) == 1 + assert processed[0].error == "Processing failed" + assert processed[0].was_processed is False + + +class TestMultipleResultsProcessing: + """Test processing multiple results in a batch.""" + + @pytest.fixture(autouse=True) + def reset_store(self): + """Reset global store before each test.""" + reset_batch_context_store() + yield + reset_batch_context_store() + + @pytest.mark.asyncio + async def test_mixed_results_processing(self): + """Process batch with mix of CCR and non-CCR results.""" + http_client = AsyncMock(spec=httpx.AsyncClient) + processor = BatchResultProcessor(http_client) + + # Set up batch context with multiple requests + store = BatchContextStore() + context = BatchContext(batch_id="batch_123", provider="anthropic") + context.add_request( + BatchRequestContext( + custom_id="req_1", + messages=[{"role": "user", "content": "Request 1"}], + model="claude-3-opus", + ) + ) + context.add_request( + BatchRequestContext( + custom_id="req_2", + messages=[{"role": "user", "content": "Request 2"}], + model="claude-3-opus", + ) + ) + context.add_request( + BatchRequestContext( + custom_id="req_3", + messages=[{"role": "user", "content": "Request 3"}], + model="claude-3-opus", + ) + ) + await store.store(context) + + with patch( + "headroom.ccr.batch_processor.get_batch_context_store", + return_value=store, + ): + results = [ + # Non-CCR result + { + "custom_id": "req_1", + "result": { + "message": {"content": [{"type": "text", "text": "Simple response"}]} + }, + }, + # Non-CCR result + { + "custom_id": "req_2", + "result": { + "message": {"content": [{"type": "text", "text": "Another response"}]} + }, + }, + # Non-CCR result (other tool) + { + "custom_id": "req_3", + "result": { + "message": { + "content": [ + { + "type": "tool_use", + "id": "tool_1", + "name": "read_file", + "input": {"path": "/test"}, + } + ] + } + }, + }, + ] + + processed = await processor.process_results("batch_123", results, "anthropic") + + assert len(processed) == 3 + # All should not be processed (no CCR tools) + assert all(not r.was_processed for r in processed) + assert processed[0].custom_id == "req_1" + assert processed[1].custom_id == "req_2" + assert processed[2].custom_id == "req_3" + + +class TestProviderSpecificFormats: + """Test provider-specific batch result formats.""" + + @pytest.fixture(autouse=True) + def reset_store(self): + """Reset global store before each test.""" + reset_batch_context_store() + yield + reset_batch_context_store() + + @pytest.mark.asyncio + async def test_anthropic_batch_format(self): + """Test full Anthropic batch result format handling.""" + http_client = AsyncMock(spec=httpx.AsyncClient) + processor = BatchResultProcessor(http_client) + + # Typical Anthropic batch result format + results = [ + { + "custom_id": "my-request-1", + "result": { + "type": "succeeded", + "message": { + "id": "msg_123", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "Response text"}], + "model": "claude-3-opus-20240229", + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 20}, + }, + }, + } + ] + + processed = await processor.process_results("batch_1", results, "anthropic") + + assert processed[0].custom_id == "my-request-1" + assert processed[0].result["result"]["message"]["content"][0]["text"] == "Response text" + + @pytest.mark.asyncio + async def test_openai_batch_format(self): + """Test full OpenAI batch result format handling.""" + http_client = AsyncMock(spec=httpx.AsyncClient) + processor = BatchResultProcessor(http_client) + + # Typical OpenAI batch result format + results = [ + { + "id": "batch_req_123", + "custom_id": "request-1", + "response": { + "status_code": 200, + "request_id": "req_abc", + "body": { + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1234567890, + "model": "gpt-4", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "OpenAI response", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + }, + }, + "error": None, + } + ] + + processed = await processor.process_results("batch_1", results, "openai") + + assert processed[0].custom_id == "request-1" + + @pytest.mark.asyncio + async def test_google_batch_format(self): + """Test full Google batch result format handling.""" + http_client = AsyncMock(spec=httpx.AsyncClient) + processor = BatchResultProcessor(http_client) + + # Typical Google batch result format + results = [ + { + "metadata": { + "key": "request-abc", + }, + "response": { + "candidates": [ + { + "content": { + "parts": [{"text": "Google response text"}], + "role": "model", + }, + "finishReason": "STOP", + "safetyRatings": [], + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 20, + }, + }, + } + ] + + processed = await processor.process_results("batch_1", results, "google") + + assert processed[0].custom_id == "request-abc" diff --git a/tests/test_compression_store.py b/tests/test_compression_store.py new file mode 100644 index 000000000..6e77f650a --- /dev/null +++ b/tests/test_compression_store.py @@ -0,0 +1,1259 @@ +"""Comprehensive unit tests for CompressionStore. + +Tests cover: +1. CompressionStore class initialization +2. Storing compressed content with hash generation +3. Retrieving content by hash +4. TTL expiration behavior +5. Memory limits and eviction +6. Statistics tracking +7. Edge cases (empty content, duplicate stores, etc.) +8. Thread safety +9. Feedback loop integration +10. Search functionality +""" + +from __future__ import annotations + +import hashlib +import json +import threading +import time +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from headroom.cache.compression_store import ( + CompressionEntry, + CompressionStore, + RetrievalEvent, + get_compression_store, + reset_compression_store, +) + +# ============================================================================= +# Fixtures +# ============================================================================= + + +@pytest.fixture(autouse=True) +def reset_global_store(): + """Reset global compression store before and after each test.""" + reset_compression_store() + yield + reset_compression_store() + + +@pytest.fixture +def store() -> CompressionStore: + """Create a fresh CompressionStore instance for testing.""" + return CompressionStore() + + +@pytest.fixture +def store_with_short_ttl() -> CompressionStore: + """Create a CompressionStore with 1 second TTL for expiration tests.""" + return CompressionStore(default_ttl=1) + + +@pytest.fixture +def store_with_small_capacity() -> CompressionStore: + """Create a CompressionStore with small capacity for eviction tests.""" + return CompressionStore(max_entries=3) + + +@pytest.fixture +def sample_items() -> list[dict[str, Any]]: + """Sample list of items for testing.""" + return [{"id": i, "name": f"item_{i}", "value": i * 10} for i in range(100)] + + +@pytest.fixture +def sample_original(sample_items: list[dict[str, Any]]) -> str: + """Sample original JSON content.""" + return json.dumps(sample_items) + + +@pytest.fixture +def sample_compressed(sample_items: list[dict[str, Any]]) -> str: + """Sample compressed JSON content (first 10 items).""" + return json.dumps(sample_items[:10]) + + +# ============================================================================= +# CompressionEntry Tests +# ============================================================================= + + +class TestCompressionEntry: + """Tests for CompressionEntry dataclass.""" + + def test_entry_creation_with_defaults(self): + """CompressionEntry can be created with minimal required fields.""" + entry = CompressionEntry( + hash="abc123", + original_content="[1,2,3]", + compressed_content="[1]", + original_tokens=100, + compressed_tokens=10, + original_item_count=3, + compressed_item_count=1, + tool_name=None, + tool_call_id=None, + query_context=None, + created_at=time.time(), + ) + assert entry.hash == "abc123" + assert entry.ttl == 300 # Default TTL + assert entry.retrieval_count == 0 + assert entry.search_queries == [] + assert entry.last_accessed is None + + def test_entry_is_expired_false_when_fresh(self): + """Fresh entries are not expired.""" + entry = CompressionEntry( + hash="abc123", + original_content="[1]", + compressed_content="[]", + original_tokens=10, + compressed_tokens=0, + original_item_count=1, + compressed_item_count=0, + tool_name=None, + tool_call_id=None, + query_context=None, + created_at=time.time(), + ttl=300, + ) + assert entry.is_expired() is False + + def test_entry_is_expired_true_after_ttl(self): + """Entries are expired after TTL passes.""" + entry = CompressionEntry( + hash="abc123", + original_content="[1]", + compressed_content="[]", + original_tokens=10, + compressed_tokens=0, + original_item_count=1, + compressed_item_count=0, + tool_name=None, + tool_call_id=None, + query_context=None, + created_at=time.time() - 10, # 10 seconds ago + ttl=5, # 5 second TTL + ) + assert entry.is_expired() is True + + def test_record_access_increments_count(self): + """record_access increments retrieval_count.""" + entry = CompressionEntry( + hash="abc123", + original_content="[1]", + compressed_content="[]", + original_tokens=10, + compressed_tokens=0, + original_item_count=1, + compressed_item_count=0, + tool_name=None, + tool_call_id=None, + query_context=None, + created_at=time.time(), + ) + assert entry.retrieval_count == 0 + + entry.record_access() + assert entry.retrieval_count == 1 + + entry.record_access() + assert entry.retrieval_count == 2 + + def test_record_access_updates_last_accessed(self): + """record_access updates last_accessed timestamp.""" + entry = CompressionEntry( + hash="abc123", + original_content="[1]", + compressed_content="[]", + original_tokens=10, + compressed_tokens=0, + original_item_count=1, + compressed_item_count=0, + tool_name=None, + tool_call_id=None, + query_context=None, + created_at=time.time(), + ) + assert entry.last_accessed is None + + before = time.time() + entry.record_access() + after = time.time() + + assert entry.last_accessed is not None + assert before <= entry.last_accessed <= after + + def test_record_access_tracks_unique_queries(self): + """record_access tracks unique search queries.""" + entry = CompressionEntry( + hash="abc123", + original_content="[1]", + compressed_content="[]", + original_tokens=10, + compressed_tokens=0, + original_item_count=1, + compressed_item_count=0, + tool_name=None, + tool_call_id=None, + query_context=None, + created_at=time.time(), + ) + + entry.record_access(query="query1") + entry.record_access(query="query2") + entry.record_access(query="query1") # Duplicate + + assert "query1" in entry.search_queries + assert "query2" in entry.search_queries + assert len(entry.search_queries) == 2 # No duplicates + + def test_record_access_limits_queries_to_10(self): + """record_access keeps only last 10 queries.""" + entry = CompressionEntry( + hash="abc123", + original_content="[1]", + compressed_content="[]", + original_tokens=10, + compressed_tokens=0, + original_item_count=1, + compressed_item_count=0, + tool_name=None, + tool_call_id=None, + query_context=None, + created_at=time.time(), + ) + + for i in range(15): + entry.record_access(query=f"query_{i}") + + assert len(entry.search_queries) == 10 + # Should have the last 10 queries + assert "query_5" in entry.search_queries + assert "query_14" in entry.search_queries + assert "query_0" not in entry.search_queries + + def test_record_access_ignores_none_query(self): + """record_access does not add None queries to list.""" + entry = CompressionEntry( + hash="abc123", + original_content="[1]", + compressed_content="[]", + original_tokens=10, + compressed_tokens=0, + original_item_count=1, + compressed_item_count=0, + tool_name=None, + tool_call_id=None, + query_context=None, + created_at=time.time(), + ) + + entry.record_access(query=None) + entry.record_access() + + assert len(entry.search_queries) == 0 + + +# ============================================================================= +# CompressionStore Initialization Tests +# ============================================================================= + + +class TestCompressionStoreInit: + """Tests for CompressionStore initialization.""" + + def test_default_initialization(self): + """CompressionStore initializes with default values.""" + store = CompressionStore() + + assert store._max_entries == 1000 + assert store._default_ttl == 300 + assert store._enable_feedback is True + assert store._backend is not None + + def test_custom_max_entries(self): + """CompressionStore accepts custom max_entries.""" + store = CompressionStore(max_entries=500) + assert store._max_entries == 500 + + def test_custom_default_ttl(self): + """CompressionStore accepts custom default_ttl.""" + store = CompressionStore(default_ttl=600) + assert store._default_ttl == 600 + + def test_feedback_can_be_disabled(self): + """CompressionStore can disable feedback tracking.""" + store = CompressionStore(enable_feedback=False) + assert store._enable_feedback is False + + def test_custom_backend(self): + """CompressionStore accepts custom backend.""" + mock_backend = MagicMock() + mock_backend.count.return_value = 0 + mock_backend.get.return_value = None + + store = CompressionStore(backend=mock_backend) + assert store._backend is mock_backend + + +# ============================================================================= +# Store Operations Tests +# ============================================================================= + + +class TestCompressionStoreOperations: + """Tests for CompressionStore store operation.""" + + def test_store_returns_24_char_hash(self, store: CompressionStore): + """store() returns a 24 character hash (96 bits for collision resistance).""" + hash_key = store.store( + original="[1,2,3]", + compressed="[1]", + ) + assert len(hash_key) == 24 + assert all(c in "0123456789abcdef" for c in hash_key) + + def test_store_hash_is_deterministic(self, store: CompressionStore): + """Same content produces same hash.""" + content = '{"id": 1, "name": "test"}' + + hash1 = store.store(original=content, compressed="{}") + hash2 = store.store(original=content, compressed="{}") + + assert hash1 == hash2 + + def test_store_hash_based_on_original_content(self, store: CompressionStore): + """Hash is computed from original content, not compressed.""" + original = '{"id": 1}' + compressed1 = '{"id": 1}' + compressed2 = "{}" + + hash1 = store.store(original=original, compressed=compressed1) + hash2 = store.store(original=original, compressed=compressed2) + + assert hash1 == hash2 # Same original = same hash + + def test_store_different_content_different_hash(self, store: CompressionStore): + """Different content produces different hash.""" + hash1 = store.store(original='{"id": 1}', compressed="{}") + hash2 = store.store(original='{"id": 2}', compressed="{}") + + assert hash1 != hash2 + + def test_store_preserves_all_metadata( + self, store: CompressionStore, sample_original: str, sample_compressed: str + ): + """store() preserves all metadata in the entry.""" + hash_key = store.store( + original=sample_original, + compressed=sample_compressed, + original_tokens=1000, + compressed_tokens=100, + original_item_count=100, + compressed_item_count=10, + tool_name="search_api", + tool_call_id="call_123", + query_context="user query", + tool_signature_hash="sig_hash_123", + compression_strategy="top_k", + ttl=600, + ) + + entry = store.retrieve(hash_key) + assert entry is not None + assert entry.original_content == sample_original + assert entry.compressed_content == sample_compressed + assert entry.original_tokens == 1000 + assert entry.compressed_tokens == 100 + assert entry.original_item_count == 100 + assert entry.compressed_item_count == 10 + assert entry.tool_name == "search_api" + assert entry.tool_call_id == "call_123" + assert entry.query_context == "user query" + assert entry.tool_signature_hash == "sig_hash_123" + assert entry.compression_strategy == "top_k" + assert entry.ttl == 600 + + def test_store_uses_default_ttl(self, store: CompressionStore): + """store() uses default TTL when not specified.""" + hash_key = store.store(original="[1]", compressed="[]") + entry = store.retrieve(hash_key) + + assert entry is not None + assert entry.ttl == 300 # Default TTL + + def test_store_accepts_custom_ttl(self, store: CompressionStore): + """store() accepts custom TTL override.""" + hash_key = store.store(original="[1]", compressed="[]", ttl=60) + entry = store.retrieve(hash_key) + + assert entry is not None + assert entry.ttl == 60 + + +# ============================================================================= +# Retrieve Operations Tests +# ============================================================================= + + +class TestCompressionStoreRetrieve: + """Tests for CompressionStore retrieve operation.""" + + def test_retrieve_existing_entry(self, store: CompressionStore): + """retrieve() returns entry for existing hash.""" + hash_key = store.store(original='{"id": 1}', compressed="{}") + + entry = store.retrieve(hash_key) + + assert entry is not None + assert entry.hash == hash_key + assert entry.original_content == '{"id": 1}' + + def test_retrieve_nonexistent_returns_none(self, store: CompressionStore): + """retrieve() returns None for nonexistent hash.""" + entry = store.retrieve("nonexistent_hash_key") + assert entry is None + + def test_retrieve_expired_entry_returns_none(self, store_with_short_ttl: CompressionStore): + """retrieve() returns None for expired entry.""" + hash_key = store_with_short_ttl.store(original="[1]", compressed="[]") + + # Should exist immediately + assert store_with_short_ttl.retrieve(hash_key) is not None + + # Wait for expiration + time.sleep(1.1) + + # Should be None after expiration + assert store_with_short_ttl.retrieve(hash_key) is None + + def test_retrieve_increments_access_count(self, store: CompressionStore): + """retrieve() increments entry access count.""" + hash_key = store.store(original="[1]", compressed="[]") + + store.retrieve(hash_key) + store.retrieve(hash_key) + entry = store.retrieve(hash_key) + + assert entry is not None + assert entry.retrieval_count >= 3 + + def test_retrieve_with_query_tracks_query(self, store: CompressionStore): + """retrieve() with query parameter tracks the query.""" + hash_key = store.store(original="[1]", compressed="[]") + + store.retrieve(hash_key, query="test query") + entry = store.retrieve(hash_key) + + assert entry is not None + assert "test query" in entry.search_queries + + def test_retrieve_returns_copy_not_reference(self, store: CompressionStore): + """retrieve() returns a copy to prevent race conditions.""" + hash_key = store.store(original="[1]", compressed="[]") + + entry1 = store.retrieve(hash_key) + entry2 = store.retrieve(hash_key) + + # Modify the returned entry's mutable field + assert entry1 is not None + assert entry2 is not None + entry1.search_queries.append("modified") + + # Should not affect the other entry + assert "modified" not in entry2.search_queries + + +# ============================================================================= +# TTL Expiration Tests +# ============================================================================= + + +class TestCompressionStoreTTL: + """Tests for CompressionStore TTL expiration behavior.""" + + def test_entry_exists_before_ttl(self, store_with_short_ttl: CompressionStore): + """Entry exists before TTL expires.""" + hash_key = store_with_short_ttl.store(original="[1]", compressed="[]") + assert store_with_short_ttl.exists(hash_key) is True + + def test_entry_not_exists_after_ttl(self, store_with_short_ttl: CompressionStore): + """Entry does not exist after TTL expires.""" + hash_key = store_with_short_ttl.store(original="[1]", compressed="[]") + + time.sleep(1.1) + + assert store_with_short_ttl.exists(hash_key) is False + + def test_get_metadata_returns_none_for_expired(self, store_with_short_ttl: CompressionStore): + """get_metadata returns None for expired entries.""" + hash_key = store_with_short_ttl.store(original="[1]", compressed="[]") + + time.sleep(1.1) + + assert store_with_short_ttl.get_metadata(hash_key) is None + + def test_search_returns_empty_for_expired(self, store_with_short_ttl: CompressionStore): + """search returns empty list for expired entries.""" + hash_key = store_with_short_ttl.store( + original=json.dumps([{"id": 1, "name": "test"}]), + compressed="[]", + ) + + time.sleep(1.1) + + results = store_with_short_ttl.search(hash_key, "test") + assert results == [] + + def test_exists_clean_expired_false_does_not_delete( + self, store_with_short_ttl: CompressionStore + ): + """exists() with clean_expired=False does not delete expired entry.""" + hash_key = store_with_short_ttl.store(original="[1]", compressed="[]") + + time.sleep(1.1) + + # Check exists without cleaning + result = store_with_short_ttl.exists(hash_key, clean_expired=False) + assert result is False + + # Entry should still be in backend (not cleaned yet) + # This is internal behavior - the entry is there but marked expired + + def test_exists_clean_expired_true_deletes(self, store_with_short_ttl: CompressionStore): + """exists() with clean_expired=True deletes expired entry.""" + hash_key = store_with_short_ttl.store(original="[1]", compressed="[]") + + time.sleep(1.1) + + # Check exists with cleaning + result = store_with_short_ttl.exists(hash_key, clean_expired=True) + assert result is False + + +# ============================================================================= +# Eviction Tests +# ============================================================================= + + +class TestCompressionStoreEviction: + """Tests for CompressionStore memory limits and eviction.""" + + def test_eviction_at_capacity(self, store_with_small_capacity: CompressionStore): + """Oldest entries are evicted when at capacity.""" + hashes = [] + for i in range(5): + h = store_with_small_capacity.store( + original=f"content_{i}", + compressed=f"compressed_{i}", + ) + hashes.append(h) + time.sleep(0.01) # Ensure different timestamps + + # Only last 3 should exist (capacity is 3) + assert not store_with_small_capacity.exists(hashes[0]) + assert not store_with_small_capacity.exists(hashes[1]) + assert store_with_small_capacity.exists(hashes[2]) + assert store_with_small_capacity.exists(hashes[3]) + assert store_with_small_capacity.exists(hashes[4]) + + def test_eviction_removes_oldest_first(self, store_with_small_capacity: CompressionStore): + """Eviction removes oldest entries first (heap-based).""" + # Fill to capacity + hashes = [] + for i in range(3): + h = store_with_small_capacity.store( + original=f"content_{i}", + compressed=f"compressed_{i}", + ) + hashes.append(h) + time.sleep(0.01) + + # All 3 should exist + for h in hashes: + assert store_with_small_capacity.exists(h) + + # Add one more - should evict oldest + new_hash = store_with_small_capacity.store( + original="content_new", + compressed="compressed_new", + ) + + # Oldest should be evicted + assert not store_with_small_capacity.exists(hashes[0]) + assert store_with_small_capacity.exists(hashes[1]) + assert store_with_small_capacity.exists(hashes[2]) + assert store_with_small_capacity.exists(new_hash) + + def test_eviction_cleans_expired_first(self): + """Eviction cleans expired entries before evicting valid ones.""" + store = CompressionStore(max_entries=3, default_ttl=1) + + # Add 2 entries that will expire + hash1 = store.store(original="content_1", compressed="c1", ttl=1) + hash2 = store.store(original="content_2", compressed="c2", ttl=1) + + time.sleep(1.1) # Wait for expiration + + # Add 2 more entries (should clean expired first, not evict new) + hash3 = store.store(original="content_3", compressed="c3", ttl=300) + hash4 = store.store(original="content_4", compressed="c4", ttl=300) + + # Expired entries should be gone + assert not store.exists(hash1) + assert not store.exists(hash2) + + # New entries should exist + assert store.exists(hash3) + assert store.exists(hash4) + + def test_heap_rebuild_on_stale_threshold(self): + """Heap is rebuilt when stale entry ratio exceeds threshold.""" + store = CompressionStore(max_entries=10) + + # Store entries and then replace them to create stale heap entries + for i in range(5): + store.store(original=f"content_{i}", compressed=f"c_{i}") + + # Replace all entries (creates stale heap entries) + for i in range(5): + store.store(original=f"content_{i}", compressed=f"updated_{i}") + + # Stale ratio should be tracked + # The heap rebuild happens automatically when threshold is exceeded + + +# ============================================================================= +# Statistics Tests +# ============================================================================= + + +class TestCompressionStoreStats: + """Tests for CompressionStore statistics tracking.""" + + def test_get_stats_entry_count(self, store: CompressionStore): + """get_stats returns correct entry count.""" + store.store(original="[1]", compressed="[]") + store.store(original="[2]", compressed="[]") + + stats = store.get_stats() + assert stats["entry_count"] == 2 + + def test_get_stats_max_entries(self, store: CompressionStore): + """get_stats includes max_entries configuration.""" + stats = store.get_stats() + assert stats["max_entries"] == 1000 + + def test_get_stats_token_totals(self, store: CompressionStore): + """get_stats calculates token totals correctly.""" + store.store( + original="[1]", + compressed="[]", + original_tokens=100, + compressed_tokens=10, + ) + store.store( + original="[2]", + compressed="[]", + original_tokens=200, + compressed_tokens=20, + ) + + stats = store.get_stats() + assert stats["total_original_tokens"] == 300 + assert stats["total_compressed_tokens"] == 30 + + def test_get_stats_retrieval_count(self, store: CompressionStore): + """get_stats tracks total retrievals.""" + hash_key = store.store(original="[1]", compressed="[]") + + store.retrieve(hash_key) + store.retrieve(hash_key) + + stats = store.get_stats() + assert stats["total_retrievals"] >= 2 + + def test_get_stats_event_count(self, store: CompressionStore): + """get_stats includes retrieval event count.""" + hash_key = store.store(original="[1]", compressed="[]") + + store.retrieve(hash_key) + store.retrieve(hash_key) + + stats = store.get_stats() + assert stats["event_count"] >= 2 + + def test_get_stats_includes_backend_stats(self, store: CompressionStore): + """get_stats includes backend-specific stats.""" + store.store(original="[1]", compressed="[]") + + stats = store.get_stats() + assert "backend" in stats + assert stats["backend"]["backend_type"] == "memory" + + +# ============================================================================= +# Get Metadata Tests +# ============================================================================= + + +class TestCompressionStoreMetadata: + """Tests for CompressionStore get_metadata operation.""" + + def test_get_metadata_returns_dict(self, store: CompressionStore): + """get_metadata returns dict with expected fields.""" + hash_key = store.store( + original="[1,2,3]", + compressed="[1]", + tool_name="test_tool", + original_item_count=3, + compressed_item_count=1, + query_context="test query", + ) + + metadata = store.get_metadata(hash_key) + + assert metadata is not None + assert metadata["hash"] == hash_key + assert metadata["tool_name"] == "test_tool" + assert metadata["original_item_count"] == 3 + assert metadata["compressed_item_count"] == 1 + assert metadata["query_context"] == "test query" + assert metadata["compressed_content"] == "[1]" + assert "created_at" in metadata + assert "ttl" in metadata + + def test_get_metadata_nonexistent_returns_none(self, store: CompressionStore): + """get_metadata returns None for nonexistent entry.""" + metadata = store.get_metadata("nonexistent") + assert metadata is None + + +# ============================================================================= +# Search Tests +# ============================================================================= + + +class TestCompressionStoreSearch: + """Tests for CompressionStore search functionality.""" + + def test_search_with_bm25_returns_matches(self, store: CompressionStore): + """search() uses BM25 to find matching items.""" + items = [ + {"id": 1, "content": "Python programming language"}, + {"id": 2, "content": "JavaScript web development"}, + {"id": 3, "content": "Python data science pandas"}, + {"id": 4, "content": "Java enterprise applications"}, + {"id": 5, "content": "Python machine learning tensorflow"}, + ] + + hash_key = store.store( + original=json.dumps(items), + compressed=json.dumps(items[:2]), + ) + + results = store.search(hash_key, "Python programming") + + assert len(results) >= 1 + result_ids = [r["id"] for r in results] + assert 1 in result_ids # "Python programming language" should match + + def test_search_respects_max_results(self, store: CompressionStore): + """search() respects max_results parameter.""" + items = [{"id": i, "content": f"item {i}"} for i in range(50)] + hash_key = store.store(original=json.dumps(items), compressed="[]") + + results = store.search(hash_key, "item", max_results=5) + + assert len(results) <= 5 + + def test_search_respects_score_threshold(self, store: CompressionStore): + """search() filters by score threshold.""" + items = [ + {"id": 1, "content": "exact match query term"}, + {"id": 2, "content": "completely unrelated content xyz"}, + ] + hash_key = store.store(original=json.dumps(items), compressed="[]") + + # High threshold should filter low-scoring items + results = store.search(hash_key, "exact match query", score_threshold=0.5) + + # Should return the exact match, filter the unrelated + if results: + assert any("exact match" in str(r) for r in results) + + def test_search_nonexistent_returns_empty(self, store: CompressionStore): + """search() returns empty list for nonexistent hash.""" + results = store.search("nonexistent", "query") + assert results == [] + + def test_search_invalid_json_returns_empty(self, store: CompressionStore): + """search() handles invalid JSON gracefully.""" + hash_key = store.store(original="not valid json", compressed="[]") + results = store.search(hash_key, "query") + assert results == [] + + def test_search_non_array_returns_empty(self, store: CompressionStore): + """search() returns empty for non-array content.""" + hash_key = store.store(original=json.dumps({"key": "value"}), compressed="{}") + results = store.search(hash_key, "query") + assert results == [] + + def test_search_empty_array_returns_empty(self, store: CompressionStore): + """search() returns empty for empty array.""" + hash_key = store.store(original="[]", compressed="[]") + results = store.search(hash_key, "query") + assert results == [] + + def test_search_logs_retrieval_event(self, store: CompressionStore): + """search() logs retrieval event with search type.""" + items = [{"id": 1, "content": "test"}] + hash_key = store.store(original=json.dumps(items), compressed="[]") + + store.search(hash_key, "test query") + + events = store.get_retrieval_events() + search_events = [e for e in events if e.retrieval_type == "search"] + assert len(search_events) >= 1 + assert search_events[-1].query == "test query" + + +# ============================================================================= +# Retrieval Events Tests +# ============================================================================= + + +class TestCompressionStoreRetrievalEvents: + """Tests for CompressionStore retrieval event tracking.""" + + def test_retrieve_logs_full_event(self, store: CompressionStore): + """retrieve() logs event with 'full' type.""" + hash_key = store.store(original="[1]", compressed="[]", tool_name="test_tool") + + store.retrieve(hash_key) + + events = store.get_retrieval_events() + full_events = [e for e in events if e.retrieval_type == "full"] + assert len(full_events) >= 1 + assert full_events[-1].tool_name == "test_tool" + + def test_get_retrieval_events_limit(self, store: CompressionStore): + """get_retrieval_events respects limit parameter.""" + hash_key = store.store(original="[1]", compressed="[]") + + for _ in range(10): + store.retrieve(hash_key) + + events = store.get_retrieval_events(limit=3) + assert len(events) <= 3 + + def test_get_retrieval_events_filter_by_tool(self, store: CompressionStore): + """get_retrieval_events filters by tool_name.""" + hash1 = store.store(original="[1]", compressed="[]", tool_name="tool_a") + hash2 = store.store(original="[2]", compressed="[]", tool_name="tool_b") + + store.retrieve(hash1) + store.retrieve(hash1) + store.retrieve(hash2) + + tool_a_events = store.get_retrieval_events(tool_name="tool_a") + tool_b_events = store.get_retrieval_events(tool_name="tool_b") + + assert len(tool_a_events) == 2 + assert len(tool_b_events) == 1 + + def test_retrieval_events_include_tool_signature_hash(self, store: CompressionStore): + """Retrieval events include tool_signature_hash for TOIN correlation.""" + hash_key = store.store( + original="[1]", + compressed="[]", + tool_signature_hash="sig_123", + ) + + store.retrieve(hash_key) + + events = store.get_retrieval_events() + assert len(events) >= 1 + assert events[-1].tool_signature_hash == "sig_123" + + +# ============================================================================= +# Edge Cases Tests +# ============================================================================= + + +class TestCompressionStoreEdgeCases: + """Tests for edge cases and error handling.""" + + def test_store_empty_content(self, store: CompressionStore): + """store() handles empty content.""" + hash_key = store.store(original="", compressed="") + + entry = store.retrieve(hash_key) + assert entry is not None + assert entry.original_content == "" + + def test_store_large_content(self, store: CompressionStore): + """store() handles large content.""" + large_content = json.dumps([{"id": i, "data": "x" * 1000} for i in range(100)]) + + hash_key = store.store(original=large_content, compressed="[]") + + entry = store.retrieve(hash_key) + assert entry is not None + assert len(entry.original_content) == len(large_content) + + def test_store_unicode_content(self, store: CompressionStore): + """store() handles unicode content correctly.""" + unicode_content = json.dumps([{"name": "cafe", "emoji": "hello"}]) + + hash_key = store.store(original=unicode_content, compressed="[]") + + entry = store.retrieve(hash_key) + assert entry is not None + assert "cafe" in entry.original_content + + def test_duplicate_store_updates_entry(self, store: CompressionStore): + """Storing same content twice updates the entry.""" + original = '{"id": 1}' + + hash1 = store.store(original=original, compressed="v1") + hash2 = store.store(original=original, compressed="v2") + + assert hash1 == hash2 + + entry = store.retrieve(hash1) + assert entry is not None + # Second store should have updated the entry + assert entry.compressed_content == "v2" + + def test_clear_removes_all_entries(self, store: CompressionStore): + """clear() removes all entries.""" + store.store(original="[1]", compressed="[]") + store.store(original="[2]", compressed="[]") + + store.clear() + + stats = store.get_stats() + assert stats["entry_count"] == 0 + + def test_clear_removes_retrieval_events(self, store: CompressionStore): + """clear() removes retrieval events.""" + hash_key = store.store(original="[1]", compressed="[]") + store.retrieve(hash_key) + + store.clear() + + events = store.get_retrieval_events() + assert len(events) == 0 + + +# ============================================================================= +# Thread Safety Tests +# ============================================================================= + + +class TestCompressionStoreThreadSafety: + """Tests for thread safety.""" + + def test_concurrent_stores(self, store: CompressionStore): + """Concurrent stores don't corrupt data.""" + hashes: list[str] = [] + lock = threading.Lock() + errors: list[str] = [] + + def store_item(i: int) -> None: + try: + h = store.store( + original=f"content_{i}", + compressed=f"compressed_{i}", + ) + with lock: + hashes.append(h) + except Exception as e: + with lock: + errors.append(str(e)) + + threads = [threading.Thread(target=store_item, args=(i,)) for i in range(20)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert errors == [] + assert len(hashes) == 20 + + def test_concurrent_retrieves(self, store: CompressionStore): + """Concurrent retrieves don't corrupt data.""" + hash_key = store.store(original="[1,2,3]", compressed="[1]") + errors: list[str] = [] + results: list[CompressionEntry | None] = [] + lock = threading.Lock() + + def retrieve_item() -> None: + try: + entry = store.retrieve(hash_key) + with lock: + results.append(entry) + except Exception as e: + with lock: + errors.append(str(e)) + + threads = [threading.Thread(target=retrieve_item) for _ in range(20)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert errors == [] + assert len(results) == 20 + for entry in results: + assert entry is not None + assert entry.original_content == "[1,2,3]" + + def test_concurrent_store_and_retrieve(self, store: CompressionStore): + """Concurrent stores and retrieves don't corrupt data.""" + errors: list[str] = [] + + def store_and_retrieve(i: int) -> None: + try: + items = [{"id": j, "batch": i} for j in range(10)] + hash_key = store.store( + original=json.dumps(items), + compressed="[]", + tool_name=f"tool_{i}", + ) + + # Immediately retrieve + entry = store.retrieve(hash_key) + if entry is None: + errors.append(f"Entry {i} not found after store") + elif f'"batch": {i}' not in entry.original_content: + errors.append(f"Entry {i} has wrong content") + except Exception as e: + errors.append(str(e)) + + threads = [threading.Thread(target=store_and_retrieve, args=(i,)) for i in range(20)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert errors == [], f"Errors during concurrent operations: {errors}" + + +# ============================================================================= +# Global Store Singleton Tests +# ============================================================================= + + +class TestGlobalStore: + """Tests for global store singleton pattern.""" + + def test_get_compression_store_returns_singleton(self): + """get_compression_store returns same instance.""" + store1 = get_compression_store() + store2 = get_compression_store() + + assert store1 is store2 + + def test_reset_compression_store_clears_data(self): + """reset_compression_store clears the global store.""" + store = get_compression_store() + store.store(original="[1]", compressed="[]") + + reset_compression_store() + + new_store = get_compression_store() + stats = new_store.get_stats() + assert stats["entry_count"] == 0 + + def test_get_compression_store_uses_params_only_on_first_call(self): + """Parameters are only used on first initialization.""" + reset_compression_store() + + store1 = get_compression_store(max_entries=500, default_ttl=600) + assert store1._max_entries == 500 + assert store1._default_ttl == 600 + + # Second call with different params should return same instance + store2 = get_compression_store(max_entries=100, default_ttl=60) + assert store2 is store1 + assert store2._max_entries == 500 # Original value + + +# ============================================================================= +# Feedback Integration Tests +# ============================================================================= + + +class TestCompressionStoreFeedback: + """Tests for feedback loop integration.""" + + def test_feedback_disabled_no_events(self): + """No events logged when feedback is disabled.""" + store = CompressionStore(enable_feedback=False) + + hash_key = store.store(original="[1]", compressed="[]") + store.retrieve(hash_key) + + # Events should still be tracked internally for the store + # but process_pending_feedback won't forward them + # Verify events are tracked even with feedback disabled + assert store.get_retrieval_events() is not None + + def test_feedback_enabled_logs_events(self): + """Events logged when feedback is enabled.""" + store = CompressionStore(enable_feedback=True) + + hash_key = store.store(original="[1]", compressed="[]", tool_name="test") + store.retrieve(hash_key) + + events = store.get_retrieval_events() + assert len(events) >= 1 + + @patch("headroom.cache.compression_feedback.get_compression_feedback") + @patch("headroom.telemetry.get_telemetry_collector") + @patch("headroom.telemetry.toin.get_toin") + def test_process_pending_feedback_forwards_events( + self, mock_toin, mock_telemetry, mock_feedback + ): + """process_pending_feedback forwards events to feedback systems.""" + mock_fb = MagicMock() + mock_tel = MagicMock() + mock_toin_instance = MagicMock() + + mock_feedback.return_value = mock_fb + mock_telemetry.return_value = mock_tel + mock_toin.return_value = mock_toin_instance + + store = CompressionStore(enable_feedback=True) + + hash_key = store.store( + original="[1]", + compressed="[]", + tool_signature_hash="sig_123", + compression_strategy="top_k", + ) + store.retrieve(hash_key) + + # Feedback should have been called + assert mock_fb.record_retrieval.called + + def test_eviction_success_creates_event(self): + """Eviction without retrieval creates success event.""" + store = CompressionStore(max_entries=2, enable_feedback=True) + + # Store entries with signature hash for eviction tracking + store.store( + original="content_0", + compressed="c0", + tool_signature_hash="sig_0", + compression_strategy="top_k", + ) + time.sleep(0.01) + + store.store( + original="content_1", + compressed="c1", + tool_signature_hash="sig_1", + compression_strategy="top_k", + ) + time.sleep(0.01) + + # This should trigger eviction of first entry + store.store( + original="content_2", + compressed="c2", + tool_signature_hash="sig_2", + compression_strategy="top_k", + ) + + # The evicted entry (content_0) was never retrieved, + # so an eviction_success event should be queued + # (tested via the pending_feedback mechanism) + + +# ============================================================================= +# RetrievalEvent Tests +# ============================================================================= + + +class TestRetrievalEvent: + """Tests for RetrievalEvent dataclass.""" + + def test_retrieval_event_creation(self): + """RetrievalEvent can be created with all fields.""" + event = RetrievalEvent( + hash="abc123", + query="test query", + items_retrieved=5, + total_items=100, + tool_name="search_api", + timestamp=time.time(), + retrieval_type="search", + tool_signature_hash="sig_123", + ) + + assert event.hash == "abc123" + assert event.query == "test query" + assert event.items_retrieved == 5 + assert event.total_items == 100 + assert event.tool_name == "search_api" + assert event.retrieval_type == "search" + assert event.tool_signature_hash == "sig_123" + + def test_retrieval_event_default_signature_hash(self): + """RetrievalEvent has None default for tool_signature_hash.""" + event = RetrievalEvent( + hash="abc123", + query=None, + items_retrieved=10, + total_items=10, + tool_name="test", + timestamp=time.time(), + retrieval_type="full", + ) + + assert event.tool_signature_hash is None + + +# ============================================================================= +# Hash Collision Detection Tests +# ============================================================================= + + +class TestHashCollisionDetection: + """Tests for hash collision detection and handling.""" + + def test_same_content_no_collision_warning( + self, store: CompressionStore, caplog: pytest.LogCaptureFixture + ): + """Same content stored twice should not warn about collision.""" + import logging + + with caplog.at_level(logging.WARNING): + store.store(original="[1,2,3]", compressed="[1]") + store.store(original="[1,2,3]", compressed="[1,2]") + + # Should not have collision warning + assert "Hash collision detected" not in caplog.text + + def test_hash_uses_sha256_truncated(self, store: CompressionStore): + """Hash is SHA256 truncated to 24 characters.""" + content = "test content" + expected_hash = hashlib.sha256(content.encode()).hexdigest()[:24] + + hash_key = store.store(original=content, compressed="[]") + + assert hash_key == expected_hash diff --git a/tests/test_integrations/langchain/test_agents.py b/tests/test_integrations/langchain/test_agents.py new file mode 100644 index 000000000..94b5c5506 --- /dev/null +++ b/tests/test_integrations/langchain/test_agents.py @@ -0,0 +1,543 @@ +"""Tests for LangChain agent tool integration. + +Tests cover: +1. ToolCompressionMetrics - Dataclass for tool compression metrics +2. ToolMetricsCollector - Collector for compression metrics +3. HeadroomToolWrapper - Wrapper for LangChain tools with compression +4. wrap_tools_with_headroom - Convenience function for wrapping multiple tools +5. get_tool_metrics / reset_tool_metrics - Global metrics access +""" + +from datetime import datetime +from unittest.mock import MagicMock, patch + +import pytest + +# Check if LangChain is available +try: + from langchain_core.tools import BaseTool, StructuredTool + + LANGCHAIN_AVAILABLE = True +except ImportError: + LANGCHAIN_AVAILABLE = False + +# Skip all tests if LangChain not installed +pytestmark = pytest.mark.skipif(not LANGCHAIN_AVAILABLE, reason="LangChain not installed") + + +@pytest.fixture +def mock_tool(): + """Create a mock LangChain tool.""" + mock = MagicMock(spec=BaseTool) + mock.name = "test_tool" + mock.description = "A test tool" + mock.invoke = MagicMock(return_value="Tool result") + return mock + + +@pytest.fixture +def mock_tool_with_large_output(): + """Create a mock tool that returns large output.""" + mock = MagicMock(spec=BaseTool) + mock.name = "search_tool" + mock.description = "Search tool with large results" + # Return > 1000 chars to trigger compression + large_output = '{"items": [' + ",".join(f'{{"id": {i}}}' for i in range(200)) + "]}" + mock.invoke = MagicMock(return_value=large_output) + return mock + + +class TestToolCompressionMetrics: + """Tests for ToolCompressionMetrics dataclass.""" + + def test_create_metrics(self): + """Create metrics with all fields.""" + from headroom.integrations.langchain.agents import ToolCompressionMetrics + + metrics = ToolCompressionMetrics( + tool_name="search", + timestamp=datetime.now(), + chars_before=5000, + chars_after=2000, + chars_saved=3000, + compression_ratio=0.4, + was_compressed=True, + ) + + assert metrics.tool_name == "search" + assert metrics.chars_before == 5000 + assert metrics.chars_after == 2000 + assert metrics.chars_saved == 3000 + assert metrics.compression_ratio == 0.4 + assert metrics.was_compressed is True + + def test_metrics_defaults(self): + """Verify no default values (all required).""" + from headroom.integrations.langchain.agents import ToolCompressionMetrics + + # All fields are required, should raise TypeError if missing + with pytest.raises(TypeError): + ToolCompressionMetrics() # type: ignore[call-arg] + + +class TestToolMetricsCollector: + """Tests for ToolMetricsCollector.""" + + def test_init_empty(self): + """Initialize with empty metrics list.""" + from headroom.integrations.langchain.agents import ToolMetricsCollector + + collector = ToolMetricsCollector() + + assert collector.metrics == [] + + def test_add_metric(self): + """Add a metric to the collector.""" + from headroom.integrations.langchain.agents import ( + ToolCompressionMetrics, + ToolMetricsCollector, + ) + + collector = ToolMetricsCollector() + metric = ToolCompressionMetrics( + tool_name="test", + timestamp=datetime.now(), + chars_before=100, + chars_after=80, + chars_saved=20, + compression_ratio=0.8, + was_compressed=True, + ) + + collector.add(metric) + + assert len(collector.metrics) == 1 + assert collector.metrics[0] is metric + + def test_add_metric_limits_to_1000(self): + """Metrics list is limited to 1000 entries.""" + from headroom.integrations.langchain.agents import ( + ToolCompressionMetrics, + ToolMetricsCollector, + ) + + collector = ToolMetricsCollector() + + # Add 1100 metrics + for i in range(1100): + metric = ToolCompressionMetrics( + tool_name=f"tool_{i}", + timestamp=datetime.now(), + chars_before=100, + chars_after=80, + chars_saved=20, + compression_ratio=0.8, + was_compressed=True, + ) + collector.add(metric) + + assert len(collector.metrics) == 1000 + # Should keep the last 1000 (most recent) + assert collector.metrics[0].tool_name == "tool_100" + assert collector.metrics[-1].tool_name == "tool_1099" + + def test_get_summary_empty(self): + """Get summary with no metrics.""" + from headroom.integrations.langchain.agents import ToolMetricsCollector + + collector = ToolMetricsCollector() + summary = collector.get_summary() + + assert summary["total_invocations"] == 0 + assert summary["total_compressions"] == 0 + assert summary["total_chars_saved"] == 0 + + def test_get_summary_with_data(self): + """Get summary with metrics.""" + from headroom.integrations.langchain.agents import ( + ToolCompressionMetrics, + ToolMetricsCollector, + ) + + collector = ToolMetricsCollector() + + # Add compressed metric + collector.add( + ToolCompressionMetrics( + tool_name="search", + timestamp=datetime.now(), + chars_before=5000, + chars_after=2000, + chars_saved=3000, + compression_ratio=0.4, + was_compressed=True, + ) + ) + + # Add uncompressed metric + collector.add( + ToolCompressionMetrics( + tool_name="simple", + timestamp=datetime.now(), + chars_before=100, + chars_after=100, + chars_saved=0, + compression_ratio=1.0, + was_compressed=False, + ) + ) + + summary = collector.get_summary() + + assert summary["total_invocations"] == 2 + assert summary["total_compressions"] == 1 + assert summary["total_chars_saved"] == 3000 + assert summary["average_compression_ratio"] == 0.4 # Only compressed + + def test_get_summary_by_tool(self): + """Get per-tool statistics.""" + from headroom.integrations.langchain.agents import ( + ToolCompressionMetrics, + ToolMetricsCollector, + ) + + collector = ToolMetricsCollector() + + # Add metrics for different tools + for _i in range(3): + collector.add( + ToolCompressionMetrics( + tool_name="search", + timestamp=datetime.now(), + chars_before=1000, + chars_after=500, + chars_saved=500, + compression_ratio=0.5, + was_compressed=True, + ) + ) + + for _i in range(2): + collector.add( + ToolCompressionMetrics( + tool_name="database", + timestamp=datetime.now(), + chars_before=100, + chars_after=100, + chars_saved=0, + compression_ratio=1.0, + was_compressed=False, + ) + ) + + summary = collector.get_summary() + + assert "by_tool" in summary + assert summary["by_tool"]["search"]["invocations"] == 3 + assert summary["by_tool"]["search"]["compressions"] == 3 + assert summary["by_tool"]["search"]["chars_saved"] == 1500 + assert summary["by_tool"]["database"]["invocations"] == 2 + assert summary["by_tool"]["database"]["compressions"] == 0 + + +class TestHeadroomToolWrapper: + """Tests for HeadroomToolWrapper.""" + + def test_init_defaults(self, mock_tool): + """Initialize with default settings.""" + from headroom.integrations.langchain.agents import HeadroomToolWrapper + + wrapper = HeadroomToolWrapper(mock_tool) + + assert wrapper.tool is mock_tool + assert wrapper.name == "test_tool" + assert wrapper.description == "A test tool" + assert wrapper.min_chars_to_compress == 1000 + + def test_init_custom_threshold(self, mock_tool): + """Initialize with custom compression threshold.""" + from headroom.integrations.langchain.agents import ( + HeadroomToolWrapper, + ToolMetricsCollector, + ) + + collector = ToolMetricsCollector() + wrapper = HeadroomToolWrapper( + mock_tool, + min_chars_to_compress=500, + metrics_collector=collector, + ) + + assert wrapper.min_chars_to_compress == 500 + assert wrapper._metrics is collector + + def test_call_small_output_no_compression(self, mock_tool): + """Small outputs are not compressed.""" + from headroom.integrations.langchain.agents import ( + HeadroomToolWrapper, + ToolMetricsCollector, + ) + + collector = ToolMetricsCollector() + wrapper = HeadroomToolWrapper( + mock_tool, + min_chars_to_compress=1000, + metrics_collector=collector, + ) + + result = wrapper("input") + + assert result == "Tool result" + assert len(collector.metrics) == 1 + assert collector.metrics[0].was_compressed is False + + def test_call_large_output_triggers_compression(self, mock_tool_with_large_output): + """Large outputs trigger compression.""" + from headroom.integrations.langchain.agents import ( + HeadroomToolWrapper, + ToolMetricsCollector, + ) + + collector = ToolMetricsCollector() + wrapper = HeadroomToolWrapper( + mock_tool_with_large_output, + min_chars_to_compress=100, + metrics_collector=collector, + ) + + # Mock compress_tool_result to return compressed output + with patch("headroom.integrations.langchain.agents.compress_tool_result") as mock_compress: + mock_compress.return_value = '{"items": [...compressed...]}' + wrapper("query") + + mock_compress.assert_called_once() + assert len(collector.metrics) == 1 + assert collector.metrics[0].was_compressed is True + + def test_call_converts_non_string_result(self, mock_tool): + """Non-string results are converted to strings.""" + from headroom.integrations.langchain.agents import HeadroomToolWrapper + + mock_tool.invoke.return_value = {"key": "value"} + + wrapper = HeadroomToolWrapper(mock_tool) + result = wrapper("input") + + assert isinstance(result, str) + assert "key" in result + + def test_invoke_alias(self, mock_tool): + """invoke() is an alias for __call__().""" + from headroom.integrations.langchain.agents import HeadroomToolWrapper + + wrapper = HeadroomToolWrapper(mock_tool) + + result1 = wrapper("input") + mock_tool.invoke.reset_mock() + result2 = wrapper.invoke("input") + + assert result1 == result2 + + def test_compression_failure_returns_original(self, mock_tool_with_large_output): + """Compression failure returns original output.""" + from headroom.integrations.langchain.agents import HeadroomToolWrapper + + wrapper = HeadroomToolWrapper( + mock_tool_with_large_output, + min_chars_to_compress=100, + ) + + with patch("headroom.integrations.langchain.agents.compress_tool_result") as mock_compress: + mock_compress.side_effect = Exception("Compression error") + result = wrapper("query") + + # Should return original output + assert "items" in result + assert "id" in result + + def test_as_langchain_tool(self, mock_tool): + """Convert wrapper to LangChain StructuredTool.""" + from headroom.integrations.langchain.agents import HeadroomToolWrapper + + wrapper = HeadroomToolWrapper(mock_tool) + lc_tool = wrapper.as_langchain_tool() + + assert isinstance(lc_tool, StructuredTool) + assert lc_tool.name == "test_tool" + assert lc_tool.description == "A test tool" + + def test_metrics_recorded_correctly(self, mock_tool_with_large_output): + """Verify metrics are recorded correctly.""" + from headroom.integrations.langchain.agents import ( + HeadroomToolWrapper, + ToolMetricsCollector, + ) + + collector = ToolMetricsCollector() + wrapper = HeadroomToolWrapper( + mock_tool_with_large_output, + min_chars_to_compress=100, + metrics_collector=collector, + ) + + original_len = len(mock_tool_with_large_output.invoke.return_value) + + with patch("headroom.integrations.langchain.agents.compress_tool_result") as mock_compress: + compressed_result = '{"items": [...]}' + mock_compress.return_value = compressed_result + wrapper("query") + + metric = collector.metrics[0] + assert metric.tool_name == "search_tool" + assert metric.chars_before == original_len + assert metric.chars_after == len(compressed_result) + assert metric.chars_saved == original_len - len(compressed_result) + + +class TestWrapToolsWithHeadroom: + """Tests for wrap_tools_with_headroom function.""" + + def test_wrap_single_tool(self, mock_tool): + """Wrap a single tool.""" + from headroom.integrations.langchain.agents import wrap_tools_with_headroom + + wrapped = wrap_tools_with_headroom([mock_tool]) + + assert len(wrapped) == 1 + assert isinstance(wrapped[0], StructuredTool) + assert wrapped[0].name == "test_tool" + + def test_wrap_multiple_tools(self, mock_tool): + """Wrap multiple tools.""" + from headroom.integrations.langchain.agents import wrap_tools_with_headroom + + tool2 = MagicMock(spec=BaseTool) + tool2.name = "tool_2" + tool2.description = "Second tool" + tool2.invoke = MagicMock(return_value="Result 2") + + wrapped = wrap_tools_with_headroom([mock_tool, tool2]) + + assert len(wrapped) == 2 + assert wrapped[0].name == "test_tool" + assert wrapped[1].name == "tool_2" + + def test_wrap_with_custom_threshold(self, mock_tool): + """Wrap with custom compression threshold.""" + from headroom.integrations.langchain.agents import wrap_tools_with_headroom + + wrapped = wrap_tools_with_headroom([mock_tool], min_chars_to_compress=500) + + assert len(wrapped) == 1 + # Invoke to verify wrapper is configured + # The wrapper should be invoked through the StructuredTool + assert wrapped[0].name == "test_tool" + + def test_wrap_with_shared_collector(self, mock_tool): + """Wrap with shared metrics collector.""" + from headroom.integrations.langchain.agents import ( + ToolMetricsCollector, + wrap_tools_with_headroom, + ) + + collector = ToolMetricsCollector() + + tool2 = MagicMock(spec=BaseTool) + tool2.name = "tool_2" + tool2.description = "Second tool" + tool2.invoke = MagicMock(return_value="Result 2") + + wrapped = wrap_tools_with_headroom( + [mock_tool, tool2], + metrics_collector=collector, + ) + + # Invoke both tools + wrapped[0].func("input1") + wrapped[1].func("input2") + + # Both should use the same collector + assert len(collector.metrics) == 2 + + def test_wrap_empty_list(self): + """Wrap empty list returns empty list.""" + from headroom.integrations.langchain.agents import wrap_tools_with_headroom + + wrapped = wrap_tools_with_headroom([]) + + assert wrapped == [] + + +class TestGlobalMetrics: + """Tests for global metrics functions.""" + + def test_get_tool_metrics(self): + """get_tool_metrics returns the global collector.""" + from headroom.integrations.langchain.agents import ( + ToolMetricsCollector, + get_tool_metrics, + ) + + collector = get_tool_metrics() + + assert isinstance(collector, ToolMetricsCollector) + + def test_reset_tool_metrics(self): + """reset_tool_metrics creates new collector.""" + from headroom.integrations.langchain.agents import ( + ToolCompressionMetrics, + get_tool_metrics, + reset_tool_metrics, + ) + + # Add a metric to the global collector + collector = get_tool_metrics() + collector.add( + ToolCompressionMetrics( + tool_name="test", + timestamp=datetime.now(), + chars_before=100, + chars_after=100, + chars_saved=0, + compression_ratio=1.0, + was_compressed=False, + ) + ) + + # Reset + reset_tool_metrics() + + # New collector should be empty + new_collector = get_tool_metrics() + assert len(new_collector.metrics) == 0 + + def test_wrapper_uses_global_metrics_by_default(self, mock_tool): + """HeadroomToolWrapper uses global metrics by default.""" + from headroom.integrations.langchain.agents import ( + HeadroomToolWrapper, + get_tool_metrics, + reset_tool_metrics, + ) + + # Reset to start fresh + reset_tool_metrics() + + wrapper = HeadroomToolWrapper(mock_tool) + wrapper("input") + + global_collector = get_tool_metrics() + assert len(global_collector.metrics) == 1 + + +class TestLangChainNotAvailable: + """Tests for behavior when LangChain is not available.""" + + def test_check_raises_import_error(self): + """_check_langchain_available raises ImportError when not available.""" + from headroom.integrations.langchain.agents import _check_langchain_available + + # When LangChain IS available, should not raise + try: + _check_langchain_available() + except ImportError: + pytest.fail("Should not raise when LangChain is available") diff --git a/tests/test_integrations/langchain/test_memory.py b/tests/test_integrations/langchain/test_memory.py new file mode 100644 index 000000000..5699e9ab3 --- /dev/null +++ b/tests/test_integrations/langchain/test_memory.py @@ -0,0 +1,499 @@ +"""Tests for LangChain memory integration with automatic compression. + +Tests cover: +1. HeadroomChatMessageHistory - Wrapper for chat message history with compression +2. Message conversion to/from OpenAI format +3. Rolling window compression behavior +4. Token counting and threshold detection +5. Compression statistics tracking +""" + +from unittest.mock import MagicMock, patch + +import pytest + +# Check if LangChain is available +try: + from langchain_core.messages import ( + AIMessage, + BaseMessage, + HumanMessage, + SystemMessage, + ToolMessage, + ) + + LANGCHAIN_AVAILABLE = True +except ImportError: + LANGCHAIN_AVAILABLE = False + +# Skip all tests if LangChain not installed +pytestmark = pytest.mark.skipif(not LANGCHAIN_AVAILABLE, reason="LangChain not installed") + + +@pytest.fixture +def mock_base_history(): + """Create a mock BaseChatMessageHistory.""" + mock = MagicMock() + mock.messages = [] + return mock + + +@pytest.fixture +def mock_provider(): + """Create a mock provider with token counter.""" + mock = MagicMock() + mock_counter = MagicMock() + mock_counter.count_text = MagicMock(side_effect=lambda text: len(text.split())) + mock.get_token_counter = MagicMock(return_value=mock_counter) + return mock + + +@pytest.fixture +def sample_langchain_messages(): + """Sample LangChain messages for testing.""" + return [ + SystemMessage(content="You are a helpful assistant."), + HumanMessage(content="Hello, how are you?"), + AIMessage(content="I am doing well, thank you!"), + HumanMessage(content="What is the weather today?"), + AIMessage(content="I don't have access to weather data."), + ] + + +class TestHeadroomChatMessageHistoryInit: + """Tests for HeadroomChatMessageHistory initialization.""" + + def test_init_defaults(self, mock_base_history): + """Initialize with default settings.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + with patch("headroom.integrations.langchain.memory.OpenAIProvider"): + history = HeadroomChatMessageHistory(mock_base_history) + + assert history._base is mock_base_history + assert history._threshold == 4000 + assert history._keep_recent_turns == 5 + assert history._model == "gpt-4o" + assert history._compression_count == 0 + assert history._total_tokens_saved == 0 + + def test_init_custom_threshold(self, mock_base_history, mock_provider): + """Initialize with custom compression threshold.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory( + mock_base_history, + compress_threshold_tokens=8000, + keep_recent_turns=10, + model="gpt-4-turbo", + provider=mock_provider, + ) + + assert history._threshold == 8000 + assert history._keep_recent_turns == 10 + assert history._model == "gpt-4-turbo" + assert history._provider is mock_provider + + +class TestHeadroomChatMessageHistoryMessages: + """Tests for message access and compression.""" + + def test_messages_returns_empty_when_no_messages(self, mock_base_history, mock_provider): + """messages property returns empty list when no messages.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + mock_base_history.messages = [] + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + messages = history.messages + + assert messages == [] + + def test_messages_returns_uncompressed_when_below_threshold( + self, mock_base_history, mock_provider, sample_langchain_messages + ): + """messages returns uncompressed when below token threshold.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + mock_base_history.messages = sample_langchain_messages + + history = HeadroomChatMessageHistory( + mock_base_history, + compress_threshold_tokens=10000, # High threshold + provider=mock_provider, + ) + + messages = history.messages + + # Should return all messages unchanged + assert len(messages) == len(sample_langchain_messages) + assert history._compression_count == 0 + + def test_messages_compresses_when_over_threshold(self, mock_base_history, mock_provider): + """messages applies compression when over token threshold.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + # Create messages that exceed threshold + mock_base_history.messages = [ + SystemMessage(content="System " * 100), + HumanMessage(content="User " * 100), + AIMessage(content="Assistant " * 100), + ] + + history = HeadroomChatMessageHistory( + mock_base_history, + compress_threshold_tokens=10, # Very low threshold + provider=mock_provider, + ) + + # Mock _apply_rolling_window to return fewer messages + with patch.object(history, "_apply_rolling_window") as mock_apply: + mock_apply.return_value = [ + SystemMessage(content="Compressed"), + ] + + _ = history.messages + + mock_apply.assert_called_once() + assert history._compression_count == 1 + + def test_messages_tracks_tokens_saved(self, mock_base_history, mock_provider): + """Compression tracks tokens saved.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + # Create messages that exceed threshold + mock_base_history.messages = [ + SystemMessage(content="Word " * 50), + HumanMessage(content="Word " * 50), + ] + + history = HeadroomChatMessageHistory( + mock_base_history, + compress_threshold_tokens=10, # Very low threshold + provider=mock_provider, + ) + + # Mock _apply_rolling_window to return fewer messages + with patch.object(history, "_apply_rolling_window") as mock_apply: + mock_apply.return_value = [ + SystemMessage(content="Short"), + ] + + _ = history.messages + + # tokens_saved should increase + assert history._total_tokens_saved > 0 + + +class TestHeadroomChatMessageHistoryAddMessage: + """Tests for add_message methods.""" + + def test_add_message(self, mock_base_history, mock_provider): + """add_message delegates to base history.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + msg = HumanMessage(content="Hello") + history.add_message(msg) + + mock_base_history.add_message.assert_called_once_with(msg) + + def test_add_user_message(self, mock_base_history, mock_provider): + """add_user_message delegates to base history.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + history.add_user_message("Hello") + + mock_base_history.add_user_message.assert_called_once_with("Hello") + + def test_add_ai_message(self, mock_base_history, mock_provider): + """add_ai_message delegates to base history.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + history.add_ai_message("Response") + + mock_base_history.add_ai_message.assert_called_once_with("Response") + + def test_clear(self, mock_base_history, mock_provider): + """clear delegates to base history.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + history.clear() + + mock_base_history.clear.assert_called_once() + + +class TestHeadroomChatMessageHistoryConversion: + """Tests for message format conversion.""" + + def test_convert_to_openai_system_message(self, mock_base_history, mock_provider): + """Convert SystemMessage to OpenAI format.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + messages = [SystemMessage(content="You are helpful.")] + result = history._convert_to_openai(messages) + + assert len(result) == 1 + assert result[0]["role"] == "system" + assert result[0]["content"] == "You are helpful." + + def test_convert_to_openai_human_message(self, mock_base_history, mock_provider): + """Convert HumanMessage to OpenAI format.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + messages = [HumanMessage(content="Hello")] + result = history._convert_to_openai(messages) + + assert result[0]["role"] == "user" + assert result[0]["content"] == "Hello" + + def test_convert_to_openai_ai_message(self, mock_base_history, mock_provider): + """Convert AIMessage to OpenAI format.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + messages = [AIMessage(content="I can help.")] + result = history._convert_to_openai(messages) + + assert result[0]["role"] == "assistant" + assert result[0]["content"] == "I can help." + + def test_convert_to_openai_ai_message_with_tool_calls(self, mock_base_history, mock_provider): + """Convert AIMessage with tool_calls to OpenAI format.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + messages = [ + AIMessage( + content="Calling tool...", + tool_calls=[{"id": "call_1", "name": "search", "args": {"q": "test"}}], + ) + ] + result = history._convert_to_openai(messages) + + assert result[0]["role"] == "assistant" + assert "tool_calls" in result[0] + assert result[0]["tool_calls"][0]["id"] == "call_1" + + def test_convert_to_openai_tool_message(self, mock_base_history, mock_provider): + """Convert ToolMessage to OpenAI format.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + messages = [ToolMessage(content='{"result": "data"}', tool_call_id="call_1")] + result = history._convert_to_openai(messages) + + assert result[0]["role"] == "tool" + assert result[0]["tool_call_id"] == "call_1" + assert result[0]["content"] == '{"result": "data"}' + + def test_convert_from_openai_system(self, mock_base_history, mock_provider): + """Convert OpenAI system message back to LangChain.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + openai_msgs = [{"role": "system", "content": "System prompt"}] + result = history._convert_from_openai(openai_msgs) + + assert len(result) == 1 + assert isinstance(result[0], SystemMessage) + assert result[0].content == "System prompt" + + def test_convert_from_openai_user(self, mock_base_history, mock_provider): + """Convert OpenAI user message back to LangChain.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + openai_msgs = [{"role": "user", "content": "Hello"}] + result = history._convert_from_openai(openai_msgs) + + assert isinstance(result[0], HumanMessage) + assert result[0].content == "Hello" + + def test_convert_from_openai_assistant(self, mock_base_history, mock_provider): + """Convert OpenAI assistant message back to LangChain.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + openai_msgs = [{"role": "assistant", "content": "Response"}] + result = history._convert_from_openai(openai_msgs) + + assert isinstance(result[0], AIMessage) + assert result[0].content == "Response" + + def test_convert_from_openai_assistant_with_tool_calls(self, mock_base_history, mock_provider): + """Convert OpenAI assistant message with tool_calls back to LangChain.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + openai_msgs = [ + { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "call_1", "name": "search", "args": {}}], + } + ] + result = history._convert_from_openai(openai_msgs) + + assert isinstance(result[0], AIMessage) + # LangChain may add a 'type' field to tool_calls, so just check key fields + assert len(result[0].tool_calls) == 1 + assert result[0].tool_calls[0]["id"] == "call_1" + assert result[0].tool_calls[0]["name"] == "search" + assert result[0].tool_calls[0]["args"] == {} + + def test_convert_from_openai_tool(self, mock_base_history, mock_provider): + """Convert OpenAI tool message back to LangChain.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory(mock_base_history, provider=mock_provider) + + openai_msgs = [{"role": "tool", "tool_call_id": "call_1", "content": '{"data": 1}'}] + result = history._convert_from_openai(openai_msgs) + + assert isinstance(result[0], ToolMessage) + assert result[0].tool_call_id == "call_1" + assert result[0].content == '{"data": 1}' + + +class TestHeadroomChatMessageHistoryTokenCounting: + """Tests for token counting.""" + + def test_count_tokens(self, mock_base_history, mock_provider): + """Count tokens using provider's tokenizer.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory( + mock_base_history, + provider=mock_provider, + model="gpt-4o", + ) + + messages = [ + HumanMessage(content="Hello world"), + AIMessage(content="Hi there"), + ] + + count = history._count_tokens(messages) + + # Mock counts words, so "Hello world" = 2, "Hi there" = 2 + assert count == 4 + mock_provider.get_token_counter.assert_called_with("gpt-4o") + + +class TestHeadroomChatMessageHistoryStats: + """Tests for compression statistics.""" + + def test_get_compression_stats_initial(self, mock_base_history, mock_provider): + """Get initial compression stats.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory( + mock_base_history, + compress_threshold_tokens=4000, + keep_recent_turns=5, + provider=mock_provider, + ) + + stats = history.get_compression_stats() + + assert stats["compression_count"] == 0 + assert stats["total_tokens_saved"] == 0 + assert stats["threshold_tokens"] == 4000 + assert stats["keep_recent_turns"] == 5 + + def test_get_compression_stats_after_compression(self, mock_base_history, mock_provider): + """Get compression stats after compression.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + mock_base_history.messages = [ + SystemMessage(content="Word " * 100), + HumanMessage(content="Word " * 100), + ] + + history = HeadroomChatMessageHistory( + mock_base_history, + compress_threshold_tokens=10, + provider=mock_provider, + ) + + # Mock _apply_rolling_window + with patch.object(history, "_apply_rolling_window") as mock_apply: + mock_apply.return_value = [SystemMessage(content="Short")] + + _ = history.messages + + stats = history.get_compression_stats() + + assert stats["compression_count"] == 1 + assert stats["total_tokens_saved"] > 0 + + +class TestHeadroomChatMessageHistoryRollingWindow: + """Tests for rolling window compression.""" + + def test_apply_rolling_window_calls_pipeline(self, mock_base_history, mock_provider): + """_apply_rolling_window uses TransformPipeline.""" + from headroom.integrations.langchain.memory import HeadroomChatMessageHistory + + history = HeadroomChatMessageHistory( + mock_base_history, + compress_threshold_tokens=1000, + keep_recent_turns=5, + provider=mock_provider, + ) + + messages = [ + HumanMessage(content="Hello"), + AIMessage(content="Hi there"), + ] + + with patch("headroom.integrations.langchain.memory.TransformPipeline") as MockPipeline: + mock_instance = MagicMock() + mock_result = MagicMock() + mock_result.messages = [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there"}, + ] + mock_instance.apply.return_value = mock_result + MockPipeline.return_value = mock_instance + + result = history._apply_rolling_window(messages) + + MockPipeline.assert_called_once() + mock_instance.apply.assert_called_once() + + # Result should be converted back to LangChain messages + assert all(isinstance(m, BaseMessage) for m in result) + + +class TestLangChainNotAvailable: + """Tests for behavior when LangChain is not available.""" + + def test_check_raises_import_error(self): + """_check_langchain_available raises ImportError when not available.""" + from headroom.integrations.langchain.memory import _check_langchain_available + + # When LangChain IS available, should not raise + try: + _check_langchain_available() + except ImportError: + pytest.fail("Should not raise when LangChain is available") diff --git a/tests/test_integrations/langchain/test_retriever.py b/tests/test_integrations/langchain/test_retriever.py new file mode 100644 index 000000000..40161f399 --- /dev/null +++ b/tests/test_integrations/langchain/test_retriever.py @@ -0,0 +1,493 @@ +"""Tests for LangChain retriever integration with document compression. + +Tests cover: +1. CompressionMetrics - Dataclass for document compression metrics +2. HeadroomDocumentCompressor - LangChain BaseDocumentCompressor implementation +3. BM25-style relevance scoring +4. Diverse document selection (MMR-style) +5. Compression statistics tracking +""" + +from unittest.mock import MagicMock + +import pytest + +# Check if LangChain is available +try: + from langchain_core.documents import Document + + LANGCHAIN_AVAILABLE = True +except ImportError: + LANGCHAIN_AVAILABLE = False + +# Skip all tests if LangChain not installed +pytestmark = pytest.mark.skipif(not LANGCHAIN_AVAILABLE, reason="LangChain not installed") + + +@pytest.fixture +def sample_documents(): + """Create sample documents for testing.""" + return [ + Document(page_content="Python is a programming language.", metadata={"id": 1}), + Document(page_content="Python is great for data science.", metadata={"id": 2}), + Document(page_content="Java is also a programming language.", metadata={"id": 3}), + Document( + page_content="Machine learning uses Python extensively.", + metadata={"id": 4}, + ), + Document(page_content="JavaScript is used for web development.", metadata={"id": 5}), + ] + + +@pytest.fixture +def many_documents(): + """Create many documents for compression testing.""" + return [ + Document( + page_content=f"Document {i} contains some text about topic {i % 5}.", + metadata={"id": i}, + ) + for i in range(50) + ] + + +class TestCompressionMetrics: + """Tests for CompressionMetrics dataclass.""" + + def test_create_metrics(self): + """Create compression metrics with all fields.""" + from headroom.integrations.langchain.retriever import CompressionMetrics + + metrics = CompressionMetrics( + documents_before=50, + documents_after=10, + documents_removed=40, + relevance_scores=[0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.15, 0.1], + ) + + assert metrics.documents_before == 50 + assert metrics.documents_after == 10 + assert metrics.documents_removed == 40 + assert len(metrics.relevance_scores) == 10 + + def test_metrics_required_fields(self): + """All fields are required.""" + from headroom.integrations.langchain.retriever import CompressionMetrics + + with pytest.raises(TypeError): + CompressionMetrics() # type: ignore[call-arg] + + +class TestHeadroomDocumentCompressorInit: + """Tests for HeadroomDocumentCompressor initialization.""" + + def test_init_defaults(self): + """Initialize with default settings.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + assert compressor.max_documents == 10 + assert compressor.min_relevance == 0.0 + assert compressor.prefer_diverse is False + assert compressor._last_metrics is None + + def test_init_custom_settings(self): + """Initialize with custom settings.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor( + max_documents=20, + min_relevance=0.5, + prefer_diverse=True, + ) + + assert compressor.max_documents == 20 + assert compressor.min_relevance == 0.5 + assert compressor.prefer_diverse is True + + +class TestHeadroomDocumentCompressorCompress: + """Tests for compress_documents method.""" + + def test_compress_empty_documents(self): + """Compress empty list returns empty list.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + result = compressor.compress_documents([], "query") + + assert result == [] + assert compressor._last_metrics is not None + assert compressor._last_metrics.documents_before == 0 + + def test_compress_fewer_than_max_documents(self, sample_documents): + """Compress when documents fewer than max returns all.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor(max_documents=10) # More than 5 docs + + result = compressor.compress_documents(sample_documents, "Python") + + assert len(result) == len(sample_documents) + assert compressor._last_metrics.documents_removed == 0 + + def test_compress_more_than_max_documents(self, many_documents): + """Compress when documents exceed max returns max_documents.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor(max_documents=10) + + result = compressor.compress_documents(many_documents, "topic 1") + + assert len(result) == 10 + assert compressor._last_metrics.documents_before == 50 + assert compressor._last_metrics.documents_after == 10 + assert compressor._last_metrics.documents_removed == 40 + + def test_compress_orders_by_relevance(self, sample_documents): + """Compressed documents are ordered by relevance.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor(max_documents=3) + + result = compressor.compress_documents(sample_documents, "Python programming") + + # Most relevant documents should come first + assert len(result) == 3 + # First doc should be highly relevant to "Python programming" + assert "Python" in result[0].page_content or "programming" in result[0].page_content + + def test_compress_with_min_relevance_filter(self): + """Documents below min_relevance are filtered out.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + documents = [ + Document(page_content="Very relevant Python tutorial"), + Document(page_content="Completely unrelated topic XYZ"), + ] + + compressor = HeadroomDocumentCompressor( + max_documents=10, + min_relevance=0.3, # Require some relevance + ) + + result = compressor.compress_documents(documents, "Python programming") + + # The very relevant doc should pass, unrelated might be filtered + assert len(result) >= 1 + # First result should be the relevant one + assert "Python" in result[0].page_content + + def test_compress_tracks_relevance_scores(self, sample_documents): + """Compression tracks relevance scores.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor(max_documents=3) + + compressor.compress_documents(sample_documents, "Python") + + assert compressor._last_metrics is not None + assert len(compressor._last_metrics.relevance_scores) == 3 + # Scores should be sorted descending + scores = compressor._last_metrics.relevance_scores + assert scores == sorted(scores, reverse=True) + + +class TestHeadroomDocumentCompressorScoring: + """Tests for document relevance scoring.""" + + def test_score_document_exact_match_boost(self): + """Exact phrase match gets relevance boost.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + doc_exact = Document(page_content="What is Python programming?") + doc_partial = Document(page_content="Programming in various languages") + + score_exact = compressor._score_document(doc_exact, "Python programming") + score_partial = compressor._score_document(doc_partial, "Python programming") + + # Exact match should score higher + assert score_exact > score_partial + + def test_score_document_term_frequency(self): + """Higher term frequency increases score.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + doc_many = Document(page_content="Python Python Python is great") + doc_one = Document(page_content="Python is a language") + + score_many = compressor._score_document(doc_many, "Python") + score_one = compressor._score_document(doc_one, "Python") + + # More mentions should score higher (BM25 diminishing returns aside) + assert score_many >= score_one + + def test_score_document_empty_query(self): + """Empty query returns zero score.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + doc = Document(page_content="Some content") + + score = compressor._score_document(doc, "") + + assert score == 0.0 + + def test_score_document_empty_content(self): + """Empty document content returns zero score.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + doc = Document(page_content="") + + score = compressor._score_document(doc, "query") + + assert score == 0.0 + + def test_score_document_case_insensitive(self): + """Scoring is case insensitive.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + doc = Document(page_content="PYTHON is GREAT") + + score = compressor._score_document(doc, "python great") + + assert score > 0.0 + + +class TestHeadroomDocumentCompressorTokenize: + """Tests for text tokenization.""" + + def test_tokenize_basic(self): + """Tokenize basic text.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + tokens = compressor._tokenize("Hello world") + + assert tokens == ["Hello", "world"] + + def test_tokenize_with_punctuation(self): + """Tokenize text with punctuation.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + tokens = compressor._tokenize("Hello, world! How are you?") + + assert "Hello" in tokens + assert "world" in tokens + assert "," not in tokens + assert "!" not in tokens + + def test_tokenize_filters_short_tokens(self): + """Tokenize filters tokens with length 1.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + tokens = compressor._tokenize("I am a developer") + + # "I" and "a" should be filtered out + assert "I" not in tokens + assert "a" not in tokens + assert "am" in tokens + assert "developer" in tokens + + +class TestHeadroomDocumentCompressorDiversity: + """Tests for diverse document selection (MMR-style).""" + + def test_compress_with_diversity(self): + """Diverse selection avoids redundant documents.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + # Create similar documents + documents = [ + Document(page_content="Python is a programming language."), + Document(page_content="Python is a great programming language."), # Very similar + Document(page_content="Python programming tutorial."), # Similar + Document(page_content="Java is a different programming language."), # Different + Document(page_content="Machine learning with TensorFlow."), # Very different + ] + + compressor = HeadroomDocumentCompressor( + max_documents=3, + prefer_diverse=True, + ) + + result = compressor.compress_documents(documents, "programming language") + + assert len(result) == 3 + # Diversity should favor the Java/ML docs over multiple Python docs + + def test_select_diverse_empty(self): + """Diverse selection with empty input.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor(prefer_diverse=True) + + result = compressor._select_diverse([], "query") + + assert result == [] + + def test_document_similarity_identical(self): + """Identical documents have similarity 1.0.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + doc1 = Document(page_content="Hello world") + doc2 = Document(page_content="Hello world") + + similarity = compressor._document_similarity(doc1, doc2) + + assert similarity == 1.0 + + def test_document_similarity_different(self): + """Different documents have low similarity.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + doc1 = Document(page_content="Python programming tutorial") + doc2 = Document(page_content="Cooking recipes for dinner") + + similarity = compressor._document_similarity(doc1, doc2) + + assert similarity < 0.2 # Very different + + def test_document_similarity_partial_overlap(self): + """Partially overlapping documents have medium similarity.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + doc1 = Document(page_content="Python programming tutorial") + doc2 = Document(page_content="Python data science tutorial") + + similarity = compressor._document_similarity(doc1, doc2) + + assert 0.2 < similarity < 0.8 # Some overlap + + def test_document_similarity_empty_content(self): + """Empty content documents have zero similarity.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + doc1 = Document(page_content="") + doc2 = Document(page_content="Some content") + + similarity = compressor._document_similarity(doc1, doc2) + + assert similarity == 0.0 + + +class TestHeadroomDocumentCompressorStats: + """Tests for compression statistics.""" + + def test_last_metrics_none_initially(self): + """last_metrics is None before any compression.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + assert compressor.last_metrics is None + + def test_last_metrics_updated_after_compression(self, sample_documents): + """last_metrics is updated after compression.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor(max_documents=3) + + compressor.compress_documents(sample_documents, "Python") + + assert compressor.last_metrics is not None + assert compressor.last_metrics.documents_before == 5 + assert compressor.last_metrics.documents_after == 3 + + def test_get_compression_stats_empty(self): + """get_compression_stats returns empty dict before compression.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor() + + stats = compressor.get_compression_stats() + + assert stats == {} + + def test_get_compression_stats_with_data(self, many_documents): + """get_compression_stats returns stats after compression.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor(max_documents=10) + + compressor.compress_documents(many_documents, "topic") + + stats = compressor.get_compression_stats() + + assert stats["documents_before"] == 50 + assert stats["documents_after"] == 10 + assert stats["documents_removed"] == 40 + assert "average_relevance" in stats + assert 0 <= stats["average_relevance"] <= 1.0 + + def test_get_compression_stats_average_relevance(self, sample_documents): + """get_compression_stats calculates average relevance correctly.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor(max_documents=2) + + compressor.compress_documents(sample_documents, "Python") + + stats = compressor.get_compression_stats() + + # Average should match manual calculation + expected_avg = sum(compressor._last_metrics.relevance_scores) / len( + compressor._last_metrics.relevance_scores + ) + assert abs(stats["average_relevance"] - expected_avg) < 0.001 + + +class TestHeadroomDocumentCompressorCallbacks: + """Tests for LangChain callbacks integration.""" + + def test_compress_ignores_callbacks(self, sample_documents): + """compress_documents accepts but ignores callbacks parameter.""" + from headroom.integrations.langchain.retriever import HeadroomDocumentCompressor + + compressor = HeadroomDocumentCompressor(max_documents=3) + + # Pass a mock callback - should not raise + mock_callback = MagicMock() + result = compressor.compress_documents( + sample_documents, "Python", callbacks=[mock_callback] + ) + + assert len(result) == 3 + + +class TestLangChainNotAvailable: + """Tests for behavior when LangChain is not available.""" + + def test_check_raises_import_error(self): + """_check_langchain_available raises ImportError when not available.""" + from headroom.integrations.langchain.retriever import _check_langchain_available + + # When LangChain IS available, should not raise + try: + _check_langchain_available() + except ImportError: + pytest.fail("Should not raise when LangChain is available") diff --git a/tests/test_integrations/langchain/test_streaming.py b/tests/test_integrations/langchain/test_streaming.py new file mode 100644 index 000000000..76143a22f --- /dev/null +++ b/tests/test_integrations/langchain/test_streaming.py @@ -0,0 +1,630 @@ +"""Tests for LangChain streaming metrics tracking. + +Tests cover: +1. StreamingMetrics - Dataclass for streaming response metrics +2. StreamingMetricsTracker - Tracker for streaming chunks +3. StreamingMetricsCallback - Context manager for streaming +4. track_streaming_response - Sync helper function +5. track_async_streaming_response - Async helper function +""" + +from datetime import datetime +from unittest.mock import MagicMock, patch + +import pytest + +# Check if LangChain is available +try: + from langchain_core.messages import AIMessageChunk + from langchain_core.outputs import ChatGenerationChunk + + LANGCHAIN_AVAILABLE = True +except ImportError: + LANGCHAIN_AVAILABLE = False + +# Skip all tests if LangChain not installed +pytestmark = pytest.mark.skipif(not LANGCHAIN_AVAILABLE, reason="LangChain not installed") + + +@pytest.fixture +def mock_provider(): + """Create a mock provider with token counter.""" + mock = MagicMock() + mock_counter = MagicMock() + # Simple token counting: split on spaces + mock_counter.count_text = MagicMock(side_effect=lambda text: len(text.split())) + mock.get_token_counter = MagicMock(return_value=mock_counter) + return mock + + +@pytest.fixture +def sample_chunks(): + """Create sample streaming chunks.""" + return [ + AIMessageChunk(content="Hello"), + AIMessageChunk(content=" "), + AIMessageChunk(content="world"), + AIMessageChunk(content="!"), + ] + + +class TestStreamingMetrics: + """Tests for StreamingMetrics dataclass.""" + + def test_create_metrics(self): + """Create metrics with all fields.""" + from headroom.integrations.langchain.streaming import StreamingMetrics + + start = datetime.now() + end = datetime.now() + + metrics = StreamingMetrics( + output_tokens=50, + chunk_count=10, + content_length=200, + start_time=start, + end_time=end, + duration_ms=150.5, + ) + + assert metrics.output_tokens == 50 + assert metrics.chunk_count == 10 + assert metrics.content_length == 200 + assert metrics.start_time == start + assert metrics.end_time == end + assert metrics.duration_ms == 150.5 + + def test_to_dict(self): + """Convert metrics to dictionary.""" + from headroom.integrations.langchain.streaming import StreamingMetrics + + start = datetime(2025, 1, 1, 12, 0, 0) + end = datetime(2025, 1, 1, 12, 0, 1) + + metrics = StreamingMetrics( + output_tokens=50, + chunk_count=10, + content_length=200, + start_time=start, + end_time=end, + duration_ms=1000.0, + ) + + result = metrics.to_dict() + + assert result["output_tokens"] == 50 + assert result["chunk_count"] == 10 + assert result["content_length"] == 200 + assert result["start_time"] == "2025-01-01T12:00:00" + assert result["end_time"] == "2025-01-01T12:00:01" + assert result["duration_ms"] == 1000.0 + + def test_to_dict_with_none_end_time(self): + """Convert metrics with None end_time.""" + from headroom.integrations.langchain.streaming import StreamingMetrics + + metrics = StreamingMetrics( + output_tokens=50, + chunk_count=10, + content_length=200, + start_time=datetime.now(), + end_time=None, + duration_ms=None, + ) + + result = metrics.to_dict() + + assert result["end_time"] is None + assert result["duration_ms"] is None + + +class TestStreamingMetricsTrackerInit: + """Tests for StreamingMetricsTracker initialization.""" + + def test_init_defaults(self): + """Initialize with default settings.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + with patch("headroom.integrations.langchain.streaming.OpenAIProvider"): + tracker = StreamingMetricsTracker() + + assert tracker._model == "gpt-4o" + assert tracker._content == "" + assert tracker._chunk_count == 0 + assert tracker._start_time is None + assert tracker._end_time is None + + def test_init_custom_settings(self, mock_provider): + """Initialize with custom settings.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker( + model="claude-3-5-sonnet-20241022", + provider=mock_provider, + ) + + assert tracker._model == "claude-3-5-sonnet-20241022" + assert tracker._provider is mock_provider + + +class TestStreamingMetricsTrackerAddChunk: + """Tests for add_chunk method.""" + + def test_add_chunk_sets_start_time(self, mock_provider): + """First chunk sets start time.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + assert tracker._start_time is None + + chunk = AIMessageChunk(content="Hello") + tracker.add_chunk(chunk) + + assert tracker._start_time is not None + + def test_add_chunk_increments_count(self, mock_provider, sample_chunks): + """Each chunk increments chunk count.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + assert tracker._chunk_count == 4 + + def test_add_chunk_accumulates_content(self, mock_provider, sample_chunks): + """Chunks accumulate content.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + assert tracker._content == "Hello world!" + + def test_add_chunk_extracts_ai_message_chunk(self, mock_provider): + """Extract content from AIMessageChunk.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + chunk = AIMessageChunk(content="Hello") + tracker.add_chunk(chunk) + + assert tracker._content == "Hello" + + def test_add_chunk_extracts_chat_generation_chunk(self, mock_provider): + """Extract content from ChatGenerationChunk.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + chunk = ChatGenerationChunk(message=AIMessageChunk(content="Hello")) + tracker.add_chunk(chunk) + + assert tracker._content == "Hello" + + def test_add_chunk_extracts_dict(self, mock_provider): + """Extract content from dict.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + chunk = {"content": "Hello"} + tracker.add_chunk(chunk) + + assert tracker._content == "Hello" + + def test_add_chunk_extracts_string(self, mock_provider): + """Extract content from string.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + tracker.add_chunk("Hello") + + assert tracker._content == "Hello" + + def test_add_chunk_handles_empty_content(self, mock_provider): + """Handle chunk with empty content.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + chunk = AIMessageChunk(content="") + tracker.add_chunk(chunk) + + assert tracker._content == "" + assert tracker._chunk_count == 1 + + def test_add_chunk_handles_none_content(self, mock_provider): + """Handle chunk with None content attribute.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + chunk = MagicMock() + chunk.content = None + tracker.add_chunk(chunk) + + assert tracker._content == "" + assert tracker._chunk_count == 1 + + +class TestStreamingMetricsTrackerFinish: + """Tests for finish method.""" + + def test_finish_sets_end_time(self, mock_provider, sample_chunks): + """finish() sets end time.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + metrics = tracker.finish() + + assert tracker._end_time is not None + assert metrics.end_time is not None + + def test_finish_calculates_duration(self, mock_provider, sample_chunks): + """finish() calculates duration.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + metrics = tracker.finish() + + assert metrics.duration_ms is not None + assert metrics.duration_ms >= 0 + + def test_finish_returns_metrics(self, mock_provider, sample_chunks): + """finish() returns StreamingMetrics.""" + from headroom.integrations.langchain.streaming import ( + StreamingMetrics, + StreamingMetricsTracker, + ) + + tracker = StreamingMetricsTracker(provider=mock_provider) + + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + metrics = tracker.finish() + + assert isinstance(metrics, StreamingMetrics) + assert metrics.chunk_count == 4 + assert metrics.content_length == len("Hello world!") + + def test_finish_with_no_chunks(self, mock_provider): + """finish() without chunks uses current time for both.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + metrics = tracker.finish() + + # start_time should be same as end_time when no chunks + assert metrics.start_time == metrics.end_time + assert metrics.duration_ms is None # No start_time was set + + +class TestStreamingMetricsTrackerProperties: + """Tests for tracker properties.""" + + def test_content_property(self, mock_provider, sample_chunks): + """content property returns accumulated content.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + assert tracker.content == "Hello world!" + + def test_output_tokens_property_empty(self, mock_provider): + """output_tokens returns 0 when no content.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + assert tracker.output_tokens == 0 + + def test_output_tokens_property_with_content(self, mock_provider, sample_chunks): + """output_tokens uses provider's token counter.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker( + model="gpt-4o", + provider=mock_provider, + ) + + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + tokens = tracker.output_tokens + + # Mock counter splits on spaces: "Hello world!" = 2 tokens + assert tokens == 2 + mock_provider.get_token_counter.assert_called_with("gpt-4o") + + def test_chunk_count_property(self, mock_provider, sample_chunks): + """chunk_count property returns number of chunks.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + assert tracker.chunk_count == 4 + + def test_duration_ms_before_finish(self, mock_provider, sample_chunks): + """duration_ms returns None before finish().""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + assert tracker.duration_ms is None + + def test_duration_ms_after_finish(self, mock_provider, sample_chunks): + """duration_ms returns value after finish().""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + tracker.finish() + + assert tracker.duration_ms is not None + assert tracker.duration_ms >= 0 + + +class TestStreamingMetricsTrackerReset: + """Tests for reset method.""" + + def test_reset_clears_state(self, mock_provider, sample_chunks): + """reset() clears all state.""" + from headroom.integrations.langchain.streaming import StreamingMetricsTracker + + tracker = StreamingMetricsTracker(provider=mock_provider) + + for chunk in sample_chunks: + tracker.add_chunk(chunk) + tracker.finish() + + tracker.reset() + + assert tracker._content == "" + assert tracker._chunk_count == 0 + assert tracker._start_time is None + assert tracker._end_time is None + + +class TestStreamingMetricsCallback: + """Tests for StreamingMetricsCallback context manager.""" + + def test_init(self, mock_provider): + """Initialize callback.""" + from headroom.integrations.langchain.streaming import StreamingMetricsCallback + + callback = StreamingMetricsCallback(model="gpt-4o", provider=mock_provider) + + assert callback._tracker._model == "gpt-4o" + assert callback._metrics is None + + def test_context_manager_enter(self, mock_provider): + """Context manager enter returns tracker.""" + from headroom.integrations.langchain.streaming import ( + StreamingMetricsCallback, + StreamingMetricsTracker, + ) + + callback = StreamingMetricsCallback(provider=mock_provider) + + with callback as tracker: + assert isinstance(tracker, StreamingMetricsTracker) + + def test_context_manager_exit_finishes_tracker(self, mock_provider, sample_chunks): + """Context manager exit finishes tracker.""" + from headroom.integrations.langchain.streaming import StreamingMetricsCallback + + callback = StreamingMetricsCallback(provider=mock_provider) + + with callback as tracker: + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + assert callback.metrics is not None + assert callback.metrics.chunk_count == 4 + + def test_tracker_property(self, mock_provider): + """tracker property returns the tracker.""" + from headroom.integrations.langchain.streaming import ( + StreamingMetricsCallback, + StreamingMetricsTracker, + ) + + callback = StreamingMetricsCallback(provider=mock_provider) + + assert isinstance(callback.tracker, StreamingMetricsTracker) + + def test_metrics_property_before_exit(self, mock_provider): + """metrics property returns None before context exit.""" + from headroom.integrations.langchain.streaming import StreamingMetricsCallback + + callback = StreamingMetricsCallback(provider=mock_provider) + + assert callback.metrics is None + + def test_metrics_property_after_exit(self, mock_provider, sample_chunks): + """metrics property returns StreamingMetrics after context exit.""" + from headroom.integrations.langchain.streaming import ( + StreamingMetrics, + StreamingMetricsCallback, + ) + + callback = StreamingMetricsCallback(provider=mock_provider) + + with callback as tracker: + for chunk in sample_chunks: + tracker.add_chunk(chunk) + + assert isinstance(callback.metrics, StreamingMetrics) + + +class TestTrackStreamingResponse: + """Tests for track_streaming_response function.""" + + def test_consumes_stream(self, mock_provider, sample_chunks): + """Function consumes entire stream.""" + from headroom.integrations.langchain.streaming import track_streaming_response + + stream = iter(sample_chunks) + + content, metrics = track_streaming_response(stream, provider=mock_provider) + + assert content == "Hello world!" + + def test_returns_content_and_metrics(self, mock_provider, sample_chunks): + """Function returns content and metrics tuple.""" + from headroom.integrations.langchain.streaming import ( + StreamingMetrics, + track_streaming_response, + ) + + stream = iter(sample_chunks) + + content, metrics = track_streaming_response(stream, provider=mock_provider) + + assert isinstance(content, str) + assert isinstance(metrics, StreamingMetrics) + + def test_with_custom_model(self, mock_provider, sample_chunks): + """Function uses custom model for token counting.""" + from headroom.integrations.langchain.streaming import track_streaming_response + + stream = iter(sample_chunks) + + content, metrics = track_streaming_response( + stream, + model="claude-3-5-sonnet-20241022", + provider=mock_provider, + ) + + mock_provider.get_token_counter.assert_called_with("claude-3-5-sonnet-20241022") + + def test_empty_stream(self, mock_provider): + """Function handles empty stream.""" + from headroom.integrations.langchain.streaming import track_streaming_response + + stream = iter([]) + + content, metrics = track_streaming_response(stream, provider=mock_provider) + + assert content == "" + assert metrics.chunk_count == 0 + + +class TestTrackAsyncStreamingResponse: + """Tests for track_async_streaming_response function.""" + + @pytest.mark.asyncio + async def test_consumes_async_stream(self, mock_provider, sample_chunks): + """Function consumes entire async stream.""" + from headroom.integrations.langchain.streaming import ( + track_async_streaming_response, + ) + + async def async_stream(): + for chunk in sample_chunks: + yield chunk + + content, metrics = await track_async_streaming_response( + async_stream(), provider=mock_provider + ) + + assert content == "Hello world!" + + @pytest.mark.asyncio + async def test_returns_content_and_metrics(self, mock_provider, sample_chunks): + """Function returns content and metrics tuple.""" + from headroom.integrations.langchain.streaming import ( + StreamingMetrics, + track_async_streaming_response, + ) + + async def async_stream(): + for chunk in sample_chunks: + yield chunk + + content, metrics = await track_async_streaming_response( + async_stream(), provider=mock_provider + ) + + assert isinstance(content, str) + assert isinstance(metrics, StreamingMetrics) + + @pytest.mark.asyncio + async def test_with_custom_model(self, mock_provider, sample_chunks): + """Function uses custom model for token counting.""" + from headroom.integrations.langchain.streaming import ( + track_async_streaming_response, + ) + + async def async_stream(): + for chunk in sample_chunks: + yield chunk + + content, metrics = await track_async_streaming_response( + async_stream(), + model="gpt-4-turbo", + provider=mock_provider, + ) + + mock_provider.get_token_counter.assert_called_with("gpt-4-turbo") + + @pytest.mark.asyncio + async def test_empty_async_stream(self, mock_provider): + """Function handles empty async stream.""" + from headroom.integrations.langchain.streaming import ( + track_async_streaming_response, + ) + + async def async_stream(): + return + yield # Make it a generator # noqa: B901 - intentionally unreachable + + content, metrics = await track_async_streaming_response( + async_stream(), provider=mock_provider + ) + + assert content == "" + assert metrics.chunk_count == 0 + + +class TestLangChainNotAvailable: + """Tests for behavior when LangChain is not available.""" + + def test_check_raises_import_error(self): + """_check_langchain_available raises ImportError when not available.""" + from headroom.integrations.langchain.streaming import _check_langchain_available + + # When LangChain IS available, should not raise + try: + _check_langchain_available() + except ImportError: + pytest.fail("Should not raise when LangChain is available") diff --git a/tests/test_log_compressor.py b/tests/test_log_compressor.py new file mode 100644 index 000000000..2b703dc7e --- /dev/null +++ b/tests/test_log_compressor.py @@ -0,0 +1,742 @@ +"""Comprehensive tests for log_compressor.py. + +Tests cover: +1. Detection of different log formats (pytest, npm, cargo, make, jest, generic) +2. Line extraction and deduplication +3. Compression ratios +4. Edge cases +""" + +from headroom.transforms.log_compressor import ( + LogCompressionResult, + LogCompressor, + LogCompressorConfig, + LogFormat, + LogLevel, + LogLine, +) + + +class TestLogFormatDetection: + """Tests for detecting different log formats.""" + + def test_detect_pytest_format(self): + """Pytest output is detected correctly.""" + content = """============================= test session starts ============================== +platform darwin -- Python 3.11.0 +collected 15 items + +tests/test_foo.py::test_basic PASSED [ 6%] +tests/test_foo.py::test_edge FAILED [ 13%] + +=================================== FAILURES =================================== +tests/test_foo.py::test_edge - AssertionError + +=========================== short test summary info ============================ +FAILED tests/test_foo.py::test_edge +========================= 1 failed, 14 passed ========================= +""" + compressor = LogCompressor() + lines = content.split("\n") + detected = compressor._detect_format(lines) + assert detected == LogFormat.PYTEST + + def test_detect_npm_format(self): + """npm output is detected correctly.""" + content = """npm WARN deprecated package@1.0.0: This package is deprecated +npm WARN deprecated another@2.0.0: Obsolete +npm ERR! code ERESOLVE +npm ERR! ERESOLVE unable to resolve dependency tree +npm info using npm@9.0.0 +> added 150 packages in 5s +""" + compressor = LogCompressor() + lines = content.split("\n") + detected = compressor._detect_format(lines) + assert detected == LogFormat.NPM + + def test_detect_cargo_format(self): + """Cargo/rustc output is detected correctly.""" + content = """ Compiling myproject v0.1.0 (/path/to/project) +warning: unused variable: `x` + --> src/main.rs:5:9 + | +5 | let x = 5; + | ^ help: if this is intentional, prefix it with an underscore: `_x` + | + = note: `#[warn(unused_variables)]` on by default + +error[E0382]: borrow of moved value: `s` + Finished dev [unoptimized + debuginfo] target(s) in 0.50s + Running `target/debug/myproject` +""" + compressor = LogCompressor() + lines = content.split("\n") + detected = compressor._detect_format(lines) + assert detected == LogFormat.CARGO + + def test_detect_make_format(self): + """make/gcc output is detected correctly.""" + content = """make[1]: Entering directory '/path/to/project' +gcc -c -o main.o main.c +gcc -c -o utils.o utils.c +make[1]: *** [Makefile:10: utils.o] Error 1 +make: *** [Makefile:5: all] Error 2 +g++ -Wall -o program main.cpp utils.cpp +""" + compressor = LogCompressor() + lines = content.split("\n") + detected = compressor._detect_format(lines) + assert detected == LogFormat.MAKE + + def test_detect_jest_format(self): + """Jest output is detected correctly.""" + content = """PASS src/components/Button.test.js +FAIL src/utils/helpers.test.ts + Test Suites: 1 failed, 1 passed, 2 total + Tests: 2 failed, 10 passed, 12 total +""" + compressor = LogCompressor() + lines = content.split("\n") + detected = compressor._detect_format(lines) + assert detected == LogFormat.JEST + + def test_detect_generic_format(self): + """Generic log format is detected for unrecognized output.""" + content = """INFO Starting application +DEBUG Initializing components +WARNING Low memory +ERROR Connection timeout +CRITICAL System failure +""" + compressor = LogCompressor() + lines = content.split("\n") + detected = compressor._detect_format(lines) + assert detected == LogFormat.GENERIC + + def test_detect_empty_returns_generic(self): + """Empty or minimal input returns GENERIC.""" + compressor = LogCompressor() + assert compressor._detect_format([]) == LogFormat.GENERIC + assert compressor._detect_format(["random line"]) == LogFormat.GENERIC + + +class TestLogLevelDetection: + """Tests for log level detection in lines.""" + + def test_detect_error_levels(self): + """ERROR, FATAL, CRITICAL are detected.""" + compressor = LogCompressor() + + error_lines = [ + "ERROR: something went wrong", + "error: file not found", + "Error: Invalid input", + "FATAL: system crash", + "fatal error occurred", + "CRITICAL: database down", + ] + + for line in error_lines: + log_lines = compressor._parse_lines([line]) + assert log_lines[0].level == LogLevel.ERROR, f"Failed for: {line}" + + def test_detect_fail_levels(self): + """FAIL, FAILED are detected.""" + compressor = LogCompressor() + + fail_lines = [ + "FAIL tests/test_foo.py", + "FAILED to connect", + "Test failed", + ] + + for line in fail_lines: + log_lines = compressor._parse_lines([line]) + assert log_lines[0].level == LogLevel.FAIL, f"Failed for: {line}" + + def test_detect_warn_levels(self): + """WARN, WARNING are detected.""" + compressor = LogCompressor() + + warn_lines = [ + "WARN: deprecated function", + "WARNING: low disk space", + "warning: unused variable", + ] + + for line in warn_lines: + log_lines = compressor._parse_lines([line]) + assert log_lines[0].level == LogLevel.WARN, f"Failed for: {line}" + + def test_detect_info_debug_trace(self): + """INFO, DEBUG, TRACE are detected.""" + compressor = LogCompressor() + + test_cases = [ + ("INFO: starting process", LogLevel.INFO), + ("info starting", LogLevel.INFO), + ("DEBUG: variable x = 5", LogLevel.DEBUG), + ("debug mode enabled", LogLevel.DEBUG), + ("TRACE: entering function", LogLevel.TRACE), + ] + + for line, expected_level in test_cases: + log_lines = compressor._parse_lines([line]) + assert log_lines[0].level == expected_level, f"Failed for: {line}" + + def test_unknown_level_default(self): + """Lines without level markers default to UNKNOWN.""" + compressor = LogCompressor() + log_lines = compressor._parse_lines(["Just some regular text"]) + assert log_lines[0].level == LogLevel.UNKNOWN + + +class TestStackTraceDetection: + """Tests for stack trace detection.""" + + def test_detect_python_traceback(self): + """Python traceback is detected.""" + content = """Traceback (most recent call last): + File "main.py", line 42, in process + result = compute(data) + File "utils.py", line 15, in compute + return data / 0 +ZeroDivisionError: division by zero +""" + compressor = LogCompressor() + log_lines = compressor._parse_lines(content.split("\n")) + + # First several lines should be marked as stack trace + stack_trace_count = sum(1 for line in log_lines if line.is_stack_trace) + assert stack_trace_count > 0 + + def test_detect_javascript_stack_trace(self): + """JavaScript stack trace is detected.""" + content = """Error: Connection failed + at Connection.connect (src/db.js:42:15) + at async main (src/index.js:10:5) +""" + compressor = LogCompressor() + log_lines = compressor._parse_lines(content.split("\n")) + + stack_trace_count = sum(1 for line in log_lines if line.is_stack_trace) + assert stack_trace_count > 0 + + def test_detect_rust_error_location(self): + """Rust error location is detected.""" + content = """error[E0382]: borrow of moved value: `s` + --> src/main.rs:5:13 + | +3 | let s = String::from("hello"); + | - move occurs +""" + compressor = LogCompressor() + log_lines = compressor._parse_lines(content.split("\n")) + + stack_trace_count = sum(1 for line in log_lines if line.is_stack_trace) + assert stack_trace_count > 0 + + +class TestLineDeduplication: + """Tests for warning/line deduplication.""" + + def test_dedupe_identical_warnings(self): + """Identical warnings are deduplicated.""" + compressor = LogCompressor() + + lines = [ + LogLine(line_number=1, content="WARNING: unused variable 'x'", level=LogLevel.WARN), + LogLine(line_number=2, content="WARNING: unused variable 'x'", level=LogLevel.WARN), + LogLine(line_number=3, content="WARNING: unused variable 'x'", level=LogLevel.WARN), + ] + + deduped = compressor._dedupe_similar(lines) + assert len(deduped) == 1 + + def test_dedupe_similar_with_numbers(self): + """Similar warnings with different numbers are deduplicated.""" + compressor = LogCompressor() + + lines = [ + LogLine(line_number=1, content="WARNING: error at line 10", level=LogLevel.WARN), + LogLine(line_number=2, content="WARNING: error at line 20", level=LogLevel.WARN), + LogLine(line_number=3, content="WARNING: error at line 30", level=LogLevel.WARN), + ] + + deduped = compressor._dedupe_similar(lines) + # Numbers normalized to "N", so all three are treated as identical pattern + assert len(deduped) == 1 + + def test_dedupe_similar_with_paths(self): + """Similar warnings with different paths are deduplicated. + + Note: The path regex /[\\w/]+/ requires paths to end with '/'. + Paths like '/path/to/' will be normalized, but '/path/to/file' won't + be fully normalized because 'file' doesn't end with '/'. + """ + compressor = LogCompressor() + + # Paths ending with / are normalized + lines = [ + LogLine(line_number=1, content="WARNING: in /path/to/ error", level=LogLevel.WARN), + LogLine(line_number=2, content="WARNING: in /other/dir/ error", level=LogLevel.WARN), + LogLine(line_number=3, content="WARNING: in /another/path/ error", level=LogLevel.WARN), + ] + + deduped = compressor._dedupe_similar(lines) + # Paths normalized to /PATH/, so all three are treated as identical pattern + assert len(deduped) == 1 + + def test_keeps_different_warnings(self): + """Different warnings are preserved.""" + compressor = LogCompressor() + + lines = [ + LogLine(line_number=1, content="WARNING: unused variable", level=LogLevel.WARN), + LogLine(line_number=2, content="WARNING: deprecated function", level=LogLevel.WARN), + LogLine(line_number=3, content="WARNING: missing docstring", level=LogLevel.WARN), + ] + + deduped = compressor._dedupe_similar(lines) + assert len(deduped) == 3 + + +class TestLineScoring: + """Tests for line importance scoring.""" + + def test_error_lines_score_highest(self): + """ERROR and FAIL lines get highest scores.""" + compressor = LogCompressor() + + error_line = LogLine(line_number=1, content="ERROR: critical", level=LogLevel.ERROR) + fail_line = LogLine(line_number=2, content="FAILED test", level=LogLevel.FAIL) + info_line = LogLine(line_number=3, content="INFO: normal", level=LogLevel.INFO) + + error_score = compressor._score_line(error_line) + fail_score = compressor._score_line(fail_line) + info_score = compressor._score_line(info_line) + + assert error_score > info_score + assert fail_score > info_score + + def test_stack_trace_boost(self): + """Stack trace lines get boosted score.""" + compressor = LogCompressor() + + regular = LogLine(line_number=1, content="some line", level=LogLevel.UNKNOWN) + stack_trace = LogLine( + line_number=2, content=" File 'x.py'", level=LogLevel.UNKNOWN, is_stack_trace=True + ) + + assert compressor._score_line(stack_trace) > compressor._score_line(regular) + + def test_summary_line_boost(self): + """Summary lines get boosted score.""" + compressor = LogCompressor() + + regular = LogLine(line_number=1, content="some line", level=LogLevel.UNKNOWN) + summary = LogLine( + line_number=2, content="10 passed, 2 failed", level=LogLevel.UNKNOWN, is_summary=True + ) + + assert compressor._score_line(summary) > compressor._score_line(regular) + + +class TestCompressionBehavior: + """Tests for overall compression behavior.""" + + def test_small_log_passthrough(self): + """Logs smaller than threshold pass through unchanged.""" + content = "INFO: Starting\nINFO: Done" + + compressor = LogCompressor(config=LogCompressorConfig(min_lines_for_ccr=100)) + result = compressor.compress(content) + + assert result.compression_ratio == 1.0 + assert result.compressed == content + assert result.original_line_count == 2 + + def test_large_log_compressed(self): + """Large logs are compressed.""" + lines = [f"INFO: Processing item {i}" for i in range(200)] + lines.append("ERROR: Failed at item 100") + content = "\n".join(lines) + + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=50, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + assert result.compression_ratio < 1.0 + assert result.compressed_line_count < result.original_line_count + # Error is preserved + assert "ERROR: Failed" in result.compressed + + def test_keeps_first_and_last_errors(self): + """First and last errors are preserved.""" + lines = [f"INFO: item {i}" for i in range(100)] + lines[10] = "ERROR: first error" + lines[50] = "ERROR: middle error" + lines[90] = "ERROR: last error" + content = "\n".join(lines) + + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=50, + keep_first_error=True, + keep_last_error=True, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + assert "first error" in result.compressed + assert "last error" in result.compressed + + def test_summary_lines_preserved(self): + """Summary lines are always preserved.""" + content = """INFO: test 1 +INFO: test 2 +======================================== +TOTAL: 10 tests passed +Build succeeded in 5.2s +""" + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=2, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + assert "========" in result.compressed + assert "TOTAL:" in result.compressed or "Build succeeded" in result.compressed + + def test_context_lines_added(self): + """Context lines around errors are included.""" + lines = [f"INFO: item {i}" for i in range(100)] + lines[50] = "ERROR: critical failure" + content = "\n".join(lines) + + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=50, + error_context_lines=2, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Should have context around the error + assert "item 48" in result.compressed or "item 49" in result.compressed + assert "item 51" in result.compressed or "item 52" in result.compressed + + +class TestCompressionRatios: + """Tests for compression ratio calculations.""" + + def test_compression_ratio_calculation(self): + """Compression ratio is calculated correctly.""" + content = "a" * 1000 # 1000 chars + compressed = "b" * 100 # 100 chars + + # Direct calculation: len(compressed) / len(content) + expected_ratio = 100 / 1000 # 0.1 + + # Result ratio is based on character counts + result = LogCompressionResult( + compressed=compressed, + original=content, + original_line_count=100, + compressed_line_count=10, + format_detected=LogFormat.GENERIC, + compression_ratio=len(compressed) / len(content), + ) + + assert result.compression_ratio == expected_ratio + + def test_tokens_saved_estimate(self): + """Token savings estimation works correctly.""" + content = "a" * 400 # ~100 tokens + compressed = "b" * 40 # ~10 tokens + + result = LogCompressionResult( + compressed=compressed, + original=content, + original_line_count=10, + compressed_line_count=1, + format_detected=LogFormat.GENERIC, + compression_ratio=0.1, + ) + + # (400 - 40) / 4 = 90 tokens saved + assert result.tokens_saved_estimate == 90 + + def test_lines_omitted_property(self): + """Lines omitted property works correctly.""" + result = LogCompressionResult( + compressed="test", + original="test\noriginal", + original_line_count=100, + compressed_line_count=10, + format_detected=LogFormat.GENERIC, + compression_ratio=0.1, + ) + + assert result.lines_omitted == 90 + + +class TestEdgeCases: + """Tests for edge cases and boundary conditions.""" + + def test_empty_input(self): + """Empty input is handled gracefully.""" + compressor = LogCompressor() + result = compressor.compress("") + + assert result.compressed == "" + assert result.original_line_count == 1 # Empty string splits to one empty line + assert result.compression_ratio == 1.0 + + def test_single_line_input(self): + """Single line input passes through.""" + compressor = LogCompressor() + result = compressor.compress("Single line of text") + + assert result.compressed == "Single line of text" + assert result.compression_ratio == 1.0 + + def test_all_errors_no_info(self): + """Log with only errors is handled.""" + lines = [f"ERROR: failure {i}" for i in range(100)] + content = "\n".join(lines) + + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=50, + max_errors=5, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Should limit to max_errors + assert result.compressed_line_count <= compressor.config.max_total_lines + + def test_unicode_content(self): + """Unicode characters are handled correctly.""" + content = """INFO: Processing 日本語 +ERROR: Failed with émoji 🚀 +WARN: Über important +""" + compressor = LogCompressor() + result = compressor.compress(content) + + # Should not crash and preserve unicode + assert ( + "日本語" in result.compressed + or "émoji" in result.compressed + or "Über" in result.compressed + ) + + def test_very_long_lines(self): + """Very long lines don't cause issues.""" + long_line = "ERROR: " + "x" * 10000 + lines = [f"INFO: line {i}" for i in range(100)] + lines[50] = long_line + content = "\n".join(lines) + + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=50, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Should complete without error + assert len(result.compressed) > 0 + + def test_mixed_line_endings(self): + """Mixed line endings are handled.""" + content = "INFO: line 1\r\nERROR: line 2\rINFO: line 3\n" + + compressor = LogCompressor() + # Should not crash + result = compressor.compress(content) + assert result.compressed is not None + + def test_binary_like_content(self): + """Content with binary-like patterns doesn't crash.""" + content = "INFO: data\x00\x01\x02ERROR: test" + + compressor = LogCompressor() + result = compressor.compress(content) + assert result.compressed is not None + + +class TestConfigOptions: + """Tests for configuration options.""" + + def test_max_errors_config(self): + """max_errors configuration limits error selection.""" + lines = [f"ERROR: error {i}" for i in range(50)] + content = "\n".join(lines) + + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=10, + max_errors=3, + max_total_lines=50, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Count error lines in output (excluding summary line) + error_count = sum(1 for line in result.compressed.split("\n") if "ERROR:" in line) + assert error_count <= 3 + compressor.config.error_context_lines * 2 + + def test_max_warnings_config(self): + """max_warnings configuration limits warning selection.""" + lines = [f"WARN: warning {i}" for i in range(50)] + content = "\n".join(lines) + + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=10, + max_warnings=2, + dedupe_warnings=False, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Warnings should be limited + warn_count = sum(1 for line in result.compressed.split("\n") if "WARN:" in line) + assert warn_count <= 2 + compressor.config.error_context_lines * 2 + + def test_max_total_lines_config(self): + """max_total_lines configuration limits output.""" + lines = [f"ERROR: error {i}" for i in range(200)] + content = "\n".join(lines) + + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=50, + max_total_lines=20, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Output lines should be limited (plus summary line) + output_lines = [line for line in result.compressed.split("\n") if line.strip()] + assert len(output_lines) <= 21 # max_total_lines + 1 summary + + def test_dedupe_warnings_disabled(self): + """dedupe_warnings=False preserves duplicate warnings.""" + lines = [ + "WARN: same warning", + "WARN: same warning", + "WARN: same warning", + ] + content = "\n".join(lines) + + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=1, + dedupe_warnings=False, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # All warnings preserved when dedupe disabled + warn_count = sum(1 for line in result.compressed.split("\n") if "WARN:" in line) + assert warn_count == 3 + + +class TestLogLineDataclass: + """Tests for LogLine dataclass behavior.""" + + def test_equality_by_line_number(self): + """LogLine equality is based on line_number.""" + line1 = LogLine(line_number=10, content="foo") + line2 = LogLine(line_number=10, content="bar") + line3 = LogLine(line_number=20, content="foo") + + assert line1 == line2 + assert line1 != line3 + + def test_hash_by_line_number(self): + """LogLine hash is based on line_number.""" + line1 = LogLine(line_number=10, content="foo") + line2 = LogLine(line_number=10, content="bar") + + assert hash(line1) == hash(line2) + + # Can be used in sets + line_set = {line1, line2} + assert len(line_set) == 1 + + def test_default_values(self): + """LogLine default values are correct.""" + line = LogLine(line_number=1, content="test") + + assert line.level == LogLevel.UNKNOWN + assert line.is_stack_trace is False + assert line.is_summary is False + assert line.score == 0.0 + + +class TestOutputFormatting: + """Tests for output formatting and stats.""" + + def test_format_output_includes_stats(self): + """Format output includes category stats.""" + lines = [ + "ERROR: error 1", + "ERROR: error 2", + "WARN: warning 1", + "INFO: info 1", + "INFO: info 2", + "INFO: info 3", + ] * 20 # Make it large enough to trigger compression + content = "\n".join(lines) + + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=50, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Stats should be populated + assert "errors" in result.stats + assert "warnings" in result.stats + assert "info" in result.stats + assert result.stats["errors"] > 0 + assert result.stats["warnings"] > 0 + + def test_format_output_summary_line(self): + """Formatted output includes summary of omitted lines.""" + lines = [f"INFO: message {i}" for i in range(200)] + lines.append("ERROR: critical") + content = "\n".join(lines) + + compressor = LogCompressor( + config=LogCompressorConfig( + min_lines_for_ccr=50, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Should have omission summary + assert "lines omitted" in result.compressed diff --git a/tests/test_search_compressor.py b/tests/test_search_compressor.py new file mode 100644 index 000000000..962209c31 --- /dev/null +++ b/tests/test_search_compressor.py @@ -0,0 +1,768 @@ +"""Comprehensive tests for search_compressor.py. + +Tests cover: +1. grep/ripgrep output parsing +2. File grouping +3. Match selection and scoring +4. Edge cases +""" + +from headroom.transforms.search_compressor import ( + FileMatches, + SearchCompressionResult, + SearchCompressor, + SearchCompressorConfig, + SearchMatch, +) + + +class TestGrepOutputParsing: + """Tests for parsing grep/ripgrep style output.""" + + def test_parse_standard_grep_format(self): + """Standard grep -n format is parsed correctly.""" + content = """src/main.py:42:def process_data(items): +src/main.py:43: \"\"\"Process items.\"\"\" +src/utils.py:15:def validate(data): +""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + + assert "src/main.py" in file_matches + assert "src/utils.py" in file_matches + assert len(file_matches["src/main.py"].matches) == 2 + assert len(file_matches["src/utils.py"].matches) == 1 + + def test_parse_ripgrep_context_format(self): + """Ripgrep with context (- separator) is parsed.""" + content = """src/main.py-40-some context before +src/main.py:42:def process_data(items): +src/main.py-43-some context after +""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + + assert "src/main.py" in file_matches + # All three lines should be parsed (both : and - separators) + assert len(file_matches["src/main.py"].matches) == 3 + + def test_parse_with_colons_in_content(self): + """Content containing colons is parsed correctly.""" + content = """src/config.py:10:DATABASE_URL = "postgres://user:pass@host:5432/db" +src/config.py:20:REDIS_URL = "redis://localhost:6379" +""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + + assert "src/config.py" in file_matches + matches = file_matches["src/config.py"].matches + + # Content after the second colon should be preserved + assert "postgres://user:pass@host:5432/db" in matches[0].content + + def test_parse_windows_paths(self): + """Windows-style paths are handled.""" + content = """C:\\Users\\dev\\src\\main.py:10:def main(): +C:\\Users\\dev\\src\\utils.py:20:def helper(): +""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + + # Windows paths may not parse correctly due to : in path + # This tests current behavior + assert len(file_matches) >= 0 # Just ensure no crash + + def test_parse_empty_content(self): + """Empty input returns empty result.""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results("") + + assert file_matches == {} + + def test_parse_whitespace_only(self): + """Whitespace-only input returns empty result.""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(" \n\n \n") + + assert file_matches == {} + + def test_parse_non_grep_content(self): + """Non-grep content returns empty result.""" + content = """This is just regular text +without any grep-style formatting +just normal lines here""" + + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + + assert file_matches == {} + + def test_parse_mixed_valid_invalid(self): + """Mixed valid and invalid lines parse valid ones.""" + content = """src/main.py:10:valid line +this is not a grep line +src/utils.py:20:another valid line +more random text +""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + + assert "src/main.py" in file_matches + assert "src/utils.py" in file_matches + assert len(file_matches) == 2 + + +class TestFileGrouping: + """Tests for grouping matches by file.""" + + def test_matches_grouped_by_file(self): + """Matches are correctly grouped by filename.""" + content = """a.py:1:line 1 +b.py:2:line 2 +a.py:3:line 3 +c.py:4:line 4 +b.py:5:line 5 +a.py:6:line 6 +""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + + assert len(file_matches) == 3 + assert len(file_matches["a.py"].matches) == 3 + assert len(file_matches["b.py"].matches) == 2 + assert len(file_matches["c.py"].matches) == 1 + + def test_file_matches_first_property(self): + """FileMatches.first returns first match.""" + fm = FileMatches( + file="test.py", + matches=[ + SearchMatch(file="test.py", line_number=10, content="first"), + SearchMatch(file="test.py", line_number=20, content="second"), + ], + ) + + assert fm.first is not None + assert fm.first.line_number == 10 + assert fm.first.content == "first" + + def test_file_matches_last_property(self): + """FileMatches.last returns last match.""" + fm = FileMatches( + file="test.py", + matches=[ + SearchMatch(file="test.py", line_number=10, content="first"), + SearchMatch(file="test.py", line_number=20, content="last"), + ], + ) + + assert fm.last is not None + assert fm.last.line_number == 20 + assert fm.last.content == "last" + + def test_file_matches_empty(self): + """FileMatches with no matches handles first/last.""" + fm = FileMatches(file="test.py", matches=[]) + + assert fm.first is None + assert fm.last is None + + +class TestMatchScoring: + """Tests for match relevance scoring.""" + + def test_score_context_word_overlap(self): + """Matches containing context words get higher scores.""" + content = """src/main.py:10:def process_data(): +src/main.py:20:def calculate_result(): +src/main.py:30:def handle_error(): +""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + compressor._score_matches(file_matches, context="error handling") + + matches = file_matches["src/main.py"].matches + error_match = next(m for m in matches if "error" in m.content) + data_match = next(m for m in matches if "data" in m.content) + + # Error match should score higher with "error" context + assert error_match.score > data_match.score + + def test_score_error_patterns_boosted(self): + """Error/exception patterns get boosted scores.""" + content = """src/main.py:10:def normal_function(): +src/main.py:20:raise ValueError("error occurred") +src/main.py:30:# TODO: fix this +""" + compressor = SearchCompressor(config=SearchCompressorConfig(boost_errors=True)) + file_matches = compressor._parse_search_results(content) + compressor._score_matches(file_matches, context="") + + matches = file_matches["src/main.py"].matches + error_match = next(m for m in matches if "error" in m.content.lower()) + normal_match = next(m for m in matches if "normal" in m.content) + + assert error_match.score > normal_match.score + + def test_score_warning_patterns(self): + """Warning patterns get boosted scores.""" + content = """src/main.py:10:def normal(): +src/main.py:20:# WARNING: deprecated +""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + compressor._score_matches(file_matches, context="") + + matches = file_matches["src/main.py"].matches + warning_match = next(m for m in matches if "WARNING" in m.content) + normal_match = next(m for m in matches if "normal" in m.content) + + assert warning_match.score > normal_match.score + + def test_score_todo_patterns(self): + """TODO/FIXME patterns get boosted scores.""" + content = """src/main.py:10:def normal(): +src/main.py:20:# FIXME: this needs work +src/main.py:30:# TODO: implement later +""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + compressor._score_matches(file_matches, context="") + + matches = file_matches["src/main.py"].matches + fixme_match = next(m for m in matches if "FIXME" in m.content) + normal_match = next(m for m in matches if "normal" in m.content) + + assert fixme_match.score > normal_match.score + + def test_score_context_keywords_config(self): + """context_keywords configuration boosts matching lines.""" + content = """src/main.py:10:def auth_handler(): +src/main.py:20:def data_processor(): +""" + config = SearchCompressorConfig(context_keywords=["auth", "security"]) + compressor = SearchCompressor(config=config) + file_matches = compressor._parse_search_results(content) + compressor._score_matches(file_matches, context="") + + matches = file_matches["src/main.py"].matches + auth_match = next(m for m in matches if "auth" in m.content) + data_match = next(m for m in matches if "data" in m.content) + + assert auth_match.score > data_match.score + + def test_score_capped_at_one(self): + """Scores are capped at 1.0.""" + content = """src/main.py:10:ERROR FATAL exception fail warning TODO FIXME +""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + compressor._score_matches(file_matches, context="error fatal exception") + + match = file_matches["src/main.py"].matches[0] + assert match.score <= 1.0 + + +class TestMatchSelection: + """Tests for selecting which matches to keep.""" + + def test_keeps_first_and_last_by_default(self): + """First and last matches are kept by default.""" + content = "\n".join([f"src/file.py:{i}:line {i}" for i in range(1, 101)]) + + compressor = SearchCompressor( + config=SearchCompressorConfig( + always_keep_first=True, + always_keep_last=True, + max_matches_per_file=5, + ) + ) + result = compressor.compress(content) + + assert "src/file.py:1:line 1" in result.compressed + assert "src/file.py:100:line 100" in result.compressed + + def test_respects_max_matches_per_file(self): + """max_matches_per_file limits matches per file.""" + content = "\n".join([f"src/file.py:{i}:line {i}" for i in range(1, 51)]) + + compressor = SearchCompressor( + config=SearchCompressorConfig( + max_matches_per_file=3, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Should have at most 3 matches + summary + file_lines = [ + line for line in result.compressed.split("\n") if line.startswith("src/file.py:") + ] + assert len(file_lines) <= 3 + + def test_respects_max_total_matches(self): + """max_total_matches limits total output.""" + # Create matches across many files + lines = [] + for f in range(20): + for i in range(10): + lines.append(f"src/file{f}.py:{i}:line content") + content = "\n".join(lines) + + compressor = SearchCompressor( + config=SearchCompressorConfig( + max_total_matches=15, + max_files=20, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Count actual match lines (not summaries) + match_lines = [ + line for line in result.compressed.split("\n") if line and not line.startswith("[") + ] + assert len(match_lines) <= 15 + + def test_respects_max_files(self): + """max_files limits number of files in output.""" + # Create matches in many files + lines = [] + for f in range(30): + lines.append(f"src/file{f}.py:1:content") + content = "\n".join(lines) + + compressor = SearchCompressor( + config=SearchCompressorConfig( + max_files=5, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Count unique files in output + output_files = set() + for line in result.compressed.split("\n"): + if ":" in line and not line.startswith("["): + parts = line.split(":") + if len(parts) >= 2: + output_files.add(parts[0]) + + assert len(output_files) <= 5 + + def test_high_scoring_files_selected_first(self): + """Files with higher-scoring matches are selected first.""" + content = """normal/file.py:1:regular content +important/file.py:1:ERROR critical failure +another/file.py:1:some code here +""" + compressor = SearchCompressor( + config=SearchCompressorConfig( + max_files=1, + boost_errors=True, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # File with ERROR should be selected + assert "important/file.py" in result.compressed + + def test_output_sorted_by_line_number(self): + """Matches in output are sorted by line number within file.""" + content = """src/file.py:50:middle line +src/file.py:10:first line +src/file.py:90:last line +""" + compressor = SearchCompressor() + result = compressor.compress(content) + + lines = result.compressed.split("\n") + line_numbers = [] + for line in lines: + if line.startswith("src/file.py:"): + parts = line.split(":") + if len(parts) >= 2 and parts[1].isdigit(): + line_numbers.append(int(parts[1])) + + assert line_numbers == sorted(line_numbers) + + +class TestCompressionBehavior: + """Tests for overall compression behavior.""" + + def test_small_results_unchanged(self): + """Small results pass through unchanged.""" + content = "src/file.py:1:def foo():\nsrc/file.py:2: pass" + + compressor = SearchCompressor() + result = compressor.compress(content) + + assert result.compression_ratio == 1.0 + assert result.compressed == content + + def test_empty_input_handled(self): + """Empty input is handled gracefully.""" + compressor = SearchCompressor() + result = compressor.compress("") + + assert result.compressed == "" + assert result.original_match_count == 0 + assert result.compression_ratio == 1.0 + + def test_compression_adds_summary(self): + """Compression adds summary for omitted matches.""" + content = "\n".join([f"src/file.py:{i}:line {i}" for i in range(1, 51)]) + + compressor = SearchCompressor( + config=SearchCompressorConfig( + max_matches_per_file=3, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Should have summary about omitted matches + assert "[... and" in result.compressed + assert "more matches" in result.compressed + + def test_compression_ratio_calculated(self): + """Compression ratio is calculated correctly.""" + content = "\n".join([f"src/file.py:{i}:line {i}" for i in range(1, 101)]) + + compressor = SearchCompressor( + config=SearchCompressorConfig( + max_matches_per_file=5, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # Ratio should be less than 1.0 for compression + assert result.compression_ratio < 1.0 + + +class TestSearchCompressionResult: + """Tests for SearchCompressionResult dataclass.""" + + def test_tokens_saved_estimate(self): + """Token savings estimation works correctly.""" + original = "a" * 400 # ~100 tokens + compressed = "b" * 40 # ~10 tokens + + result = SearchCompressionResult( + compressed=compressed, + original=original, + original_match_count=100, + compressed_match_count=10, + files_affected=5, + compression_ratio=0.1, + ) + + # (400 - 40) / 4 = 90 tokens saved + assert result.tokens_saved_estimate == 90 + + def test_matches_omitted_property(self): + """matches_omitted property calculates correctly.""" + result = SearchCompressionResult( + compressed="test", + original="original", + original_match_count=100, + compressed_match_count=15, + files_affected=10, + compression_ratio=0.15, + ) + + assert result.matches_omitted == 85 + + def test_default_summaries_empty(self): + """Default summaries is empty dict.""" + result = SearchCompressionResult( + compressed="test", + original="original", + original_match_count=1, + compressed_match_count=1, + files_affected=1, + compression_ratio=1.0, + ) + + assert result.summaries == {} + + +class TestEdgeCases: + """Tests for edge cases and boundary conditions.""" + + def test_single_match_passthrough(self): + """Single match passes through unchanged.""" + content = "src/file.py:10:single match" + + compressor = SearchCompressor() + result = compressor.compress(content) + + assert result.compressed == content + assert result.original_match_count == 1 + assert result.compressed_match_count == 1 + + def test_unicode_content(self): + """Unicode characters in content are handled.""" + content = """src/main.py:10:msg = "こんにちは" +src/main.py:20:emoji = "🎉" +src/main.py:30:umlaut = "über" +""" + compressor = SearchCompressor() + result = compressor.compress(content) + + assert "こんにちは" in result.compressed + assert "🎉" in result.compressed + assert "über" in result.compressed + + def test_very_long_lines(self): + """Very long content lines are handled.""" + long_content = "x" * 10000 + content = f"src/file.py:1:{long_content}" + + compressor = SearchCompressor() + result = compressor.compress(content) + + assert len(result.compressed) > 0 + assert long_content in result.compressed + + def test_many_files_few_matches(self): + """Many files with one match each are handled.""" + lines = [f"src/file{i}.py:1:single match" for i in range(100)] + content = "\n".join(lines) + + compressor = SearchCompressor( + config=SearchCompressorConfig( + max_files=10, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + assert result.files_affected == 100 + # Output should be limited to max_files + output_files = set() + for line in result.compressed.split("\n"): + if ":" in line and not line.startswith("["): + parts = line.split(":") + if len(parts) >= 2: + output_files.add(parts[0]) + assert len(output_files) <= 10 + + def test_special_characters_in_path(self): + """Special characters in file paths are handled.""" + content = """src/my-file.py:10:content +src/my_file.py:20:content +src/my.file.py:30:content +src/file (1).py:40:content +""" + compressor = SearchCompressor() + result = compressor.compress(content) + + assert "my-file.py" in result.compressed + assert "my_file.py" in result.compressed + + def test_line_number_zero(self): + """Line number 0 is handled (edge case).""" + content = "src/file.py:0:line at position 0" + + compressor = SearchCompressor() + result = compressor.compress(content) + + assert ":0:" in result.compressed + + def test_negative_line_number_skipped(self): + """Negative line numbers don't match the pattern.""" + content = "src/file.py:-1:invalid" + + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + + # Pattern requires \d+ which is positive integers only + assert len(file_matches) == 0 + + +class TestContextIntegration: + """Tests for context-aware compression.""" + + def test_context_influences_selection(self): + """Context string influences which matches are selected.""" + lines = [] + for i in range(50): + lines.append(f"src/utils.py:{i}:def helper_{i}():") + + # Add some specific matches + lines.append("src/auth.py:100:def authenticate_user():") + lines.append("src/auth.py:200:def validate_token():") + + content = "\n".join(lines) + + compressor = SearchCompressor( + config=SearchCompressorConfig( + max_total_matches=5, + context_keywords=["auth", "token", "validate"], + enable_ccr=False, + ) + ) + result = compressor.compress(content, context="find authentication code") + + # Auth-related matches should be included + assert "authenticate" in result.compressed or "token" in result.compressed + + def test_short_context_words_ignored(self): + """Context words <= 2 chars are ignored for scoring.""" + content = """src/file.py:10:a = 1 +src/file.py:20:do something important +""" + compressor = SearchCompressor() + file_matches = compressor._parse_search_results(content) + compressor._score_matches(file_matches, context="a") + + # Short context word "a" shouldn't cause errors or abnormal scoring + matches = file_matches["src/file.py"].matches + assert all(m.score <= 1.0 for m in matches) + + +class TestOutputFormatting: + """Tests for output format and structure.""" + + def test_output_maintains_grep_format(self): + """Output maintains file:line:content format.""" + content = """src/file.py:10:def foo(): +src/file.py:20:def bar(): +""" + compressor = SearchCompressor() + result = compressor.compress(content) + + for line in result.compressed.split("\n"): + if line and not line.startswith("["): + assert line.count(":") >= 2 + parts = line.split(":", 2) + assert parts[1].isdigit() + + def test_summaries_track_omitted_per_file(self): + """Summaries dict tracks omissions per file.""" + content = "\n".join([f"src/file.py:{i}:line {i}" for i in range(1, 51)]) + + compressor = SearchCompressor( + config=SearchCompressorConfig( + max_matches_per_file=3, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + assert "src/file.py" in result.summaries + assert "more matches" in result.summaries["src/file.py"] + + def test_files_sorted_in_output(self): + """Files are sorted alphabetically in output.""" + content = """z_file.py:1:content +a_file.py:1:content +m_file.py:1:content +""" + compressor = SearchCompressor() + result = compressor.compress(content) + + lines = [ + line for line in result.compressed.split("\n") if line and not line.startswith("[") + ] + files = [line.split(":")[0] for line in lines] + + assert files == sorted(files) + + +class TestSearchMatchDataclass: + """Tests for SearchMatch dataclass.""" + + def test_default_score_zero(self): + """Default score is 0.0.""" + match = SearchMatch(file="test.py", line_number=1, content="test") + assert match.score == 0.0 + + def test_match_attributes(self): + """Match attributes are set correctly.""" + match = SearchMatch( + file="src/main.py", + line_number=42, + content="def process():", + score=0.8, + ) + + assert match.file == "src/main.py" + assert match.line_number == 42 + assert match.content == "def process():" + assert match.score == 0.8 + + +class TestConfigOptions: + """Tests for configuration options.""" + + def test_disable_keep_first(self): + """always_keep_first=False doesn't force first match.""" + content = "\n".join([f"src/file.py:{i}:line {i}" for i in range(1, 51)]) + + compressor = SearchCompressor( + config=SearchCompressorConfig( + always_keep_first=False, + always_keep_last=True, + max_matches_per_file=2, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # First line not guaranteed to be present + # But last should be + assert "src/file.py:50:line 50" in result.compressed + + def test_disable_keep_last(self): + """always_keep_last=False doesn't force last match.""" + content = "\n".join([f"src/file.py:{i}:line {i}" for i in range(1, 51)]) + + compressor = SearchCompressor( + config=SearchCompressorConfig( + always_keep_first=True, + always_keep_last=False, + max_matches_per_file=2, + enable_ccr=False, + ) + ) + result = compressor.compress(content) + + # First line should be present + assert "src/file.py:1:line 1" in result.compressed + + def test_disable_error_boost(self): + """boost_errors=False doesn't prioritize error patterns.""" + content = """src/file.py:1:ERROR critical failure +src/file.py:2:normal code line +""" + compressor = SearchCompressor( + config=SearchCompressorConfig( + boost_errors=False, + ) + ) + file_matches = compressor._parse_search_results(content) + compressor._score_matches(file_matches, context="") + + matches = file_matches["src/file.py"].matches + # Without boost, both should have similar (low) scores + error_match = next(m for m in matches if "ERROR" in m.content) + assert error_match.score == 0.0 # No boost applied + + def test_min_matches_for_ccr(self): + """min_matches_for_ccr threshold is respected.""" + content = "\n".join([f"src/file.py:{i}:line {i}" for i in range(1, 6)]) + + # With threshold of 10, CCR should not activate for 5 matches + compressor = SearchCompressor( + config=SearchCompressorConfig( + min_matches_for_ccr=10, + enable_ccr=True, + ) + ) + result = compressor.compress(content) + + assert result.cache_key is None