learningcircuit--local-deep-research/tests/langchain_integration/test_custom_langchain_llm.py

"""
Test custom LangChain LLM integration with LDR.

This tests the integration of custom LangChain LLMs with Local Deep Research,
ensuring that users can provide their own LLM implementations.
"""

import os
import pytest
from unittest.mock import Mock, patch
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage
from langchain_core.outputs import ChatResult, ChatGeneration
from typing import Any, List, Optional

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


class CustomTestLLM(BaseChatModel):
    """Custom Chat LLM for testing."""

    @property
    def _llm_type(self) -> str:
        """Return identifier of llm."""
        return "custom_test_llm"

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager=None,
        **kwargs: Any,
    ) -> ChatResult:
        """Generate a chat response based on message content."""
        prompt = messages[-1].content if messages else ""
        # Simple test response that varies based on prompt content
        if "quantum" in prompt.lower():
            response = "Quantum computing uses quantum bits (qubits) that can exist in superposition states."
        elif "machine learning" in prompt.lower():
            response = "Machine learning is a subset of AI that enables systems to learn from data."
        elif "climate" in prompt.lower():
            response = "Climate change is primarily driven by greenhouse gas emissions."
        else:
            response = f"This is a response from the custom LLM about: {prompt[:50]}..."
        return ChatResult(
            generations=[ChatGeneration(message=AIMessage(content=response))]
        )


