learningcircuit--local-deep-research/examples/llm_integration/mock_llm_example.py

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