learningcircuit--local-deep-research/tests/advanced_search_system/knowledge/test_followup_context_manager.py

"""
Tests for FollowUpContextHandler.

Tests cover:
- Initialization with model and settings
- Context building from research data
- Findings extraction from various data formats
- Source extraction and deduplication
- Entity extraction via LLM
- Summary generation for context and prompts
- Gap identification
- Settings snapshot formatting
- LLM context formatting
"""

from unittest.mock import Mock

import pytest

from local_deep_research.advanced_search_system.knowledge.followup_context_manager import (
    FollowUpContextHandler,
)


class TestFollowUpContextHandlerInit:
    """Tests for FollowUpContextHandler initialization."""

    def test_init_stores_model(self):
        """Handler stores the model reference."""
        mock_model = Mock()
        handler = FollowUpContextHandler(mock_model)
        assert handler.model is mock_model

    def test_init_with_settings_snapshot(self):
        """Handler stores settings snapshot."""
        mock_model = Mock()
        settings = {"key": "value"}
        handler = FollowUpContextHandler(mock_model, settings_snapshot=settings)
        assert handler.settings_snapshot == settings

    def test_init_without_settings_snapshot(self):
        """Handler initializes empty settings snapshot when None."""
        mock_model = Mock()
        handler = FollowUpContextHandler(mock_model, settings_snapshot=None)
        assert handler.settings_snapshot == {}

    def test_init_creates_empty_cache(self):
        """Handler initializes empty research cache."""
        mock_model = Mock()
        handler = FollowUpContextHandler(mock_model)
        assert handler.past_research_cache == {}


class TestBuildContext:
    """Tests for build_context method."""

    @pytest.fixture
    def mock_model(self):
        """Create mock model."""
        mock = Mock()
        mock.invoke.return_value = Mock(
            content="Entity1\nEntity2\nSummary text"
        )
        return mock

    @pytest.fixture
    def handler(self, mock_model):
        """Create handler instance."""
        return FollowUpContextHandler(mock_model)

    @pytest.fixture
    def sample_research_data(self):
        """Create sample research data."""
        return {
            "research_id": "research_123",
            "query": "What is machine learning?",
            "formatted_findings": "ML is a field of AI...",
            "report_content": "Detailed report about ML...",
            "resources": [
                {"url": "https://example.com/1", "title": "Source 1"}
            ],
            "all_links_of_system": [
                {"url": "https://example.com/2", "title": "Source 2"}
            ],
            "strategy": "detailed",
            "mode": "research",
            "created_at": "2024-01-14",
            "research_meta": {"duration": 120},
        }

    def test_build_context_returns_dict(self, handler, sample_research_data):
        """build_context returns a dictionary."""
        result = handler.build_context(
            sample_research_data, "What are ML applications?"
        )
        assert isinstance(result, dict)

    def test_build_context_surfaces_chat_past_findings(self, handler):
        """End-to-end of the fix: a chat-style context (only past_findings,
        no formatted_findings/report_content) surfaces its summary as
        past_findings, so it reaches generate_contextualized_query instead of
        being replaced by the 'No previous findings available' sentinel."""
        chat_context = {
            "research_id": "r1",
            "original_query": "What is machine learning?",
            "past_findings": "Summary of the conversation, focused on cost.",
        }
        result = handler.build_context(chat_context, "How much does it cost?")
        assert (
            result["past_findings"]
            == "Summary of the conversation, focused on cost."
        )

    def test_build_context_includes_parent_research_id(
        self, handler, sample_research_data
    ):
        """Context includes parent_research_id."""
        result = handler.build_context(
            sample_research_data, "Follow-up question?"
        )
        assert result["parent_research_id"] == "research_123"

    def test_build_context_includes_original_query(
        self, handler, sample_research_data
    ):
        """Context includes original_query."""
        result = handler.build_context(
            sample_research_data, "Follow-up question?"
        )
        assert result["original_query"] == "What is machine learning?"

    def test_build_context_includes_follow_up_query(
        self, handler, sample_research_data
    ):
        """Context includes follow_up_query."""
        result = handler.build_context(
            sample_research_data, "What are ML applications?"
        )
        assert result["follow_up_query"] == "What are ML applications?"

    def test_build_context_includes_all_keys(
        self, handler, sample_research_data
    ):
        """Context includes all expected keys."""
        result = handler.build_context(sample_research_data, "Follow-up?")
        expected_keys = [
            "parent_research_id",
            "original_query",
            "follow_up_query",
            "past_findings",
            "past_sources",
            "key_entities",
            "summary",
            "report_content",
            "formatted_findings",
            "all_links_of_system",
            "metadata",
        ]
        for key in expected_keys:
            assert key in result


