""" Mock LLM example for testing Local Deep Research without API costs. This example shows how to create mock LLMs that return predefined responses, useful for: - Testing research pipelines - Development without API keys - Debugging specific scenarios - CI/CD pipelines """ import json from typing import Any, Dict, List, Optional from langchain_core.documents import Document from langchain_core.language_models import BaseChatModel from langchain_core.messages import AIMessage, BaseMessage from langchain_core.outputs import ChatGeneration, ChatResult from langchain_core.retrievers import BaseRetriever from local_deep_research.api import ( create_settings_snapshot, generate_report, quick_summary, ) _DEFAULT_RESPONSES: Dict[str, str] = { "default": "This is a mock response for testing purposes.", "quantum": "Quantum computing uses quantum mechanics principles like superposition and entanglement to process information in fundamentally new ways.", "climate": "Climate change refers to long-term shifts in global temperatures and weather patterns, primarily driven by human activities.", "ai": "Artificial Intelligence encompasses machine learning, neural networks, and systems that can perform tasks requiring human intelligence.", "summary": "Based on the search results, here is a comprehensive summary of the findings.", "report": "# Research Report\n\n## Executive Summary\n\nThis report provides detailed analysis.\n\n## Findings\n\n1. Key finding one\n2. Key finding two", } class MockLLM(BaseChatModel): """Mock LLM that returns predefined responses based on queries.""" response_map: Optional[Dict[str, str]] = None call_history: Optional[List[Dict]] = None def model_post_init(self, __context: Any) -> None: super().model_post_init(__context) if self.response_map is None: self.response_map = dict(_DEFAULT_RESPONSES) if self.call_history is None: self.call_history = [] def _generate( self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[Any] = None, **kwargs: Any, ) -> ChatResult: """Generate mock response based on query content.""" # Extract query query = messages[-1].content.lower() if messages else "" # Log the call self.call_history.append( { "messages": [ {"role": m.__class__.__name__, "content": m.content} for m in messages ], "kwargs": kwargs, } ) # Find matching response response = self.response_map.get("default", "Mock response") for key, value in self.response_map.items(): if key in query: response = value break # Create response message = AIMessage(content=response) generation = ChatGeneration(message=message) return ChatResult(generations=[generation]) @property def _llm_type(self) -> str: return "mock" def get_call_history(self) -> List[Dict]: """Get history of all calls made to this LLM.""" return self.call_history def clear_history(self): """Clear call history.""" self.call_history = [] class ScenarioMockLLM(BaseChatModel): """Mock LLM that simulates specific scenarios for testing.""" scenario: str = "success" call_count: int = 0 def _generate( self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[Any] = None, **kwargs: Any, ) -> ChatResult: """Generate response based on scenario.""" self.call_count += 1 if self.scenario == "success": response = self._success_response(messages) elif self.scenario == "partial_failure": response = self._partial_failure_response() elif self.scenario == "empty": response = "" elif self.scenario == "verbose": response = self._verbose_response(messages) elif self.scenario == "json": response = self._json_response(messages) else: response = f"Unknown scenario: {self.scenario}" message = AIMessage(content=response) generation = ChatGeneration(message=message) return ChatResult(generations=[generation]) def _success_response(self, messages): """Generate successful response.""" query = messages[-1].content if messages else "query" return f"Successfully analyzed: {query}. Found 5 relevant sources with high confidence." def _partial_failure_response(self): """Generate partial failure response.""" if self.call_count % 3 == 0: return "Unable to process query due to insufficient data." return "Partial results found. Limited information available." def _verbose_response(self, messages): """Generate verbose response for testing truncation.""" query = messages[-1].content if messages else "query" return f""" Detailed Analysis of: {query} Section 1: Introduction {"-" * 50} This is a comprehensive analysis with multiple sections. Section 2: Methodology {"-" * 50} We used advanced techniques to analyze this query. Section 3: Findings {"-" * 50} Finding 1: Important discovery about the topic. Finding 2: Another significant insight. Finding 3: Additional relevant information. Section 4: Conclusion {"-" * 50} In conclusion, this analysis provides valuable insights. """ + "\n".join([f"Additional point {i}" for i in range(20)]) def _json_response(self, messages): """Generate JSON response for testing parsing.""" query = messages[-1].content if messages else "query" data = { "query": query, "findings": [ {"id": 1, "content": "First finding", "confidence": 0.9}, {"id": 2, "content": "Second finding", "confidence": 0.85}, ], "summary": "JSON-formatted response for testing", "metadata": {"timestamp": "2024-01-01T00:00:00Z", "version": "1.0"}, } return json.dumps(data, indent=2) @property def _llm_type(self) -> str: return f"scenario_{self.scenario}" class MockRetriever(BaseRetriever): """Offline retriever returning canned documents. Registered as a search engine to keep the pipeline fully offline — otherwise falling back to the default engines hits live services. """ def _get_relevant_documents(self, query, *, run_manager=None): return [ Document( page_content=f"Mock document 1 about {query}", metadata={"source": "mock://doc1"}, ), Document( page_content=f"Mock document 2 about {query}", metadata={"source": "mock://doc2"}, ), ] def test_basic_mock(): """Test basic mock functionality.""" print("Testing Basic Mock LLM") print("-" * 40) mock_llm = MockLLM() snapshot = create_settings_snapshot( provider="mock", overrides={"search.tool": "mock_retriever"}, ) result = quick_summary( query="Tell me about quantum computing", llms={"mock": mock_llm}, retrievers={"mock_retriever": MockRetriever()}, settings_snapshot=snapshot, iterations=1, ) print(f"Result: {result['summary']}") print(f"Call history: {len(mock_llm.get_call_history())} calls") print() def test_scenario_mocks(): """Test different scenario mocks.""" print("Testing Scenario Mocks") print("-" * 40) scenarios = ["success", "partial_failure", "empty", "verbose", "json"] for scenario in scenarios: print(f"\nScenario: {scenario}") mock_llm = ScenarioMockLLM(scenario=scenario) try: snapshot = create_settings_snapshot( provider=f"mock_{scenario}", overrides={"search.tool": "mock_retriever"}, ) result = quick_summary( query="Test query for scenario", llms={f"mock_{scenario}": mock_llm}, retrievers={"mock_retriever": MockRetriever()}, settings_snapshot=snapshot, iterations=1, ) print(f"Summary preview: {result['summary'][:100]}...") print(f"Calls made: {mock_llm.call_count}") except Exception as e: print(f"Error in scenario {scenario}: {e}") def test_mock_in_pipeline(): """Test mock LLM in a full research pipeline.""" print("\nTesting Mock in Research Pipeline") print("-" * 40) # Create specialized mocks for different stages response_map = { "questions": "Generated questions: 1) What is X? 2) How does Y work? 3) What are the benefits?", "analysis": "Analysis complete. Key findings: A, B, and C.", "synthesis": "Synthesis: Combining all findings into coherent summary.", "report": "# Final Report\n\n## Summary\n\nAll findings have been compiled.", } mock_llm = MockLLM(response_map=response_map) # Test with report generation snapshot = create_settings_snapshot( provider="pipeline_mock", overrides={"search.tool": "mock_retriever"}, ) report = generate_report( query="Create a comprehensive report", llms={"pipeline_mock": mock_llm}, retrievers={"mock_retriever": MockRetriever()}, settings_snapshot=snapshot, searches_per_section=1, ) print(f"Report generated: {len(report.get('content', ''))} characters") print(f"Total LLM calls: {len(mock_llm.get_call_history())}") # Analyze call patterns print("\nCall Analysis:") for i, call in enumerate(mock_llm.get_call_history()[:5]): # First 5 calls last_message = ( call["messages"][-1]["content"] if call["messages"] else "No message" ) print(f" Call {i + 1}: {last_message[:50]}...") def test_mock_with_custom_retriever(): """Test mock LLM with custom retriever.""" print("\nTesting Mock LLM with Custom Retriever") print("-" * 40) mock_llm = MockLLM( response_map={ "default": "Analyzed documents and found relevant information.", "summary": "Summary: Based on internal documents, the answer is clear.", } ) snapshot = create_settings_snapshot( provider="mock", overrides={"search.tool": "mock_retriever"}, ) result = quick_summary( query="Internal policy question", llms={"mock": mock_llm}, retrievers={"mock_retriever": MockRetriever()}, settings_snapshot=snapshot, ) print(f"Result: {result['summary']}") print(f"Sources: {result.get('sources', [])}") def main(): """Run all mock examples.""" test_basic_mock() test_scenario_mocks() test_mock_in_pipeline() test_mock_with_custom_retriever() print("\n" + "=" * 60) print("Mock LLM Testing Complete!") print( "Use these patterns to test your research pipelines without API costs." ) if __name__ == "__main__": main()