""" Advanced tests for ChatContextManager. Tests context window limits, summarization, and context building. """ from src.local_deep_research.chat.context import ChatContextManager class TestContextBuilding: """Tests for building research context.""" def test_build_research_context_structure(self): """Test that research context has expected structure.""" messages = [ {"role": "user", "content": "Test query", "message_type": "query"} ] manager = ChatContextManager("test-session", messages, {}) context = manager.build_research_context() expected_keys = [ "session_id", "original_query", "accumulated_findings", "past_findings", "key_entities", "topics", "is_multi_turn", "turn_count", ] for key in expected_keys: assert key in context, f"Missing key: {key}" def test_is_multi_turn_false_for_empty(self): """Test is_multi_turn is False for empty messages.""" manager = ChatContextManager("test-session", [], {}) context = manager.build_research_context() assert context["is_multi_turn"] is False def test_is_multi_turn_false_with_only_user_messages(self): """Test is_multi_turn is False when only user messages (no assistant response yet).""" messages = [ {"role": "user", "content": "Query", "message_type": "query"} ] manager = ChatContextManager("test-session", messages, {}) context = manager.build_research_context() assert context["is_multi_turn"] is False def test_is_multi_turn_true_with_assistant_response(self): """Test is_multi_turn is True when an assistant response exists.""" messages = [ {"role": "user", "content": "Query", "message_type": "query"}, { "role": "assistant", "content": "Response", "message_type": "response", }, ] manager = ChatContextManager("test-session", messages, {}) context = manager.build_research_context() assert context["is_multi_turn"] is True def test_turn_count_accuracy(self): """Test that turn_count matches message count.""" messages = [ {"role": "user", "content": f"Msg {i}", "message_type": "query"} for i in range(5) ] manager = ChatContextManager("test-session", messages, {}) context = manager.build_research_context() assert context["turn_count"] == 5 class TestFindingsExtraction: """Tests for extracting findings from history.""" def test_extract_findings_from_assistant_messages(self): """Test extracting findings from assistant messages.""" messages = [ {"role": "user", "content": "Query", "message_type": "query"}, { "role": "assistant", "content": "Finding 1", "message_type": "response", "research_id": "r1", }, { "role": "user", "content": "Follow up", "message_type": "followup", }, { "role": "assistant", "content": "Finding 2", "message_type": "response", "research_id": "r2", }, ] manager = ChatContextManager("test-session", messages, {}) findings = manager._extract_findings_from_history() assert "Finding 1" in findings assert "Finding 2" in findings def test_extract_findings_skips_user_messages(self): """Test that user messages are not included in findings.""" messages = [ {"role": "user", "content": "User query", "message_type": "query"}, { "role": "assistant", "content": "Assistant finding", "message_type": "response", "research_id": "r1", }, ] manager = ChatContextManager("test-session", messages, {}) findings = manager._extract_findings_from_history() assert "User query" not in findings assert "Assistant finding" in findings def test_extract_findings_limits_to_max(self): """Test that findings are limited to MAX_FINDINGS_TO_INCLUDE.""" messages = [] for i in range(10): messages.append( { "role": "assistant", "content": f"Finding {i}", "message_type": "response", "research_id": f"r{i}", } ) manager = ChatContextManager("test-session", messages, {}) findings = manager._extract_findings_from_history() # Should only have last MAX_FINDINGS_TO_INCLUDE assert "Finding 9" in findings assert "Finding 0" not in findings class TestContextUpdates: """Tests for extracting context updates.""" def test_extract_context_updates_creates_summary(self): """Test that context updates include summary creation.""" manager = ChatContextManager("test-session", [], {}) updates = manager.extract_context_updates( "New research content here with details" ) assert "summary_addition" in updates assert len(updates["summary_addition"]) > 0 def test_create_summary_truncates_long_content(self): """Test that summary creation truncates long content.""" manager = ChatContextManager("test-session", [], {}) long_content = "A" * 1000 summary = manager._create_summary(long_content) assert len(summary) <= 310 # 300 + "..." assert summary.endswith("...") def test_create_summary_handles_empty_content(self): """Test that empty content returns empty summary.""" manager = ChatContextManager("test-session", [], {}) summary = manager._create_summary("") assert summary == "" class TestStepMessageFiltering: """Tests that step messages are excluded from context building.""" def test_step_messages_excluded_from_turn_count(self): """Step messages should not count as conversation turns.""" messages = [ {"role": "user", "content": "Query", "message_type": "query"}, { "role": "assistant", "content": "Searching...", "message_type": "step", }, { "role": "assistant", "content": "Found results", "message_type": "response", "research_id": "r1", }, ] manager = ChatContextManager("test-session", messages, {}) context = manager.build_research_context() # step is filtered → only 2 messages visible assert context["turn_count"] == 2 def test_step_messages_excluded_from_is_multi_turn(self): """is_multi_turn should not be True from step messages alone.""" messages = [ {"role": "user", "content": "Query", "message_type": "query"}, { "role": "assistant", "content": "Searching...", "message_type": "step", }, ] manager = ChatContextManager("test-session", messages, {}) context = manager.build_research_context() # Only user message remains after filtering → not multi-turn assert context["is_multi_turn"] is False def test_step_messages_excluded_from_findings(self): """Step messages with research_id should not appear in findings.""" messages = [ { "role": "assistant", "content": "Searching now...", "message_type": "step", "research_id": "r1", }, { "role": "assistant", "content": "Real finding", "message_type": "response", "research_id": "r1", }, ] manager = ChatContextManager("test-session", messages, {}) findings = manager._extract_findings_from_history() assert "Searching now..." not in findings assert "Real finding" in findings def test_messages_without_message_type_pass_through(self): """Legacy messages with no message_type field are not filtered out.""" messages = [ {"role": "user", "content": "Old query"}, {"role": "assistant", "content": "Old response"}, ] manager = ChatContextManager("test-session", messages, {}) context = manager.build_research_context() assert context["turn_count"] == 2