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

"""Tests for custom LLM integration with API functions."""

from typing import List
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.api import (
    detailed_research,
    generate_report,
    quick_summary,
)
from local_deep_research.llm import clear_llm_registry, is_llm_registered


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

    identifier: str = Field(default="custom")
    messages_received: List[List[BaseMessage]] = Field(default_factory=list)

    def _generate(self, messages, **kwargs):
        """Generate response and track messages."""
        self.messages_received.append(messages)

        response = f"Response from {self.identifier}: {len(messages)} messages"
        message = AIMessage(content=response)
        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])

    @property
    def _llm_type(self):
        return f"custom_{self.identifier}"


@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 mock_search_system():
    """Create a mock search system."""
    system = MagicMock()
    system.analyze_topic.return_value = {
        "current_knowledge": "Mock research summary",
        "findings": ["Finding 1", "Finding 2"],
        "iterations": 2,
        "questions": {"iteration_1": ["Q1", "Q2"]},
        "formatted_findings": "Formatted findings",
        "all_links_of_system": ["http://example.com"],
    }
    system.model = MagicMock()
    return system


def test_quick_summary_with_custom_llm(mock_search_system):
    """Test quick_summary with a custom LLM."""
    llm = CustomTestLLM(identifier="quick")

    with patch(
        "local_deep_research.api.research_functions._init_search_system"
    ) as mock_init:
        mock_init.return_value = mock_search_system

        result = quick_summary(
            query="Test query",
            llms={"my_llm": llm},
            provider="my_llm",
            temperature=0.5,
        )

        # Verify LLM was registered
        assert is_llm_registered("my_llm")

        # Verify result structure
        assert "summary" in result
        assert result["summary"] == "Mock research summary"
        assert len(result["findings"]) == 2

        # Verify init was called
        assert mock_init.called
        # The parameters are passed through kwargs
        # Just verify the LLM was registered and used


def test_multiple_llms_registration(mock_search_system):
    """Test registering multiple LLMs at once."""
    llm1 = CustomTestLLM(identifier="llm1")
    llm2 = CustomTestLLM(identifier="llm2")
    llm3 = CustomTestLLM(identifier="llm3")

    llms = {"model1": llm1, "model2": llm2, "model3": llm3}

    with patch(
        "local_deep_research.api.research_functions._init_search_system"
    ) as mock_init:
        mock_init.return_value = mock_search_system

        quick_summary(
            query="Test multiple LLMs",
            llms=llms,
            provider="model2",  # Use the second one
        )

        # Verify all were registered
        assert is_llm_registered("model1")
        assert is_llm_registered("model2")
        assert is_llm_registered("model3")

        # Verify init was called
        assert mock_init.called
        # The LLMs were registered and that's what matters


def test_detailed_research_with_custom_llm(mock_search_system):
    """Test detailed_research with custom LLM."""
    llm = CustomTestLLM(identifier="detailed")

    with patch(
        "local_deep_research.api.research_functions._init_search_system"
    ) as mock_init:
        mock_init.return_value = mock_search_system

        result = detailed_research(
            query="Detailed test query",
            llms={"detail_llm": llm},
            provider="detail_llm",
            iterations=3,
            research_id="test-123",
        )

        # Verify result
        assert result["summary"] == "Mock research summary"
        assert "findings" in result

    # Verify research context was set - patch before calling the function
    with (
        patch(
            "local_deep_research.api.research_functions._init_search_system"
        ) as mock_init,
        patch(
            "local_deep_research.api.research_functions.set_search_context"
        ) as mock_context,
    ):
        mock_init.return_value = mock_search_system

        detailed_research(
            query="Context test",
            llms={"ctx_llm": llm},
            provider="ctx_llm",
            research_id="ctx-123",
        )

        # Check context was set with correct research_id
        assert mock_context.called
        context_call = mock_context.call_args[0][0]
        assert context_call["research_id"] == "ctx-123"
        assert context_call["research_mode"] == "detailed"


