/tests/langchain_integration/test_custom_langchain_retriever.py

"""
Test custom LangChain retriever integration with LDR.

This tests the integration of custom LangChain retrievers with Local Deep Research,
ensuring that users can provide their own retriever implementations as search engines.
"""

import os
import pytest
from unittest.mock import Mock, patch
from typing import List
from langchain_core.retrievers import BaseRetriever, Document
from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun
from pydantic import Field

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


class CustomTestRetriever(BaseRetriever):
    """Custom retriever for testing."""

    documents: List[Document] = Field(default_factory=list)

    def __init__(self, documents: List[Document] = None, **kwargs):
        """Initialize with optional documents."""
        super().__init__(**kwargs)
        self.documents = documents or self._get_default_documents()

    def _get_default_documents(self) -> List[Document]:
        """Get default test documents."""
        return [
            Document(
                page_content="Machine learning is a method of data analysis that automates analytical model building.",
                metadata={
                    "source": "internal_kb",
                    "title": "ML Basics",
                    "url": "internal://ml-basics",
                    "doc_id": "1",
                },
            ),
            Document(
                page_content="Deep learning is a subset of machine learning that uses neural networks with multiple layers.",
                metadata={
                    "source": "internal_kb",
                    "title": "Deep Learning Overview",
                    "url": "internal://dl-overview",
                    "doc_id": "2",
                },
            ),
            Document(
                page_content="Natural language processing enables computers to understand and process human language.",
                metadata={
                    "source": "internal_kb",
                    "title": "NLP Introduction",
                    "url": "internal://nlp-intro",
                    "doc_id": "3",
                },
            ),
        ]

    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun = None,
    ) -> List[Document]:
        """Get documents relevant to the query."""
        # Simple keyword matching for testing
        relevant_docs = []
        query_lower = query.lower()

        for doc in self.documents:
            content_lower = doc.page_content.lower()
            if any(word in content_lower for word in query_lower.split()):
                relevant_docs.append(doc)

        # If no matches, return first document as fallback
        if not relevant_docs and self.documents:
            relevant_docs = [self.documents[0]]

        return relevant_docs

    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun = None,
    ) -> List[Document]:
        """Async version."""
        return self._get_relevant_documents(query, run_manager=run_manager)


