""" Comprehensive tests for PrecisionExtractionHandler. Tests question type detection, various extraction methods, and edge cases. """ import pytest from unittest.mock import Mock class TestPrecisionExtractionHandlerInit: """Tests for PrecisionExtractionHandler initialization.""" def test_init_with_llm_only(self, mock_llm): """Test initialization with just an LLM.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) handler = PrecisionExtractionHandler(llm=mock_llm) assert handler.llm == mock_llm assert handler.settings_snapshot == {} def test_init_creates_answer_patterns(self, mock_llm): """Test initialization creates regex patterns.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) handler = PrecisionExtractionHandler(llm=mock_llm) assert "full_name" in handler.answer_patterns assert "year" in handler.answer_patterns assert "number" in handler.answer_patterns assert "dimension" in handler.answer_patterns assert "score" in handler.answer_patterns def test_init_with_settings_snapshot( self, mock_llm, settings_with_fact_checking ): """Test initialization with settings snapshot.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) handler = PrecisionExtractionHandler( llm=mock_llm, settings_snapshot=settings_with_fact_checking ) assert handler.settings_snapshot == settings_with_fact_checking class TestPrecisionExtractionIdentifyQuestionType: """Tests for _identify_question_type method.""" @pytest.fixture def handler(self, mock_llm): """Create a PrecisionExtractionHandler.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) return PrecisionExtractionHandler(llm=mock_llm) def test_identifies_full_name_question(self, handler): """Test detection of full name questions.""" assert ( handler._identify_question_type( "What is the full name of the founder?" ) == "full_name" ) def test_identifies_name_question_with_name_of(self, handler): """Test detection of name questions with 'name of'.""" assert ( handler._identify_question_type("What is the name of the CEO?") == "name" ) def test_identifies_name_question_with_who_was(self, handler): """Test detection of name questions with 'who was'.""" assert ( handler._identify_question_type("Who was the inventor?") == "name" ) def test_identifies_name_question_with_who_is(self, handler): """Test detection of name questions with 'who is'.""" assert ( handler._identify_question_type("Who is the current president?") == "name" ) def test_identifies_location_question_with_where(self, handler): """Test detection of location questions with 'where'.""" assert ( handler._identify_question_type("Where did it happen?") == "location" ) def test_identifies_location_question_with_city(self, handler): """Test detection of location questions mentioning city.""" assert ( handler._identify_question_type( "In which city was the company founded?" ) == "location" ) def test_identifies_location_question_with_country(self, handler): """Test detection of location questions mentioning country.""" assert ( handler._identify_question_type("What country is it located in?") == "location" ) def test_identifies_temporal_question_with_when(self, handler): """Test detection of temporal questions with 'when'.""" assert ( handler._identify_question_type("When was the company founded?") == "temporal" ) def test_identifies_temporal_question_with_year(self, handler): """Test detection of temporal questions with 'year'.""" assert ( handler._identify_question_type("In what year did the war begin?") == "temporal" ) def test_identifies_temporal_question_with_date(self, handler): """Test detection of temporal questions with 'date'.""" assert ( handler._identify_question_type("What is the date of the event?") == "temporal" ) def test_identifies_number_question_with_how_many(self, handler): """Test detection of number questions with 'how many'.""" assert ( handler._identify_question_type("How many employees work there?") == "number" ) def test_identifies_number_question_with_how_much(self, handler): """Test detection of number questions with 'how much'.""" assert ( handler._identify_question_type("How much did it cost?") == "number" ) def test_identifies_score_question(self, handler): """Test detection of score questions.""" assert ( handler._identify_question_type("What was the final score?") == "score" ) def test_identifies_score_question_with_result(self, handler): """Test detection of score questions with 'result'.""" assert ( handler._identify_question_type("What was the result of the match?") == "score" ) def test_identifies_dimension_question_with_height(self, handler): """Test detection of dimension questions with 'height'.""" assert ( handler._identify_question_type( "What is the height of the building?" ) == "dimension" ) def test_identifies_dimension_question_with_tall(self, handler): """Test detection of dimension questions with 'tall'.""" assert ( handler._identify_question_type("How tall is the tower?") == "dimension" ) def test_identifies_dimension_question_with_length(self, handler): """Test detection of dimension questions with 'length'.""" assert ( handler._identify_question_type("What is the length of the bridge?") == "dimension" ) def test_identifies_single_choice_question(self, handler): """Test detection of single choice questions.""" assert ( handler._identify_question_type("Which one of these is correct?") == "single_choice" ) def test_returns_general_for_ambiguous_question(self, handler): """Test that ambiguous questions return 'general'.""" assert ( handler._identify_question_type("What happened in the meeting?") == "general" ) def test_case_insensitive_detection(self, handler): """Test case-insensitive question type detection.""" assert ( handler._identify_question_type("WHAT IS THE FULL NAME?") == "full_name" ) assert ( handler._identify_question_type("WHERE DID IT HAPPEN?") == "location" ) class TestPrecisionExtractionRegexPatterns: """Tests for regex pattern matching.""" @pytest.fixture def handler(self, mock_llm): """Create a PrecisionExtractionHandler.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) return PrecisionExtractionHandler(llm=mock_llm) def test_full_name_pattern_matches_two_words(self, handler): """Test full name pattern matches two-word names.""" matches = handler.answer_patterns["full_name"].findall( "John Smith went to the store." ) assert "John Smith" in matches def test_full_name_pattern_matches_three_words(self, handler): """Test full name pattern matches three-word names.""" matches = handler.answer_patterns["full_name"].findall( "John Michael Smith was born." ) assert "John Michael Smith" in matches def test_year_pattern_matches_1900s(self, handler): """Test year pattern matches 1900s years.""" matches = handler.answer_patterns["year"].findall("Founded in 1995.") assert "1995" in matches def test_year_pattern_matches_2000s(self, handler): """Test year pattern matches 2000s years.""" matches = handler.answer_patterns["year"].findall("Launched in 2024.") assert "2024" in matches def test_number_pattern_matches_integer(self, handler): """Test number pattern matches integers.""" matches = handler.answer_patterns["number"].findall( "There are 42 employees." ) assert "42" in matches def test_number_pattern_matches_decimal(self, handler): """Test number pattern matches decimals.""" matches = handler.answer_patterns["number"].findall( "The cost is 19.99 dollars." ) assert "19.99" in matches def test_dimension_pattern_matches_meters(self, handler): """Test dimension pattern matches meters.""" matches = handler.answer_patterns["dimension"].findall( "The tower is 324 meters tall." ) assert ("324", "meters") in matches def test_dimension_pattern_matches_feet(self, handler): """Test dimension pattern matches feet.""" matches = handler.answer_patterns["dimension"].findall( "It is 100 feet high." ) assert ("100", "feet") in matches def test_dimension_pattern_matches_kg(self, handler): """Test dimension pattern matches kilograms.""" matches = handler.answer_patterns["dimension"].findall("Weighs 50 kg.") assert ("50", "kg") in matches def test_score_pattern_matches_hyphen_format(self, handler): """Test score pattern matches X-Y format.""" matches = handler.answer_patterns["score"].findall( "Final score was 3-2." ) assert ("3", "2") in matches def test_score_pattern_matches_dash_format(self, handler): """Test score pattern matches X–Y format (en-dash).""" matches = handler.answer_patterns["score"].findall( "Final score was 3–2." ) assert ("3", "2") in matches class TestPrecisionExtractionMethods: """Tests for individual extraction methods.""" @pytest.fixture def handler(self, mock_llm): """Create a PrecisionExtractionHandler.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) return PrecisionExtractionHandler(llm=mock_llm) def test_extract_full_name_with_llm(self, handler, mock_llm): """Test _extract_full_name calls LLM.""" mock_llm.invoke.return_value = Mock( content="Full name: John Michael Smith" ) result = handler._extract_full_name( "John Smith was a scientist.", "Who is it?", "Sources" ) mock_llm.invoke.assert_called_once() assert "John Michael Smith" in result def test_extract_full_name_fallback_to_longest(self, handler, mock_llm): """Test _extract_full_name falls back to longest name when LLM doesn't identify.""" mock_llm.invoke.return_value = Mock(content="No full name identified") content = "John Michael Smith and John Smith were scientists." result = handler._extract_full_name(content, "Who is it?", "Sources") # Should pick the longer name assert "John Michael Smith" in result def test_extract_full_name_handles_exception(self, handler, mock_llm): """Test _extract_full_name handles exceptions gracefully.""" mock_llm.invoke.side_effect = Exception("LLM error") result = handler._extract_full_name( "John Smith content", "Who?", "Sources" ) # Should return original content assert "John Smith content" == result def test_extract_best_name_returns_most_frequent(self, handler): """Test _extract_best_name returns most frequent name.""" content = "John Smith said hello. Smith was happy. John Smith left." result = handler._extract_best_name(content, "Who is it?", "Sources") assert "John Smith" in result def test_extract_dimension_with_llm(self, handler, mock_llm): """Test _extract_dimension calls LLM.""" mock_llm.invoke.return_value = Mock(content="324 meters") result = handler._extract_dimension( "The tower is 324 meters tall.", "How tall?", "Sources" ) mock_llm.invoke.assert_called_once() assert "324" in result def test_extract_dimension_fallback_to_pattern(self, handler, mock_llm): """Test _extract_dimension falls back to regex pattern.""" mock_llm.invoke.return_value = Mock(content="No clear answer") content = "The tower is 324 meters tall." result = handler._extract_dimension(content, "How tall?", "Sources") # Should still extract the measurement assert "324" in result or "meters" in result or content == result def test_extract_score_with_multiple_scores(self, handler, mock_llm): """Test _extract_score when multiple scores present.""" mock_llm.invoke.return_value = Mock(content="3-2") content = "Halftime was 1-1. Final score was 3-2." result = handler._extract_score( content, "What was the final score?", "Sources" ) mock_llm.invoke.assert_called_once() assert "3-2" in result def test_extract_score_no_scores_returns_content(self, handler, mock_llm): """Test _extract_score returns content when no scores found.""" content = "The game was exciting." result = handler._extract_score( content, "What was the score?", "Sources" ) assert result == content def test_extract_temporal_with_llm(self, handler, mock_llm): """Test _extract_temporal calls LLM.""" mock_llm.invoke.return_value = Mock(content="1998") content = "Founded in 1998, acquired in 2015." result = handler._extract_temporal( content, "When was it founded?", "Sources" ) mock_llm.invoke.assert_called_once() assert "1998" in result def test_extract_temporal_no_years_returns_content(self, handler, mock_llm): """Test _extract_temporal returns content when no years found.""" content = "It happened long ago." result = handler._extract_temporal(content, "When?", "Sources") assert result == content def test_extract_number_with_llm(self, handler, mock_llm): """Test _extract_number calls LLM.""" mock_llm.invoke.return_value = Mock(content="42") content = "There are 42 employees and 10 managers." result = handler._extract_number( content, "How many employees?", "Sources" ) mock_llm.invoke.assert_called_once() assert "42" in result def test_extract_single_answer_removes_alternatives( self, handler, mock_llm ): """Test _extract_single_answer removes comma/and/or alternatives.""" mock_llm.invoke.return_value = Mock( content="Option A, Option B, and Option C" ) result = handler._extract_single_answer( "Multiple options content", "Which one?", "Sources" ) # Should only have first option assert "Option A" in result assert "Option B" not in result.split(".")[0] def test_extract_key_facts_from_previous_knowledge(self, handler, mock_llm): """Test _extract_key_facts extracts from previous knowledge.""" mock_llm.invoke.return_value = Mock( content="Key facts: Founded 1998, Location: NYC" ) result = handler._extract_key_facts( "Previous knowledge about the company", "name" ) mock_llm.invoke.assert_called_once() assert len(result) <= 500 def test_extract_key_facts_handles_exception(self, handler, mock_llm): """Test _extract_key_facts handles exception gracefully.""" mock_llm.invoke.side_effect = Exception("LLM error") result = handler._extract_key_facts("Previous knowledge", "name") # Should return truncated previous knowledge assert "Previous knowledge" in result def test_extract_full_name_handles_string_response(self, handler, mock_llm): """_extract_full_name handles a raw string LLM return (str branch).""" mock_llm.invoke.return_value = "Full name: John Michael Smith" result = handler._extract_full_name( "John Smith was a scientist.", "Who is it?", "Sources" ) assert "John Michael Smith" in result def test_extract_single_answer_handles_string_response( self, handler, mock_llm ): """_extract_single_answer handles a raw string LLM return (str branch).""" mock_llm.invoke.return_value = "Option A, Option B, and Option C" result = handler._extract_single_answer( "Multiple options content", "Which one?", "Sources" ) assert result.startswith("Option A.") def test_extract_dimension_handles_string_response(self, handler, mock_llm): """_extract_dimension handles a raw string LLM return (str branch).""" mock_llm.invoke.return_value = "324 meters" result = handler._extract_dimension( "The tower is 324 meters tall.", "How tall?", "Sources" ) assert result.startswith("324 meters.") def test_extract_score_handles_string_response(self, handler, mock_llm): """_extract_score handles a raw string LLM return (str branch).""" mock_llm.invoke.return_value = "3-2" result = handler._extract_score( "Halftime was 1-1. Final score was 3-2.", "What was the final score?", "Sources", ) assert result.startswith("3-2.") def test_extract_temporal_handles_string_response(self, handler, mock_llm): """_extract_temporal handles a raw string LLM return (str branch).""" mock_llm.invoke.return_value = "1998" result = handler._extract_temporal( "Founded in 1998, acquired in 2015.", "When was it founded?", "Sources", ) assert result.startswith("1998.") def test_extract_number_handles_string_response(self, handler, mock_llm): """_extract_number handles a raw string LLM return (str branch).""" mock_llm.invoke.return_value = "42" result = handler._extract_number( "There are 42 employees and 10 managers.", "How many employees?", "Sources", ) assert result.startswith("42.") def test_extract_key_facts_handles_string_response(self, handler, mock_llm): """_extract_key_facts handles a raw string LLM return (str branch).""" mock_llm.invoke.return_value = "Key facts: Founded 1998, Location: NYC" result = handler._extract_key_facts( "Previous knowledge about the company", "name" ) assert "Founded 1998" in result assert len(result) <= 500 class TestPrecisionExtractionAnalyzeInitial: """Tests for analyze_initial method.""" def test_analyze_initial_returns_content_and_documents( self, mock_llm, sample_search_results ): """Test analyze_initial returns proper structure.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) handler = PrecisionExtractionHandler(llm=mock_llm) result = handler.analyze_initial( "What is the topic?", sample_search_results ) assert "content" in result assert "documents" in result assert len(result["documents"]) == 3 def test_analyze_initial_prompt_includes_question_type( self, mock_llm, sample_search_results ): """Test analyze_initial prompt includes detected question type.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) mock_llm.invoke.return_value = Mock( content="The full name is John Smith. [1]" ) handler = PrecisionExtractionHandler(llm=mock_llm) handler.analyze_initial( "What is the full name of the founder?", sample_search_results ) call_args = mock_llm.invoke.call_args_list[0][0][0] assert "full_name" in call_args def test_analyze_initial_prompt_includes_precision_instructions( self, mock_llm, sample_search_results ): """Test analyze_initial prompt includes precision extraction instructions.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) mock_llm.invoke.return_value = Mock(content="Answer [1]") handler = PrecisionExtractionHandler(llm=mock_llm) handler.analyze_initial("What is it?", sample_search_results) call_args = mock_llm.invoke.call_args_list[0][0][0] assert "PRECISION" in call_args assert "EXACT answer" in call_args def test_analyze_initial_applies_extraction_for_full_name( self, mock_llm, name_search_results ): """Test analyze_initial applies extraction for full_name questions.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) # First call returns LLM response, second call extracts name mock_llm.invoke.side_effect = [ Mock(content="The person is John Smith."), Mock(content="Full name: John Michael William Smith"), ] handler = PrecisionExtractionHandler(llm=mock_llm) handler.analyze_initial("What is the full name?", name_search_results) # Should have called extraction assert mock_llm.invoke.call_count >= 1 def test_analyze_initial_handles_string_response( self, mock_llm_string_response, sample_search_results ): """Test analyze_initial handles string LLM response.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) handler = PrecisionExtractionHandler(llm=mock_llm_string_response) result = handler.analyze_initial("What is it?", sample_search_results) assert "content" in result def test_analyze_initial_strips_think_tags( self, mock_llm, sample_search_results ): """analyze_initial strips reasoning (general question, no extraction).""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) mock_llm.invoke.return_value = Mock( content="decidingGeneral answer with citation [1]." ) handler = PrecisionExtractionHandler(llm=mock_llm) # general-type question -> _apply_precision_extraction leaves content as-is result = handler.analyze_initial( "Tell me about the topic.", sample_search_results ) assert "" not in result["content"] assert "deciding" not in result["content"] assert result["content"] == "General answer with citation [1]." class TestPrecisionExtractionAnalyzeFollowup: """Tests for analyze_followup method.""" def test_analyze_followup_returns_content_and_documents( self, mock_llm, sample_search_results, sample_previous_knowledge ): """Test analyze_followup returns proper structure.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) handler = PrecisionExtractionHandler(llm=mock_llm) result = handler.analyze_followup( "Follow-up question?", sample_search_results, sample_previous_knowledge, nr_of_links=3, ) assert "content" in result assert "documents" in result def test_analyze_followup_extracts_key_facts( self, mock_llm, sample_search_results, sample_previous_knowledge ): """Test analyze_followup extracts key facts from previous knowledge.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) mock_llm.invoke.return_value = Mock(content="Key facts extracted") handler = PrecisionExtractionHandler(llm=mock_llm) handler.analyze_followup( "What year was it founded?", sample_search_results, sample_previous_knowledge, nr_of_links=0, ) # Should have called LLM for key facts extraction first_call = mock_llm.invoke.call_args_list[0][0][0] assert ( "key facts" in first_call.lower() or "extract" in first_call.lower() ) def test_analyze_followup_applies_nr_of_links_offset( self, mock_llm, sample_search_results, sample_previous_knowledge ): """Test analyze_followup applies nr_of_links offset to document indices.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) handler = PrecisionExtractionHandler(llm=mock_llm) result = handler.analyze_followup( "Question?", sample_search_results, sample_previous_knowledge, nr_of_links=10, ) # First document should have index 11 (10 + 1) assert result["documents"][0].metadata["index"] == 11 def test_analyze_followup_handles_string_response( self, mock_llm_string_response, sample_search_results, sample_previous_knowledge, ): """analyze_followup returns raw string content for a general-type question.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) handler = PrecisionExtractionHandler(llm=mock_llm_string_response) result = handler.analyze_followup( "Tell me more about the topic.", sample_search_results, sample_previous_knowledge, nr_of_links=0, ) # general-type question -> precision extraction leaves content unchanged assert result["content"] == "Test string response with citation [1]." assert "documents" in result class TestPrecisionExtractionEdgeCases: """Tests for edge cases and error handling.""" def test_empty_search_results(self, mock_llm, empty_search_results): """Test handling of empty search results.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) handler = PrecisionExtractionHandler(llm=mock_llm) result = handler.analyze_initial("What is it?", empty_search_results) assert result["documents"] == [] def test_string_search_results(self, mock_llm, string_search_results): """Test handling of string search results.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) handler = PrecisionExtractionHandler(llm=mock_llm) result = handler.analyze_initial("What is it?", string_search_results) assert result["documents"] == [] def test_with_output_instructions( self, mock_llm, sample_search_results, settings_with_output_instructions ): """Test that output instructions are included in prompts.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) mock_llm.invoke.return_value = Mock(content="Response") handler = PrecisionExtractionHandler( llm=mock_llm, settings_snapshot=settings_with_output_instructions ) handler.analyze_initial("What is it?", sample_search_results) call_args = mock_llm.invoke.call_args_list[0][0][0] assert "formal academic English" in call_args def test_general_question_type_no_extraction( self, mock_llm, sample_search_results ): """Test that general question type doesn't apply special extraction.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) mock_llm.invoke.return_value = Mock(content="General response") handler = PrecisionExtractionHandler(llm=mock_llm) result = handler.analyze_initial( "What happened?", sample_search_results ) # Should just return the LLM response without extraction modifications assert result["content"] == "General response" def test_apply_precision_extraction_returns_content_for_unknown_type( self, mock_llm ): """Test _apply_precision_extraction returns content for unknown question type.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) handler = PrecisionExtractionHandler(llm=mock_llm) result = handler._apply_precision_extraction( "Original content", "Query?", "unknown_type", "Sources" ) assert result == "Original content" class TestPrecisionExtractionDimensionContextAwareness: """Tests for dimension extraction context awareness.""" @pytest.fixture def handler(self, mock_llm): """Create a PrecisionExtractionHandler.""" from local_deep_research.citation_handlers.precision_extraction_handler import ( PrecisionExtractionHandler, ) return PrecisionExtractionHandler(llm=mock_llm) def test_dimension_types_include_height(self, handler): """Test dimension type detection includes height keywords.""" question_type = handler._identify_question_type( "What is the height of the building?" ) assert question_type == "dimension" def test_dimension_types_include_size(self, handler): """Test dimension type detection includes size keywords.""" question_type = handler._identify_question_type( "What is the size of the building?" ) assert question_type == "dimension" def test_dimension_types_include_length(self, handler): """Test dimension type detection includes length keywords.""" question_type = handler._identify_question_type( "What is the length of the bridge?" ) assert question_type == "dimension" def test_dimension_types_include_width(self, handler): """Test dimension type detection includes width keywords.""" question_type = handler._identify_question_type( "How wide is the river?" ) assert question_type == "dimension"