class TestExtractFindings:
    """Tests for _extract_findings method."""

    @pytest.fixture
    def handler(self):
        """Create handler without model."""
        return FollowUpContextHandler(None)

    def test_extract_from_formatted_findings(self, handler):
        """Extract from formatted_findings field."""
        data = {"formatted_findings": "Formatted findings text"}
        result = handler._extract_findings(data)
        assert "Formatted findings text" in result

    def test_extract_from_report_content_when_no_formatted(self, handler):
        """Extract from report_content when no formatted_findings."""
        data = {"report_content": "Report content text here"}
        result = handler._extract_findings(data)
        assert "Report content text" in result

    def test_report_content_truncated_to_2000(self, handler):
        """Report content is truncated to 2000 characters."""
        long_content = "A" * 5000
        data = {"report_content": long_content}
        result = handler._extract_findings(data)
        assert len(result) == 2000

    def test_returns_default_when_no_findings(self, handler):
        """Returns default message when no findings available."""
        data = {}
        result = handler._extract_findings(data)
        assert result == "No previous findings available"

    def test_honors_chat_past_findings_when_no_report(self, handler):
        """Multi-turn chat supplies prior findings under ``past_findings``
        (it has no formatted_findings/report_content); that must be used as
        the findings rather than the "nothing available" sentinel."""
        data = {"past_findings": "Condensed summary of the conversation."}
        result = handler._extract_findings(data)
        assert result == "Condensed summary of the conversation."

    def test_report_content_takes_precedence_over_past_findings(self, handler):
        """Non-chat callers populate report_content; past_findings must not
        override it (keeps the standalone follow-up path unchanged)."""
        data = {
            "report_content": "Full parent report.",
            "past_findings": "should not be used",
        }
        result = handler._extract_findings(data)
        assert "Full parent report." in result
        assert "should not be used" not in result

    def test_prefers_formatted_over_report(self, handler):
        """Prefers formatted_findings over report_content."""
        data = {
            "formatted_findings": "Formatted text",
            "report_content": "Report text",
        }
        result = handler._extract_findings(data)
        assert "Formatted text" in result


class TestExtractSources:
    """Tests for _extract_sources method."""

    @pytest.fixture
    def handler(self):
        """Create handler without model."""
        return FollowUpContextHandler(None)

    def test_extract_from_resources(self, handler):
        """Extract sources from resources field."""
        data = {"resources": [{"url": "https://a.com", "title": "A"}]}
        result = handler._extract_sources(data)
        assert len(result) == 1
        assert result[0]["url"] == "https://a.com"

    def test_extract_from_all_links_of_system(self, handler):
        """Extract sources from all_links_of_system field."""
        data = {"all_links_of_system": [{"url": "https://b.com", "title": "B"}]}
        result = handler._extract_sources(data)
        assert len(result) == 1

    def test_extract_from_past_links(self, handler):
        """Extract sources from past_links field."""
        data = {"past_links": [{"url": "https://c.com", "title": "C"}]}
        result = handler._extract_sources(data)
        assert len(result) == 1

    def test_deduplicates_by_url(self, handler):
        """Deduplicates sources by URL."""
        data = {
            "resources": [{"url": "https://a.com", "title": "A1"}],
            "all_links_of_system": [{"url": "https://a.com", "title": "A2"}],
        }
        result = handler._extract_sources(data)
        assert len(result) == 1

    def test_includes_sources_without_url(self, handler):
        """Includes sources that have no URL."""
        data = {"resources": [{"title": "No URL source"}]}
        result = handler._extract_sources(data)
        assert len(result) == 1

    def test_combines_all_source_fields(self, handler):
        """Combines sources from all fields."""
        data = {
            "resources": [{"url": "https://a.com", "title": "A"}],
            "all_links_of_system": [{"url": "https://b.com", "title": "B"}],
            "past_links": [{"url": "https://c.com", "title": "C"}],
        }
        result = handler._extract_sources(data)
        assert len(result) == 3

    def test_returns_empty_list_when_no_sources(self, handler):
        """Returns empty list when no sources available."""
        data = {}
        result = handler._extract_sources(data)
        assert result == []


