learningcircuit--local-deep-research/tests/retriever_integration/test_api_integration.py

"""Tests for retriever integration with LDR API functions."""

from unittest.mock import Mock, patch

from langchain_core.retrievers import BaseRetriever, Document

from local_deep_research.api import (
    detailed_research,
    generate_report,
    quick_summary,
)

# Import retriever_registry from the same module path that the API functions use
# This ensures we're using the same registry instance


class SimpleRetriever(BaseRetriever):
    """Simple retriever for testing API integration."""

    content: str = "Test content"

    class Config:
        """Pydantic config."""

        arbitrary_types_allowed = True

    def _get_relevant_documents(self, query: str):
        return [
            Document(
                page_content=f"{self.content} for query: {query}",
                metadata={"title": f"Result for {query}", "source": "test"},
            )
        ]

    async def _aget_relevant_documents(self, query: str):
        return self._get_relevant_documents(query)


class TestAPIIntegration:
    """Test retriever integration with API functions."""

    def setup_method(self):
        """Clear registry before each test."""
        # Import the registry to ensure we're using the right instance
        from local_deep_research.web_search_engines.retriever_registry import (
            retriever_registry,
        )

        retriever_registry.clear()

    def teardown_method(self):
        """Clear registry after each test."""
        # Import the registry to ensure we're using the right instance
        from local_deep_research.web_search_engines.retriever_registry import (
            retriever_registry,
        )

        retriever_registry.clear()

    @patch("local_deep_research.api.research_functions._init_search_system")
    def test_quick_summary_with_single_retriever(self, mock_init):
        """Test quick_summary with a single retriever."""
        # Mock the search system
        mock_system = Mock()
        mock_system.analyze_topic.return_value = {
            "current_knowledge": "Summary based on retriever content",
            "findings": ["Finding 1", "Finding 2"],
            "iterations": 1,
            "questions": {},
            "all_links_of_system": [],
        }
        mock_init.return_value = mock_system

        # Create and use retriever
        retriever = SimpleRetriever(content="Test content")
        result = quick_summary(
            query="test query",
            retrievers={"test_kb": retriever},
            search_tool="test_kb",
        )

        # Verify result without checking registry state
        # The important thing is that the function works correctly
        assert "summary" in result
        assert result["summary"] == "Summary based on retriever content"

        # Verify that the search system was initialized with correct search_tool
        mock_init.assert_called_once()
        call_kwargs = mock_init.call_args[1]
        assert call_kwargs.get("search_tool") == "test_kb"

    @patch("local_deep_research.api.research_functions._init_search_system")
    def test_quick_summary_with_multiple_retrievers(self, mock_init):
        """Test quick_summary with multiple retrievers."""
        mock_system = Mock()
        mock_system.analyze_topic.return_value = {
            "current_knowledge": "Combined summary",
            "findings": [],
            "iterations": 1,
            "questions": {},
            "all_links_of_system": [],
        }
        mock_init.return_value = mock_system

        # Create multiple retrievers
        retrievers = {
            "kb1": SimpleRetriever(content="Knowledge base 1"),
            "kb2": SimpleRetriever(content="Knowledge base 2"),
            "kb3": SimpleRetriever(content="Knowledge base 3"),
        }

        quick_summary(
            query="multi retriever test",
            retrievers=retrievers,
            search_tool="kb1",
        )

        # Verify that the function works correctly without checking registry state
        # The search_tool="kb1" should work with the registered retrievers
        mock_init.assert_called_once()
        call_kwargs = mock_init.call_args[1]
        assert call_kwargs.get("search_tool") == "kb1"

    @patch("local_deep_research.api.research_functions._init_search_system")
    def test_detailed_research_with_retrievers(self, mock_init):
        """Test detailed_research with retrievers."""
        mock_system = Mock()
        mock_system.analyze_topic.return_value = {
            "current_knowledge": "Detailed analysis",
            "findings": ["Finding 1", "Finding 2", "Finding 3"],
            "iterations": 3,
            "questions": {"iteration_1": ["Q1", "Q2"]},
            "formatted_findings": "Formatted findings",
            "all_links_of_system": ["source1", "source2"],
        }
        mock_init.return_value = mock_system

        retriever = SimpleRetriever(content="Test content")
        result = detailed_research(
            query="detailed test",
            retrievers={"detail_kb": retriever},
            search_tool="detail_kb",
            iterations=3,
        )

        # Verify result structure
        assert result["query"] == "detailed test"
        assert result["summary"] == "Detailed analysis"
        assert len(result["findings"]) == 3
        assert result["metadata"]["search_tool"] == "detail_kb"
        assert result["metadata"]["iterations_requested"] == 3

    @patch("local_deep_research.api.research_functions._init_search_system")
    @patch(
        "local_deep_research.api.research_functions.IntegratedReportGenerator"
    )
    def test_generate_report_with_retrievers(
        self, mock_report_gen_class, mock_init
    ):
        """Test generate_report with retrievers."""
        # Mock system
        mock_system = Mock()
        mock_system.analyze_topic.return_value = {
            "findings": "initial findings"
        }
        mock_init.return_value = mock_system

        # Mock report generator
        mock_report_gen = Mock()
        mock_report_gen.generate_report.return_value = {
            "content": "# Report\nGenerated report content",
            "metadata": {"query": "report test"},
        }
        mock_report_gen_class.return_value = mock_report_gen

        retriever = SimpleRetriever(content="Test content")
        result = generate_report(
            query="report test",
            retrievers={"report_kb": retriever},
            search_tool="report_kb",
        )

        # Verify result without checking registry state
        assert "content" in result
        assert "# Report" in result["content"]

        # Verify that the search system was initialized with correct search_tool
        mock_init.assert_called_once()
        call_kwargs = mock_init.call_args[1]
        assert call_kwargs.get("search_tool") == "report_kb"

    def test_retriever_selection(self):
        """Test selecting specific retriever from multiple."""
        # This test actually runs the retriever to verify selection works
        retrievers = {
            "kb1": SimpleRetriever(content="Content 1"),
            "kb2": SimpleRetriever(content="Content 2"),
            "kb3": SimpleRetriever(content="Content 3"),
        }

        # Register all retrievers
        # Import the registry to ensure we're using the right instance
        from local_deep_research.web_search_engines.retriever_registry import (
            retriever_registry,
        )

        retriever_registry.register_multiple(retrievers)

        # Verify we can get specific retriever
        selected = retriever_registry.get("kb2")
        assert selected == retrievers["kb2"]

        # Verify it returns correct content
        docs = selected.invoke("test")
        assert "Content 2" in docs[0].page_content

    @patch("local_deep_research.api.research_functions._init_search_system")
    def test_research_id_generation(self, mock_init):
        """Test that research_id is generated if not provided."""
        mock_system = Mock()
        mock_system.analyze_topic.return_value = {
            "current_knowledge": "Summary",
            "findings": [],
            "iterations": 1,
            "questions": {},
            "all_links_of_system": [],
        }
        mock_init.return_value = mock_system

        quick_summary(
            query="test", retrievers={"test": SimpleRetriever(content="Test")}
        )

        # Should have generated a research_id (UUID format)
        assert mock_init.called

    def test_empty_retrievers_dict(self):
        """Test handling of empty retrievers dictionary."""
        # Should not crash with empty dict
        with patch(
            "local_deep_research.api.research_functions._init_search_system"
        ) as mock_init:
            mock_system = Mock()
            mock_system.analyze_topic.return_value = {
                "current_knowledge": "Summary",
                "findings": [],
                "iterations": 1,
                "questions": {},
                "all_links_of_system": [],
            }
            mock_init.return_value = mock_system

            quick_summary(
                query="test",
                retrievers={},  # Empty dict
                search_tool="searxng",  # Use regular search engine
            )

            # Verify that the search system was initialized with regular search engine
            mock_init.assert_called_once()
            call_kwargs = mock_init.call_args[1]
            assert call_kwargs.get("search_tool") == "searxng"