""" 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