class TestExtractEntities:
    """Tests for _extract_entities method."""

    @pytest.fixture
    def mock_model(self):
        """Create mock model."""
        mock = Mock()
        mock.invoke.return_value = Mock(content="Entity1\nEntity2\nEntity3")
        return mock

    @pytest.fixture
    def handler(self, mock_model):
        """Create handler with mock model."""
        return FollowUpContextHandler(mock_model)

    def test_extracts_entities_from_findings(self, handler):
        """Extracts entities using LLM."""
        data = {"formatted_findings": "Research findings about Python and ML"}
        result = handler._extract_entities(data)
        assert len(result) == 3
        assert "Entity1" in result

    def test_limits_to_10_entities(self, handler, mock_model):
        """Limits entities to 10."""
        mock_model.invoke.return_value = Mock(
            content="\n".join([f"Entity{i}" for i in range(15)])
        )
        data = {"formatted_findings": "Some findings"}
        result = handler._extract_entities(data)
        assert len(result) <= 10

    def test_returns_empty_when_no_findings_and_no_model(self):
        """Returns empty list when no findings and no model."""
        handler = FollowUpContextHandler(None)
        data = {}
        result = handler._extract_entities(data)
        assert result == []

    def test_returns_empty_when_no_model(self):
        """Returns empty list when no model available."""
        handler = FollowUpContextHandler(None)
        data = {"formatted_findings": "Some findings"}
        result = handler._extract_entities(data)
        assert result == []

    def test_handles_llm_exception(self, handler, mock_model):
        """Handles LLM exception gracefully."""
        mock_model.invoke.side_effect = RuntimeError("Connection failed")
        data = {"formatted_findings": "Some findings"}
        result = handler._extract_entities(data)
        assert result == []


class TestCreateSummary:
    """Tests for _create_summary method."""

    @pytest.fixture
    def mock_model(self):
        """Create mock model."""
        mock = Mock()
        mock.invoke.return_value = Mock(content="Summary of findings")
        return mock

    @pytest.fixture
    def handler(self, mock_model):
        """Create handler with mock model."""
        return FollowUpContextHandler(mock_model)

    def test_creates_summary_with_llm(self, handler):
        """Creates summary using LLM."""
        data = {
            "formatted_findings": "Research findings",
            "query": "Original question?",
        }
        handler._create_summary(data, "Follow-up question?")
        assert handler.model.invoke.called

    def test_includes_original_query_in_prompt(self, handler):
        """Includes original query in prompt."""
        data = {
            "formatted_findings": "Research findings",
            "query": "Original question?",
        }
        handler._create_summary(data, "Follow-up?")
        call_args = handler.model.invoke.call_args[0][0]
        assert "Original question?" in call_args


class TestSummarizeForFollowup:
    """Tests for summarize_for_followup method."""

    @pytest.fixture
    def mock_model(self):
        """Create mock model."""
        mock = Mock()
        mock.invoke.return_value = Mock(content="Concise summary")
        return mock

    @pytest.fixture
    def handler(self, mock_model):
        """Create handler with mock model."""
        return FollowUpContextHandler(mock_model)

    def test_summarize_returns_string(self, handler):
        """summarize_for_followup returns a string."""
        result = handler.summarize_for_followup(
            "Long findings text...", "Follow-up question?"
        )
        assert isinstance(result, str)

    def test_summarize_respects_max_length(self, handler, mock_model):
        """Result respects max_length constraint."""
        mock_model.invoke.return_value = Mock(content="A" * 2000)
        result = handler.summarize_for_followup(
            "Findings", "Question?", max_length=500
        )
        assert len(result) <= 503  # 500 + "..."

    def test_short_findings_returned_as_is(self, handler):
        """Short findings are returned as-is without LLM call."""
        result = handler.summarize_for_followup(
            "Short text", "Question?", max_length=1000
        )
        # Model should not be called for short text
        assert result == "Short text"


class TestGenerateSummary:
    """Tests for _generate_summary method."""

    @pytest.fixture
    def mock_model(self):
        """Create mock model."""
        mock = Mock()
        mock.invoke.return_value = Mock(content="Generated summary")
        return mock

    @pytest.fixture
    def handler(self, mock_model):
        """Create handler with mock model."""
        return FollowUpContextHandler(mock_model)

    def test_returns_empty_for_empty_findings(self, handler):
        """Returns empty string for empty findings."""
        result = handler._generate_summary("", "Query")
        assert result == ""

    def test_returns_findings_when_short_enough(self, handler):
        """Returns findings as-is when within max_length."""
        result = handler._generate_summary(
            "Short findings", "Query", max_length=1000
        )
        assert result == "Short findings"

    def test_context_purpose_includes_original_query(self, handler):
        """Context purpose prompt includes original query."""
        handler._generate_summary(
            "Long " * 500,  # Force LLM call
            "Follow-up?",
            original_query="Original?",
            purpose="context",
        )
        call_args = handler.model.invoke.call_args[0][0]
        assert "Original?" in call_args

    def test_prompt_purpose_omits_original_query(self, handler):
        """Prompt purpose omits original query reference."""
        handler._generate_summary(
            "Long " * 500,
            "Follow-up?",
            original_query="Original?",
            purpose="prompt",
        )
        call_args = handler.model.invoke.call_args[0][0]
        # Prompt purpose uses simpler format
        assert "Follow-up?" in call_args

    def test_fallback_without_model(self):
        """Fallback truncation when no model available."""
        handler = FollowUpContextHandler(None)
        result = handler._generate_summary("A" * 1000, "Query", max_length=100)
        assert len(result) <= 103  # 100 + "..."

    def test_handles_llm_exception(self, handler, mock_model):
        """Handles LLM exception with truncation fallback."""
        mock_model.invoke.side_effect = RuntimeError("Error")
        result = handler._generate_summary("A" * 1000, "Query", max_length=100)
        assert len(result) <= 103