@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 TestCustomLangChainLLM:
    """Test suite for custom LangChain LLM integration."""

    @pytest.fixture
    def mock_session(self):
        """Create a mock database session."""
        session = Mock()
        return session

    @pytest.fixture
    def settings_snapshot(self):
        """Create a settings snapshot for testing."""
        return {
            "llm.provider": {"value": "custom", "type": "str"},
            "llm.model": {"value": "custom_test_llm", "type": "str"},
            "llm.temperature": {"value": 0.7, "type": "float"},
            "llm.custom.api_key": {"value": "test-key", "type": "str"},
            "research.iterations": {"value": 2, "type": "int"},
            "research.questions_per_iteration": {"value": 3, "type": "int"},
            "research.search_engines": {"value": ["wikipedia"], "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": "wikipedia", "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"},
        }

    def test_custom_llm_basic_usage(self, settings_snapshot):
        """Test basic usage of custom LLM with quick_summary."""
        # Create custom LLM instance
        custom_llm = CustomTestLLM()

        # Mock the search results
        with patch(
            "local_deep_research.api.research_functions.get_search"
        ) as _mock_get_search:
            mock_search_engine = Mock()
            mock_search_engine.run.return_value = [
                {
                    "url": "https://example.com/quantum",
                    "title": "Quantum Computing Basics",
                    "content": "Quantum computing is a revolutionary technology...",
                    "source": "wikipedia",
                }
            ]

            _mock_get_search.return_value = mock_search_engine

            # Run quick summary with custom LLM
            result = quick_summary(
                query="What is quantum computing?",
                research_id=12345,
                llms={"custom": custom_llm},
                settings_snapshot=settings_snapshot,
                search_tool="wikipedia",
                iterations=1,
                questions_per_iteration=2,
            )

        # Verify results
        assert result is not None
        assert "research_id" in result
        assert result["research_id"] == 12345
        assert "summary" in result
        assert "quantum" in result["summary"].lower()
        assert "sources" in result
        assert len(result["sources"]) > 0

    def test_custom_llm_with_detailed_research(self, settings_snapshot):
        """Test custom LLM with detailed_research function."""
        custom_llm = CustomTestLLM()

        with patch(
            "local_deep_research.api.research_functions.get_search"
        ) as mock_search:
            # Mock multiple search results
            mock_search_engine = Mock()
            mock_search_engine.run.return_value = [
                {
                    "url": "https://example.com/ml1",
                    "title": "Machine Learning Introduction",
                    "content": "Machine learning is transforming industries...",
                    "source": "wikipedia",
                },
                {
                    "url": "https://example.com/ml2",
                    "title": "ML Applications",
                    "content": "Applications of machine learning include...",
                    "source": "wikipedia",
                },
            ]

            mock_search.return_value = mock_search_engine

            result = detailed_research(
                query="Explain machine learning applications",
                research_id="test-67890",
                llms={"custom": custom_llm},
                settings_snapshot=settings_snapshot,
                search_tool="wikipedia",
                iterations=2,
                questions_per_iteration=3,
            )

        assert result is not None
        assert result["research_id"] == "test-67890"
        assert "machine learning" in result["summary"].lower()
        assert len(result["sources"]) >= 2
        assert "findings" in result

    def test_custom_llm_with_custom_factory(self, settings_snapshot):
        """Test using a custom LLM factory function."""

        def create_custom_llm(
            model_name=None, temperature=None, settings_snapshot=None
        ):
            """Factory function for creating custom LLM."""
            # Access settings from snapshot
            api_key = settings_snapshot.get("llm.custom.api_key", {}).get(
                "value"
            )
            assert api_key == "test-key"  # Verify settings access

            # Create and configure custom LLM
            llm = CustomTestLLM()
            # In real implementation, would use api_key and other settings
            return llm

        # Use factory to create LLM
        custom_llm = create_custom_llm(
            model_name="custom_test_llm",
            temperature=0.7,
            settings_snapshot=settings_snapshot,
        )

        with patch(
            "local_deep_research.api.research_functions.get_search"
        ) as _mock_get_search:
            mock_search_engine = Mock()
            mock_search_engine.run.return_value = [
                {
                    "url": "https://example.com/climate",
                    "title": "Climate Change Overview",
                    "content": "Climate change affects global weather patterns...",
                    "source": "wikipedia",
                }
            ]

            _mock_get_search.return_value = mock_search_engine

            result = quick_summary(
                query="Impact of climate change",
                research_id="test-11111",
                llms={"custom": custom_llm},
                settings_snapshot=settings_snapshot,
                search_tool="wikipedia",
                iterations=1,
            )

        assert result["research_id"] == "test-11111"
        assert "climate" in result["summary"].lower()

    def test_custom_llm_error_handling(self, settings_snapshot):
        """Test error handling with custom LLM that raises errors."""

        class FailingLLM(BaseChatModel):
            """LLM that raises errors for testing."""

            @property
            def _llm_type(self) -> str:
                return "failing_llm"

            def _generate(
                self,
                messages: List[BaseMessage],
                stop=None,
                run_manager=None,
                **kwargs,
            ) -> ChatResult:
                raise RuntimeError("LLM call failed")

        failing_llm = FailingLLM()

        # Override provider to match the "failing" LLM registration key
        failing_settings = {
            **settings_snapshot,
            "llm.provider": {"value": "failing", "type": "str"},
        }

        with patch(
            "local_deep_research.api.research_functions.get_search"
        ) as _mock_get_search:
            mock_search_engine = Mock()
            mock_search_engine.run.return_value = [
                {
                    "url": "https://example.com/test",
                    "title": "Test Article",
                    "content": "Test content...",
                    "source": "wikipedia",
                }
            ]

            _mock_get_search.return_value = mock_search_engine

            # The search strategy catches LLM errors gracefully,
            # so the research completes but with degraded results
            try:
                result = quick_summary(
                    query="Test query",
                    llms={"failing": failing_llm},
                    settings_snapshot=failing_settings,
                    search_tool="wikipedia",
                    iterations=1,
                )
                # If it completes, verify we got some result back
                assert result is not None
                assert "summary" in result
            except RuntimeError as e:
                # If it does raise, verify it's the expected error
                assert "LLM call failed" in str(e)

    def test_custom_llm_streaming(self, settings_snapshot):
        """Test custom LLM with streaming support."""

        class StreamingTestLLM(BaseChatModel):
            """Chat LLM with streaming support for testing."""

            @property
            def _llm_type(self) -> str:
                return "streaming_test_llm"

            def _generate(
                self,
                messages: List[BaseMessage],
                stop=None,
                run_manager=None,
                **kwargs,
            ) -> ChatResult:
                response = (
                    "This is a streaming response about quantum computing."
                )
                if run_manager:
                    for token in response.split():
                        run_manager.on_llm_new_token(token + " ")
                return ChatResult(
                    generations=[
                        ChatGeneration(message=AIMessage(content=response))
                    ]
                )

        streaming_llm = StreamingTestLLM()

        # Override provider to match the "streaming" LLM registration key
        streaming_settings = {
            **settings_snapshot,
            "llm.provider": {"value": "streaming", "type": "str"},
        }

        with patch(
            "local_deep_research.api.research_functions.get_search"
        ) as _mock_get_search:
            mock_search_engine = Mock()
            mock_search_engine.run.return_value = [
                {
                    "url": "https://example.com/quantum",
                    "title": "Quantum Info",
                    "content": "Quantum information...",
                    "source": "wikipedia",
                }
            ]

            _mock_get_search.return_value = mock_search_engine

            result = quick_summary(
                query="Quantum computing basics",
                research_id="test-99999",
                llms={"streaming": streaming_llm},
                settings_snapshot=streaming_settings,
                search_tool="wikipedia",
                iterations=1,
            )

        assert result["research_id"] == "test-99999"
        assert "streaming response" in result["summary"]