def test_generate_report_with_custom_llm():
    """Test generate_report with custom LLM."""
    llm = CustomTestLLM(identifier="report")

    # Patch the entire flow to avoid real execution
    with patch(
        "local_deep_research.api.research_functions._init_search_system"
    ) as mock_init:
        # Set up the mock system with a properly mocked model
        mock_system = MagicMock()
        mock_system.analyze_topic.return_value = {
            "current_knowledge": "Initial findings",
            "findings": ["Finding 1", "Finding 2"],
        }

        # Mock the model's invoke method properly
        mock_model = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "Report structure"
        mock_model.invoke.return_value = mock_response
        mock_system.model = mock_model

        mock_init.return_value = mock_system

        # Call the function
        result = generate_report(
            query="Report test query",
            llms={"report_llm": llm},
            provider="report_llm",
            searches_per_section=3,
        )

        # Verify custom LLM was registered
        assert is_llm_registered("report_llm")

        # Verify init was called
        assert mock_init.called

        # Verify we got a report back
        assert "content" in result
        assert isinstance(result["content"], str)


def test_llm_factory_in_api():
    """Test using a factory function with API."""
    factory_calls = []

    def create_custom_llm(model_name=None, temperature=0.7, **kwargs):
        factory_calls.append(
            {
                "model_name": model_name,
                "temperature": temperature,
                "extra": kwargs,
            }
        )
        return CustomTestLLM(identifier=f"factory-{model_name}")

    with patch(
        "local_deep_research.api.research_functions._init_search_system"
    ) as mock_init:
        mock_system = MagicMock()
        mock_system.analyze_topic.return_value = {
            "current_knowledge": "Factory test",
            "findings": [],
            "iterations": 1,
            "questions": {},
            "formatted_findings": "",
            "all_links_of_system": [],
        }
        mock_init.return_value = mock_system

        # Call quick_summary with factory
        quick_summary(
            query="Factory test",
            llms={"factory_llm": create_custom_llm},
            provider="factory_llm",
            model_name="test-v1",
            temperature=0.2,
        )

        # Verify the factory was registered
        assert is_llm_registered("factory_llm")

        # Verify init was called
        assert mock_init.called
        # The factory was registered and that's what matters


def test_combining_custom_llms_and_retrievers():
    """Test using both custom LLMs and retrievers."""
    llm = CustomTestLLM(identifier="combined")
    mock_retriever = MagicMock()

    with patch(
        "local_deep_research.api.research_functions._init_search_system"
    ) as mock_init:
        with patch(
            "local_deep_research.web_search_engines.retriever_registry.retriever_registry"
        ) as mock_reg:
            mock_system = MagicMock()
            mock_system.analyze_topic.return_value = {
                "current_knowledge": "Combined test"
            }
            mock_init.return_value = mock_system

            quick_summary(
                query="Combined test",
                llms={"custom_llm": llm},
                retrievers={"custom_retriever": mock_retriever},
                provider="custom_llm",
                search_tool="custom_retriever",
            )

            # Verify both were registered
            assert is_llm_registered("custom_llm")
            mock_reg.register_multiple.assert_called_once()

            # Verify init was called
            assert mock_init.called
            # The LLMs and retrievers were registered and that's what matters


def test_api_without_custom_llms():
    """Test that API still works without custom LLMs."""
    with patch(
        "local_deep_research.api.research_functions._init_search_system"
    ) as mock_init:
        mock_system = MagicMock()
        mock_system.analyze_topic.return_value = {
            "current_knowledge": "No custom LLM"
        }
        mock_init.return_value = mock_system

        # Call without llms parameter
        result = quick_summary(
            query="No custom LLM test",
            provider="ollama",  # Use built-in provider
            model_name="llama2",
        )

        # Should work normally
        assert "summary" in result

        # Verify init was called
        assert mock_init.called
        # Since no custom LLMs were provided, none should be registered