class TestIdentifyGaps:
    """Tests for identify_gaps method."""

    @pytest.fixture
    def mock_model(self):
        """Create mock model."""
        mock = Mock()
        mock.invoke.return_value = Mock(content="Gap 1\nGap 2\nGap 3")
        return mock

    @pytest.fixture
    def handler(self, mock_model):
        """Create handler with mock model."""
        return FollowUpContextHandler(mock_model)

    def test_identifies_gaps(self, handler):
        """Identifies gaps using LLM."""
        data = {"formatted_findings": "Some research findings"}
        result = handler.identify_gaps(data, "Follow-up question?")
        assert len(result) == 3
        assert "Gap 1" in result

    def test_limits_to_5_gaps(self, handler, mock_model):
        """Limits gaps to 5."""
        mock_model.invoke.return_value = Mock(
            content="\n".join([f"Gap {i}" for i in range(10)])
        )
        data = {"formatted_findings": "Findings"}
        result = handler.identify_gaps(data, "Question?")
        assert len(result) <= 5

    def test_returns_empty_when_no_findings_and_no_model(self):
        """Returns empty list when no findings and no model."""
        handler = FollowUpContextHandler(None)
        data = {}
        result = handler.identify_gaps(data, "Question?")
        assert result == []

    def test_returns_empty_when_no_model(self):
        """Returns empty list when no model."""
        handler = FollowUpContextHandler(None)
        data = {"formatted_findings": "Findings"}
        result = handler.identify_gaps(data, "Question?")
        assert result == []

    def test_handles_llm_exception(self, handler, mock_model):
        """Handles LLM exception gracefully."""
        mock_model.invoke.side_effect = RuntimeError("Error")
        data = {"formatted_findings": "Findings"}
        result = handler.identify_gaps(data, "Question?")
        assert result == []


class TestFormatForSettingsSnapshot:
    """Tests for format_for_settings_snapshot method."""

    @pytest.fixture
    def handler(self):
        """Create handler without model."""
        return FollowUpContextHandler(None)

    def test_returns_minimal_metadata(self, handler):
        """Returns minimal metadata dictionary."""
        context = {
            "parent_research_id": "research_123",
            "past_findings": "Some findings",
            "extra_data": "Should not be included",
        }
        result = handler.format_for_settings_snapshot(context)
        assert "followup_metadata" in result
        assert (
            result["followup_metadata"]["parent_research_id"] == "research_123"
        )
        assert result["followup_metadata"]["is_followup"] is True

    def test_has_context_true_when_findings_present(self, handler):
        """has_context is True when past_findings present."""
        context = {"past_findings": "Some findings"}
        result = handler.format_for_settings_snapshot(context)
        assert result["followup_metadata"]["has_context"] is True

    def test_has_context_false_when_no_findings(self, handler):
        """has_context is False when no past_findings."""
        context = {}
        result = handler.format_for_settings_snapshot(context)
        assert result["followup_metadata"]["has_context"] is False