@pytest.mark.skipif(
    os.environ.get("CI") == "true"
    or os.environ.get("GITHUB_ACTIONS") == "true",
    reason="Langchain integration tests skipped in CI - testing advanced features",
)
class TestCustomLangChainRetriever:
    """Test suite for custom LangChain retriever integration."""

    @pytest.fixture
    def settings_snapshot(self):
        """Create a settings snapshot for testing."""
        return {
            "llm.provider": {"value": "none", "type": "str"},
            "llm.model": {"value": "mock-model", "type": "str"},
            "llm.temperature": {"value": 0.7, "type": "float"},
            "research.iterations": {"value": 2, "type": "int"},
            "research.questions_per_iteration": {"value": 3, "type": "int"},
            "research.search_engines": {"value": ["custom_kb"], "type": "list"},
            "research.local_context": {"value": 2000, "type": "int"},
            "research.web_context": {"value": 2000, "type": "int"},
            "llm.context_window_unrestricted": {"value": False, "type": "bool"},
            "llm.context_window_size": {"value": 8192, "type": "int"},
            "llm.local_context_window_size": {"value": 4096, "type": "int"},
            "llm.supports_max_tokens": {"value": True, "type": "bool"},
            "llm.max_tokens": {"value": 4096, "type": "int"},
            "rate_limiting.llm_enabled": {"value": False, "type": "bool"},
            "search.tool": {"value": "searxng", "type": "str"},
            "search.max_results": {"value": 10, "type": "int"},
            "search.cross_engine_max_results": {"value": 100, "type": "int"},
            "search.cross_engine_use_reddit": {"value": False, "type": "bool"},
            "search.cross_engine_min_date": {"value": None, "type": "str"},
            "search.region": {"value": "us", "type": "str"},
            "search.time_period": {"value": "y", "type": "str"},
            "search.safe_search": {"value": True, "type": "bool"},
            "search.snippets_only": {"value": True, "type": "bool"},
            "search.search_language": {"value": "English", "type": "str"},
            "search.max_filtered_results": {"value": 20, "type": "int"},
        }

    @pytest.fixture
    def mock_llm(self):
        """Create a mock LLM for testing."""
        llm = Mock()
        llm.invoke.return_value = Mock(
            content="This is a summary of machine learning concepts."
        )
        llm.batch.return_value = [
            Mock(content="What are the key algorithms in machine learning?"),
            Mock(content="How does deep learning differ from traditional ML?"),
            Mock(content="What are the applications of NLP?"),
        ]
        return llm

    def test_custom_retriever_basic_usage(self, settings_snapshot, mock_llm):
        """Test basic usage of custom retriever with quick_summary."""
        # Create custom retriever
        custom_retriever = CustomTestRetriever()

        # Patch get_llm where it is imported (in research_functions), not where
        # it is defined (in llm_config), so the already-imported reference is replaced.
        with patch(
            "local_deep_research.api.research_functions.get_llm",
            return_value=mock_llm,
        ):
            # Run quick summary with custom retriever
            result = quick_summary(
                query="What is machine learning?",
                research_id="test-12345",
                retrievers={"custom_kb": custom_retriever},
                search_tool="custom_kb",
                settings_snapshot=settings_snapshot,
                iterations=1,
                questions_per_iteration=2,
            )

        # Verify results
        assert result is not None
        assert result["research_id"] == "test-12345"
        assert "summary" in result
        assert "sources" in result

    def test_custom_retriever_with_multiple_retrievers(
        self, settings_snapshot, mock_llm
    ):
        """Test using multiple custom retrievers."""
        # Create different retrievers with different document sets
        tech_docs = [
            Document(
                page_content="Python is a popular programming language for machine learning.",
                metadata={"source": "tech_kb", "title": "Python for ML"},
            ),
            Document(
                page_content="TensorFlow and PyTorch are leading deep learning frameworks.",
                metadata={"source": "tech_kb", "title": "DL Frameworks"},
            ),
        ]

        research_docs = [
            Document(
                page_content="Recent research shows promising results in transformer models.",
                metadata={
                    "source": "research_kb",
                    "title": "Transformer Research",
                },
            ),
            Document(
                page_content="Attention mechanisms have revolutionized NLP.",
                metadata={"source": "research_kb", "title": "Attention in NLP"},
            ),
        ]

        tech_retriever = CustomTestRetriever(documents=tech_docs)
        research_retriever = CustomTestRetriever(documents=research_docs)

        with patch(
            "local_deep_research.api.research_functions.get_llm",
            return_value=mock_llm,
        ):
            # Use a specific registered retriever as search_tool.
            result = detailed_research(
                query="What are the latest developments in deep learning?",
                research_id="test-67890",
                retrievers={
                    "tech_kb": tech_retriever,
                    "research_kb": research_retriever,
                },
                search_tool="tech_kb",
                settings_snapshot=settings_snapshot,
                iterations=2,
            )

        assert result["research_id"] == "test-67890"
        assert "sources" in result
        assert "summary" in result

    def test_custom_retriever_with_web_search_hybrid(
        self, settings_snapshot, mock_llm
    ):
        """Test that a custom retriever works as the search tool alongside a patched LLM."""
        custom_retriever = CustomTestRetriever()

        # Patch get_llm where it is imported in research_functions.
        with patch(
            "local_deep_research.api.research_functions.get_llm",
            return_value=mock_llm,
        ):
            # Use the registered retriever directly as the search tool.
            result = quick_summary(
                query="Machine learning trends and internal best practices",
                research_id="test-11111",
                retrievers={"internal_kb": custom_retriever},
                search_tool="internal_kb",
                settings_snapshot=settings_snapshot,
                iterations=1,
            )

        assert result["research_id"] == "test-11111"
        assert "summary" in result
        assert "sources" in result

    def test_custom_retriever_empty_results(self, settings_snapshot, mock_llm):
        """Test handling when custom retriever returns no results."""
        # Create retriever with no documents
        empty_retriever = CustomTestRetriever(documents=[])

        with patch(
            "local_deep_research.api.research_functions.get_llm",
            return_value=mock_llm,
        ):
            # Should handle gracefully even with no results
            result = quick_summary(
                query="Non-existent topic",
                research_id="test-22222",
                retrievers={"empty_kb": empty_retriever},
                search_tool="empty_kb",
                settings_snapshot=settings_snapshot,
                iterations=1,
            )

        assert result["research_id"] == "test-22222"
        assert "summary" in result  # Should still generate a summary

    def test_custom_retriever_with_llm_integration(self, settings_snapshot):
        """Test custom retriever with both custom LLM and retriever."""
        from tests.langchain_integration.test_custom_langchain_llm import (
            CustomTestLLM,
        )

        # Create custom components
        custom_llm = CustomTestLLM()
        custom_retriever = CustomTestRetriever()

        # CustomTestLLM extends LLM (not BaseChatModel), so it cannot go
        # through the provider-based registry path in get_llm.  Patch get_llm
        # to return the custom LLM directly.
        with patch(
            "local_deep_research.api.research_functions.get_llm",
            return_value=custom_llm,
        ):
            result = quick_summary(
                query="Explain machine learning concepts",
                research_id="test-33333",
                retrievers={"custom_kb": custom_retriever},
                search_tool="custom_kb",
                settings_snapshot=settings_snapshot,
                iterations=1,
            )

        assert result["research_id"] == "test-33333"
        assert "summary" in result

    def test_custom_retriever_error_handling(self, settings_snapshot, mock_llm):
        """Test that a failing retriever is handled gracefully.

        RetrieverSearchEngine.run() catches exceptions from the underlying
        retriever and returns an empty result list, so the error should NOT
        propagate to the caller.
        """

        class FailingRetriever(BaseRetriever):
            """Retriever that raises errors for testing."""

            def _get_relevant_documents(
                self, query: str, **kwargs
            ) -> List[Document]:
                raise RuntimeError("Retriever failed")

            async def _aget_relevant_documents(
                self, query: str, **kwargs
            ) -> List[Document]:
                raise RuntimeError("Retriever failed")

        failing_retriever = FailingRetriever()

        with patch(
            "local_deep_research.api.research_functions.get_llm",
            return_value=mock_llm,
        ):
            # The error is caught by RetrieverSearchEngine.run(), which returns
            # an empty list.  The system should complete without raising.
            result = quick_summary(
                query="Test query",
                retrievers={"failing_kb": failing_retriever},
                search_tool="failing_kb",
                settings_snapshot=settings_snapshot,
                iterations=1,
            )

        # System handled the error gracefully
        assert result is not None
        assert "summary" in result

    def test_custom_retriever_with_report_generation(
        self, settings_snapshot, mock_llm
    ):
        """Test custom retriever with report generation."""
        custom_retriever = CustomTestRetriever()

        # Report generation calls the LLM many times (analyse topic, determine
        # structure, research & generate each section, etc.).  Use return_value
        # instead of side_effect so the mock never runs out of responses.
        mock_llm.invoke.return_value = Mock(
            content=(
                "# Report Section\n\n"
                "Machine learning is transforming industries. "
                "Key findings include accelerating adoption."
            )
        )

        with patch(
            "local_deep_research.api.research_functions.get_llm",
            return_value=mock_llm,
        ):
            result = generate_report(
                query="Machine learning adoption strategy",
                retrievers={"internal_docs": custom_retriever},
                search_tool="internal_docs",
                settings_snapshot=settings_snapshot,
                report_type="research_report",
                iterations=1,
            )

        # generate_report returns report content, not research_id
        assert result is not None
        assert "content" in result or "report" in result

    def test_custom_retriever_metadata_handling(
        self, settings_snapshot, mock_llm
    ):
        """Test that retriever metadata is properly preserved."""
        # Create retriever with rich metadata
        docs_with_metadata = [
            Document(
                page_content="Advanced ML techniques for production systems.",
                metadata={
                    "source": "internal_kb",
                    "title": "ML in Production",
                    "author": "ML Team",
                    "date": "2024-01-15",
                    "tags": ["ml", "production", "best-practices"],
                    "department": "Engineering",
                },
            )
        ]

        metadata_retriever = CustomTestRetriever(documents=docs_with_metadata)

        with patch(
            "local_deep_research.api.research_functions.get_llm",
            return_value=mock_llm,
        ):
            result = quick_summary(
                query="ML production best practices",
                research_id="test-55555",
                retrievers={"internal": metadata_retriever},
                search_tool="internal",
                settings_snapshot=settings_snapshot,
                iterations=1,
            )

        assert result["research_id"] == "test-55555"
        assert "sources" in result