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