class TestGetRelevantContextForLLM:
    """Tests for get_relevant_context_for_llm method."""

    @pytest.fixture
    def handler(self):
        """Create handler without model."""
        return FollowUpContextHandler(None)

    def test_includes_original_query(self, handler):
        """Includes original query in output."""
        context = {"original_query": "Original question?"}
        result = handler.get_relevant_context_for_llm(context)
        assert "Original question?" in result

    def test_includes_follow_up_query(self, handler):
        """Includes follow-up query in output."""
        context = {"follow_up_query": "Follow-up question?"}
        result = handler.get_relevant_context_for_llm(context)
        assert "Follow-up question?" in result

    def test_includes_summary(self, handler):
        """Includes summary in output."""
        context = {"summary": "Summary of findings"}
        result = handler.get_relevant_context_for_llm(context)
        assert "Summary of findings" in result

    def test_includes_key_entities(self, handler):
        """Includes key entities in output."""
        context = {"key_entities": ["Entity1", "Entity2", "Entity3"]}
        result = handler.get_relevant_context_for_llm(context)
        assert "Entity1" in result

    def test_limits_entities_to_5(self, handler):
        """Limits entities to first 5."""
        context = {"key_entities": [f"Entity{i}" for i in range(10)]}
        result = handler.get_relevant_context_for_llm(context)
        assert "Entity4" in result
        assert "Entity5" not in result or result.count("Entity") <= 5

    def test_includes_source_count(self, handler):
        """Includes source count in output."""
        context = {"past_sources": [{"url": "a"}, {"url": "b"}, {"url": "c"}]}
        result = handler.get_relevant_context_for_llm(context)
        assert "3" in result

    def test_truncates_to_max_tokens(self, handler):
        """Truncates output to approximate max_tokens."""
        context = {
            "original_query": "A" * 1000,
            "follow_up_query": "B" * 1000,
            "summary": "C" * 10000,
        }
        result = handler.get_relevant_context_for_llm(context, max_tokens=500)
        # max_tokens * 4 = 2000 chars
        assert len(result) <= 2003  # 2000 + "..."


class TestExtractMetadata:
    """Tests for _extract_metadata method."""

    @pytest.fixture
    def handler(self):
        """Create handler without model."""
        return FollowUpContextHandler(None)

    def test_extracts_strategy(self, handler):
        """Extracts strategy from research data."""
        data = {"strategy": "detailed"}
        result = handler._extract_metadata(data)
        assert result["strategy"] == "detailed"

    def test_extracts_mode(self, handler):
        """Extracts mode from research data."""
        data = {"mode": "research"}
        result = handler._extract_metadata(data)
        assert result["mode"] == "research"

    def test_extracts_created_at(self, handler):
        """Extracts created_at from research data."""
        data = {"created_at": "2024-01-14"}
        result = handler._extract_metadata(data)
        assert result["created_at"] == "2024-01-14"

    def test_extracts_research_meta(self, handler):
        """Extracts research_meta from research data."""
        data = {"research_meta": {"duration": 120}}
        result = handler._extract_metadata(data)
        assert result["research_meta"]["duration"] == 120

    def test_returns_empty_strings_for_missing_fields(self, handler):
        """Returns empty strings for missing fields."""
        data = {}
        result = handler._extract_metadata(data)
        assert result["strategy"] == ""
        assert result["mode"] == ""
        assert result["created_at"] == ""
        assert result["research_meta"] == {}


class TestIntegration:
    """Integration tests for complete workflows."""

    @pytest.fixture
    def mock_model(self):
        """Create mock model with realistic responses."""
        mock = Mock()

        def invoke_side_effect(prompt):
            if "entities" in prompt.lower():
                return Mock(content="Python\nMachine Learning\nAI")
            if "summary" in prompt.lower() or "summarize" in prompt.lower():
                return Mock(content="A concise summary of the research.")
            if "gaps" in prompt.lower():
                return Mock(content="Gap 1\nGap 2")
            return Mock(content="Default response")

        mock.invoke.side_effect = invoke_side_effect
        return mock

    @pytest.fixture
    def handler(self, mock_model):
        """Create handler with mock model."""
        return FollowUpContextHandler(mock_model)

    def test_full_workflow(self, handler):
        """Test complete workflow from research data to context."""
        research_data = {
            "research_id": "research_123",
            "query": "What is machine learning?",
            "formatted_findings": "Machine learning is a branch of AI that enables computers to learn from data.",
            "resources": [{"url": "https://example.com", "title": "ML Guide"}],
            "strategy": "detailed",
            "mode": "research",
        }

        context = handler.build_context(
            research_data, "How is ML used in healthcare?"
        )

        # Verify context structure
        assert context["parent_research_id"] == "research_123"
        assert context["original_query"] == "What is machine learning?"
        assert context["follow_up_query"] == "How is ML used in healthcare?"
        assert len(context["past_sources"]) == 1
        assert len(context["key_entities"]) > 0

    def test_full_workflow_with_gaps(self, handler):
        """Test gap identification in complete workflow."""
        research_data = {
            "formatted_findings": "Some research findings about ML."
        }

        gaps = handler.identify_gaps(research_data, "What about deep learning?")
        assert isinstance(gaps, list)
        assert len(gaps) <= 5