learningcircuit--local-deep-research/tests/test_llm/test_llm_integration.py

"""Tests for LLM integration with the broader system."""

from typing import List, Optional
from unittest.mock import MagicMock, patch

import pytest
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from pydantic import Field

from local_deep_research.config.llm_config import get_llm
from local_deep_research.llm import clear_llm_registry, register_llm


class TestLLM(BaseChatModel):
    """Test LLM implementation."""

    response_text: str = Field(default="Test response")
    call_count: int = Field(default=0)
    last_messages: Optional[List[BaseMessage]] = Field(default=None)

    def _generate(self, messages, **kwargs):
        """Generate test response."""
        self.call_count += 1
        self.last_messages = messages

        message = AIMessage(content=self.response_text)
        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])

    @property
    def _llm_type(self):
        return "test"


@pytest.fixture(autouse=True)
def clear_registry():
    """Clear the registry before and after each test."""
    clear_llm_registry()
    yield
    clear_llm_registry()


@pytest.fixture
def full_settings_snapshot():
    """Provide a complete settings snapshot for tests."""

    def _text_param(value):
        return {"value": value, "ui_element": "text"}

    def _number_param(value):
        return {"value": value, "ui_element": "number"}

    def _bool_param(value):
        return {"value": value, "ui_element": "checkbox"}

    return {
        # Real runs always carry a primary engine; the inference
        # egress PEP now fails closed without it.
        "search.tool": _text_param("searxng"),
        "llm.model": _text_param("test-model"),
        "llm.temperature": _number_param(0.7),
        "llm.provider": _text_param("test"),
        "llm.supports_max_tokens": _bool_param(True),
        "llm.max_tokens": _number_param(100000),
        "llm.local_context_window_size": _number_param(4096),
        "llm.context_window_unrestricted": _bool_param(True),
        "llm.context_window_size": _number_param(128000),
        "llm.ollama.url": _text_param("http://localhost:11434"),
        "llm.openai.api_key": _text_param(None),
        "llm.anthropic.api_key": _text_param(None),
        "llm.openai_endpoint.api_key": _text_param(None),
        "llm.openai_endpoint.url": _text_param("https://openrouter.ai/api/v1"),
        "rate_limiting.llm_enabled": _bool_param(False),
    }


def test_get_llm_with_custom_provider(full_settings_snapshot):
    """Test that get_llm returns custom LLMs when registered."""
    # Register a custom LLM
    custom_llm = TestLLM(response_text="Custom response")
    register_llm("custom_provider", custom_llm)

    # Get LLM with custom provider
    with patch(
        "local_deep_research.config.llm_config.wrap_llm_without_think_tags"
    ) as mock_wrap:
        # Configure mock to return the LLM passed to it
        mock_wrap.side_effect = lambda llm, **kwargs: llm

        get_llm(
            provider="custom_provider",
            temperature=0.5,
            settings_snapshot=full_settings_snapshot,
        )

        # Verify wrap was called with our custom LLM
        assert mock_wrap.called
        assert mock_wrap.call_args[0][0] is custom_llm


def test_get_llm_with_factory_function(full_settings_snapshot):
    """Test that get_llm works with factory functions."""
    # Create a factory that tracks calls
    factory_calls = []

    def llm_factory(model_name=None, temperature=0.7, **kwargs):
        factory_calls.append(
            {
                "model_name": model_name,
                "temperature": temperature,
                "kwargs": kwargs,
            }
        )
        return TestLLM(response_text=f"Factory LLM: {model_name}")

    # Register the factory
    register_llm("factory_provider", llm_factory)

    # Get LLM with factory provider
    with patch(
        "local_deep_research.config.llm_config.wrap_llm_without_think_tags"
    ) as mock_wrap:
        mock_wrap.side_effect = lambda llm, **kwargs: llm

        get_llm(
            provider="factory_provider",
            model_name="test-model",
            temperature=0.3,
            settings_snapshot=full_settings_snapshot,
        )

        # Verify factory was called with correct parameters
        assert len(factory_calls) == 1
        assert factory_calls[0]["model_name"] == "test-model"
        assert factory_calls[0]["temperature"] == 0.3


def test_api_integration_with_custom_llm():
    """Test that API functions work with custom LLMs."""
    from local_deep_research.api import quick_summary

    # Create a custom LLM
    custom_llm = TestLLM(response_text="API test response")

    # Mock the necessary components
    with patch(
        "local_deep_research.api.research_functions._init_search_system"
    ) as mock_init:
        # Create a mock search system
        mock_system = MagicMock()
        mock_system.analyze_topic.return_value = {
            "current_knowledge": "Test summary",
            "findings": [],
            "iterations": 1,
            "questions": {},
            "formatted_findings": "Test findings",
            "all_links_of_system": [],
        }
        mock_init.return_value = mock_system

        # Call quick_summary with custom LLM
        result = quick_summary(
            query="Test query",
            llms={"test_llm": custom_llm},
            provider="test_llm",
        )

        # Verify the LLM was registered and used
        assert result["summary"] == "Test summary"
        assert mock_init.called

        # Verify init was called
        # The LLM was registered and used successfully


def test_multiple_custom_llms(full_settings_snapshot):
    """Test registering and using multiple custom LLMs."""
    llm1 = TestLLM(response_text="Response 1")
    llm2 = TestLLM(response_text="Response 2")

    register_llm("provider1", llm1)
    register_llm("provider2", llm2)

    with patch(
        "local_deep_research.config.llm_config.wrap_llm_without_think_tags"
    ) as mock_wrap:
        mock_wrap.side_effect = lambda llm, **kwargs: llm

        # Get first LLM
        get_llm(provider="provider1", settings_snapshot=full_settings_snapshot)
        assert mock_wrap.call_args[0][0] is llm1

        # Get second LLM
        get_llm(provider="provider2", settings_snapshot=full_settings_snapshot)
        assert mock_wrap.call_args[0][0] is llm2


def test_custom_llm_with_research_context(full_settings_snapshot):
    """Test that custom LLMs receive research context properly."""
    custom_llm = TestLLM()
    register_llm("context_test", custom_llm)

    research_id = "test-research-123"
    research_context = {"query": "test", "mode": "quick"}

    with patch(
        "local_deep_research.config.llm_config.wrap_llm_without_think_tags"
    ) as mock_wrap:
        mock_wrap.side_effect = lambda llm, **kwargs: llm

        get_llm(
            provider="context_test",
            research_id=research_id,
            research_context=research_context,
            settings_snapshot=full_settings_snapshot,
        )

        # Verify context was passed to wrapper
        wrap_kwargs = mock_wrap.call_args[1]
        assert wrap_kwargs["research_id"] == research_id
        assert wrap_kwargs["research_context"] == research_context
        assert wrap_kwargs["provider"] == "context_test"


def test_factory_error_handling(full_settings_snapshot):
    """Test error handling when factory fails."""

    def failing_factory(**kwargs):
        raise ValueError("Factory error")

    register_llm("failing_factory", failing_factory)

    # Should raise the factory error
    with pytest.raises(ValueError, match="Factory error"):
        get_llm(
            provider="failing_factory", settings_snapshot=full_settings_snapshot
        )


def test_invalid_provider_after_checking_registry(
    monkeypatch, full_settings_snapshot
):
    """Test that invalid provider error is raised for non-existent providers."""
    # Don't register anything

    with pytest.raises(ValueError, match="Invalid provider: fake_provider"):
        get_llm(
            provider="fake_provider", settings_snapshot=full_settings_snapshot
        )