learningcircuit--local-deep-research/tests/test_api_key_configuration.py

"""
Tests for API key configuration and LLM execution.
Tests that API keys can be set and used properly for different LLM providers.
"""

import pytest
from unittest.mock import Mock, patch

from langchain_core.language_models import BaseChatModel

from local_deep_research.config.llm_config import (
    get_llm,
)
from local_deep_research.llm.providers.implementations.openai import (
    OpenAIProvider,
)
from local_deep_research.llm.providers.implementations.anthropic import (
    AnthropicProvider,
)
from local_deep_research.llm.providers.implementations.custom_openai_endpoint import (
    CustomOpenAIEndpointProvider,
)
from local_deep_research.settings import SettingsManager


def _llm_mock(**kwargs):
    """Build a Mock that satisfies the BaseChatModel isinstance check.

    The registered-LLM branch in get_llm() validates that create_llm()
    returned a real BaseChatModel; with spec=BaseChatModel the Mock
    passes isinstance() while still capturing call args.
    """
    return Mock(spec=BaseChatModel, **kwargs)


@pytest.fixture(autouse=True)
def _restore_auto_discovered_providers():
    """Make sure auto-discovered providers are registered before each test.

    The previous version of this fixture cleared the registry to force the
    procedural ``if/elif`` chain in ``get_llm`` — that chain has been
    deleted, so tests now go through the live registered-LLM path. ChatXxx
    patches in each test target the provider class module (e.g.
    ``...implementations.openai.ChatOpenAI``) which is where the LLM is
    actually constructed.

    ``force_refresh=True`` is required because the discovery singleton
    short-circuits subsequent calls otherwise — and other test modules in
    the suite may have called ``clear_llm_registry()`` between tests.
    """
    from local_deep_research.llm.providers import discover_providers

    discover_providers(force_refresh=True)
    yield


class TestAPIKeyConfiguration:
    """Test API key configuration for different providers."""

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

    @pytest.fixture
    def settings_manager(self, mock_db_session):
        """Create a settings manager with mock session."""
        return SettingsManager(mock_db_session)

    def _get_base_settings(self):
        """Get base settings that all tests need."""
        return {
            "search.tool": "searxng",
            "llm.supports_max_tokens": True,
            "llm.max_tokens": 4096,
            "llm.context_window_unrestricted": False,
            "llm.context_window_size": 8192,
            "llm.local_context_window_size": 4096,
            "rate_limiting.llm_enabled": False,
        }

    @pytest.fixture
    def settings_snapshot_with_openai(self):
        """Create a settings snapshot with OpenAI API key."""
        settings = self._get_base_settings()
        settings.update(
            {
                "llm.provider": "openai",
                "llm.model": "gpt-4",
                "llm.temperature": 0.7,
                "llm.openai.api_key": "test-openai-api-key",
            }
        )
        return settings

    @pytest.fixture
    def settings_snapshot_with_anthropic(self):
        """Create a settings snapshot with Anthropic API key."""
        settings = self._get_base_settings()
        settings.update(
            {
                "llm.provider": "anthropic",
                "llm.model": "claude-3-opus-20240229",
                "llm.temperature": 0.7,
                "llm.anthropic.api_key": "test-anthropic-api-key",
                "llm.context_window_size": 200000,
            }
        )
        return settings

    @pytest.fixture
    def settings_snapshot_with_openai_endpoint(self):
        """Create a settings snapshot with OpenAI endpoint API key."""
        settings = self._get_base_settings()
        settings.update(
            {
                "llm.provider": "openai_endpoint",
                "llm.model": "claude-3-opus",
                "llm.temperature": 0.5,
                "llm.openai_endpoint.api_key": "test-openrouter-api-key",
                "llm.openai_endpoint.url": "https://openrouter.ai/api/v1",
                "llm.context_window_unrestricted": True,
                "llm.context_window_size": 128000,
                # "library" is the only engine the factory can instantiate
                # from this minimal snapshot (searxng et al. need an instance
                # URL), and the research path does create the engine. But
                # "library" is a PRIVATE engine, so under the default adaptive
                # egress scope it would resolve to PRIVATE_ONLY and force
                # local LLM — denying the remote openai_endpoint provider this
                # test configures. Pin the scope to "unprotected" (the escape
                # hatch; "both" is retired per ADR-0007) so the test stays
                # about LLM-provider config, not egress.
                "search.tool": "library",
                "policy.egress_scope": "unprotected",
                "search.max_results": 10,
                "search.cross_engine_max_results": 100,
                "search.cross_engine_use_reddit": False,
                "search.cross_engine_min_date": None,
                "search.region": "us",
                "search.time_period": "y",
                "search.safe_search": True,
                "search.snippets_only": True,
                "search.search_language": "English",
                "search.max_filtered_results": 20,
                "research.iterations": 2,
                "research.questions_per_iteration": 3,
                "research.local_context": 2000,
                "research.web_context": 2000,
            }
        )
        return settings

    def test_openai_api_key_configuration(self, settings_snapshot_with_openai):
        """Test that OpenAI API key can be configured and used."""
        with patch(
            "local_deep_research.llm.providers.implementations.openai.ChatOpenAI"
        ) as mock_openai:
            # Create a mock LLM instance
            mock_llm_instance = _llm_mock()
            mock_openai.return_value = mock_llm_instance

            # Get LLM with OpenAI settings
            get_llm(settings_snapshot=settings_snapshot_with_openai)

            # Verify ChatOpenAI was called with correct parameters
            mock_openai.assert_called_once()
            call_args = mock_openai.call_args

            assert call_args.kwargs["model"] == "gpt-4"
            assert call_args.kwargs["api_key"] == "test-openai-api-key"
            assert call_args.kwargs["temperature"] == 0.7
            assert "max_tokens" in call_args.kwargs

    def test_anthropic_api_key_configuration(
        self, settings_snapshot_with_anthropic
    ):
        """Test that Anthropic API key can be configured and used."""
        with patch(
            "local_deep_research.llm.providers.implementations.anthropic.ChatAnthropic"
        ) as mock_anthropic:
            # Create a mock LLM instance
            mock_llm_instance = _llm_mock()
            mock_anthropic.return_value = mock_llm_instance

            # Get LLM with Anthropic settings
            get_llm(settings_snapshot=settings_snapshot_with_anthropic)

            # Verify ChatAnthropic was called with correct parameters
            mock_anthropic.assert_called_once()
            call_args = mock_anthropic.call_args

            assert call_args.kwargs["model"] == "claude-3-opus-20240229"
            assert (
                call_args.kwargs["anthropic_api_key"]
                == "test-anthropic-api-key"
            )
            assert call_args.kwargs["temperature"] == 0.7
            assert "max_tokens" in call_args.kwargs

    def test_openai_endpoint_configuration(
        self, settings_snapshot_with_openai_endpoint
    ):
        """Test that OpenAI endpoint (OpenRouter) API key can be configured and used."""
        # Patch the openai_base SSRF guard to a passthrough so this config
        # test does not depend on live DNS resolution of the endpoint host
        # (matches tests/llm/test_provider_base_url_ssrf.py convention).
        with (
            patch(
                "local_deep_research.llm.providers.openai_base.ChatOpenAI"
            ) as mock_openai,
            patch(
                "local_deep_research.llm.providers.openai_base.assert_base_url_safe",
                side_effect=lambda url, **_kwargs: url,
            ),
        ):
            # Create a mock LLM instance
            mock_llm_instance = _llm_mock()
            mock_openai.return_value = mock_llm_instance

            # Get LLM with OpenAI endpoint settings
            get_llm(settings_snapshot=settings_snapshot_with_openai_endpoint)

            # Verify ChatOpenAI was called with correct parameters for endpoint
            mock_openai.assert_called_once()
            call_args = mock_openai.call_args

            assert call_args.kwargs["model"] == "claude-3-opus"
            assert call_args.kwargs["api_key"] == "test-openrouter-api-key"
            assert (
                call_args.kwargs["base_url"] == "https://openrouter.ai/api/v1"
            )
            assert call_args.kwargs["temperature"] == 0.5

    def test_llm_execution_with_api_key(self, settings_snapshot_with_openai):
        """Test that LLM can actually be invoked with API key."""
        with patch(
            "local_deep_research.llm.providers.implementations.openai.ChatOpenAI"
        ) as mock_openai:
            # Create a mock LLM instance with invoke method
            mock_llm_instance = _llm_mock()
            mock_response = Mock()
            mock_response.content = "Test response from OpenAI"
            mock_llm_instance.invoke.return_value = mock_response
            mock_openai.return_value = mock_llm_instance

            # Get LLM and test invocation
            llm = get_llm(settings_snapshot=settings_snapshot_with_openai)
            response = llm.invoke("Test prompt")

            # Verify the LLM was invoked
            mock_llm_instance.invoke.assert_called_once_with("Test prompt")
            assert "Test response from OpenAI" in response.content

    def test_multiple_provider_switching(
        self,
        settings_snapshot_with_openai,
        settings_snapshot_with_anthropic,
        settings_snapshot_with_openai_endpoint,
    ):
        """Test switching between different providers with their API keys."""
        # Test OpenAI
        with patch(
            "local_deep_research.llm.providers.implementations.openai.ChatOpenAI"
        ) as mock_openai:
            mock_openai.return_value = _llm_mock()
            get_llm(settings_snapshot=settings_snapshot_with_openai)
            assert mock_openai.called
            assert (
                mock_openai.call_args.kwargs["api_key"] == "test-openai-api-key"
            )

        # Test Anthropic
        with patch(
            "local_deep_research.llm.providers.implementations.anthropic.ChatAnthropic"
        ) as mock_anthropic:
            mock_anthropic.return_value = _llm_mock()
            get_llm(settings_snapshot=settings_snapshot_with_anthropic)
            assert mock_anthropic.called
            assert (
                mock_anthropic.call_args.kwargs["anthropic_api_key"]
                == "test-anthropic-api-key"
            )

        # Test OpenAI Endpoint
        # Patch the openai_base SSRF guard to a passthrough so this config
        # test does not depend on live DNS resolution of the endpoint host.
        with (
            patch(
                "local_deep_research.llm.providers.openai_base.ChatOpenAI"
            ) as mock_openai_endpoint,
            patch(
                "local_deep_research.llm.providers.openai_base.assert_base_url_safe",
                side_effect=lambda url, **_kwargs: url,
            ),
        ):
            mock_openai_endpoint.return_value = _llm_mock()
            get_llm(settings_snapshot=settings_snapshot_with_openai_endpoint)
            assert mock_openai_endpoint.called
            assert (
                mock_openai_endpoint.call_args.kwargs["api_key"]
                == "test-openrouter-api-key"
            )
            assert (
                mock_openai_endpoint.call_args.kwargs["base_url"]
                == "https://openrouter.ai/api/v1"
            )

    def test_research_with_api_configured_llm(
        self, settings_snapshot_with_openai_endpoint
    ):
        """Test that research can use LLM with configured API key."""
        from local_deep_research.api.research_functions import quick_summary

        # Patch the openai_base SSRF guard to a passthrough so this config
        # test does not depend on live DNS resolution of the endpoint host.
        with (
            patch(
                "local_deep_research.llm.providers.openai_base.ChatOpenAI"
            ) as mock_openai,
            patch(
                "local_deep_research.llm.providers.openai_base.assert_base_url_safe",
                side_effect=lambda url, **_kwargs: url,
            ),
        ):
            # Setup mock LLM
            mock_llm_instance = _llm_mock()
            mock_response = Mock()
            mock_response.content = "Research summary about test topic"
            mock_llm_instance.invoke.return_value = mock_response
            mock_openai.return_value = mock_llm_instance

            # Mock the search system to avoid network calls
            with patch(
                "local_deep_research.api.research_functions.AdvancedSearchSystem"
            ) as mock_search_system:
                mock_system_instance = Mock()
                mock_system_instance.analyze_topic.return_value = {
                    "current_knowledge": "Research summary about test topic",
                    "sources": ["https://example.com/test-topic"],
                    "all_links_of_system": ["https://example.com/test-topic"],
                }
                mock_search_system.return_value = mock_system_instance

                # Run research with API-configured LLM
                quick_summary(
                    query="Test research query",
                    settings_snapshot=settings_snapshot_with_openai_endpoint,
                )

                # Verify LLM was created with API key
                assert mock_openai.called
                assert (
                    mock_openai.call_args.kwargs["api_key"]
                    == "test-openrouter-api-key"
                )

    def test_api_availability_checks(self):
        """Test the availability check functions for different providers."""
        # Test OpenAI availability with settings snapshot
        settings_with_openai = {
            "llm.openai.api_key": {"value": "test-openai-key", "type": "str"}
        }
        assert (
            OpenAIProvider.is_available(settings_snapshot=settings_with_openai)
            is True
        )

        settings_without_openai = {
            "llm.openai.api_key": {"value": None, "type": "str"}
        }
        assert (
            OpenAIProvider.is_available(
                settings_snapshot=settings_without_openai
            )
            is False
        )

        # Test Anthropic availability
        settings_with_anthropic = {
            "llm.anthropic.api_key": {
                "value": "test-anthropic-key",
                "type": "str",
            }
        }
        assert (
            AnthropicProvider.is_available(
                settings_snapshot=settings_with_anthropic
            )
            is True
        )

        settings_without_anthropic = {
            "llm.anthropic.api_key": {"value": None, "type": "str"}
        }
        assert (
            AnthropicProvider.is_available(
                settings_snapshot=settings_without_anthropic
            )
            is False
        )

        # Test OpenAI endpoint availability
        settings_with_endpoint = {
            "llm.openai_endpoint.api_key": {
                "value": "test-endpoint-key",
                "type": "str",
            }
        }
        assert (
            CustomOpenAIEndpointProvider.is_available(
                settings_snapshot=settings_with_endpoint
            )
            is True
        )

        settings_without_endpoint = {
            "llm.openai_endpoint.api_key": {"value": None, "type": "str"}
        }
        assert (
            CustomOpenAIEndpointProvider.is_available(
                settings_snapshot=settings_without_endpoint
            )
            is False
        )


class TestLLMIntegration:
    """Integration tests for LLM execution with real-like scenarios."""

    def _get_base_settings(self):
        """Get base settings that all tests need."""
        return {
            "search.tool": "searxng",
            "llm.supports_max_tokens": True,
            "llm.max_tokens": 4096,
            "llm.context_window_unrestricted": False,
            "llm.context_window_size": 8192,
            "llm.local_context_window_size": 4096,
            "rate_limiting.llm_enabled": False,
        }

    def test_llm_with_token_counting(self):
        """Test LLM execution with token counting enabled."""
        settings_snapshot = self._get_base_settings()
        settings_snapshot.update(
            {
                "llm.provider": "openai",
                "llm.model": "gpt-4",
                "llm.temperature": 0.7,
                "llm.openai.api_key": "test-openai-api-key",
            }
        )

        with patch(
            "local_deep_research.llm.providers.implementations.openai.ChatOpenAI"
        ) as mock_openai:
            # Setup mock LLM with callbacks
            mock_llm_instance = _llm_mock()
            mock_llm_instance.callbacks = []
            mock_response = Mock()
            mock_response.content = "Test response"
            mock_llm_instance.invoke.return_value = mock_response
            mock_openai.return_value = mock_llm_instance

            # Get LLM with research_id for token counting
            get_llm(
                settings_snapshot=settings_snapshot,
                research_id="test-research-123",
                research_context={"phase": "testing"},
            )

            # Verify LLM was created with research_id
            assert mock_openai.called

    def test_llm_error_handling(self):
        """Test LLM error handling when API calls fail."""
        settings_snapshot = self._get_base_settings()
        settings_snapshot.update(
            {
                "llm.provider": "openai",
                "llm.model": "gpt-4",
                "llm.temperature": 0.7,
                "llm.openai.api_key": "test-openai-api-key",
            }
        )

        with patch(
            "local_deep_research.llm.providers.implementations.openai.ChatOpenAI"
        ) as mock_openai:
            # Setup mock LLM that raises error
            mock_llm_instance = _llm_mock()
            mock_llm_instance.invoke.side_effect = Exception(
                "API rate limit exceeded"
            )
            mock_openai.return_value = mock_llm_instance

            # Get LLM and test error handling
            llm = get_llm(settings_snapshot=settings_snapshot)

            with pytest.raises(Exception) as exc_info:
                llm.invoke("Test prompt")

            assert "API rate limit exceeded" in str(exc_info.value)

    def test_custom_endpoint_url_configuration(self):
        """Test configuring custom endpoint URLs for different providers."""
        settings_snapshot = self._get_base_settings()
        settings_snapshot.update(
            {
                "llm.provider": "openai_endpoint",
                "llm.model": "custom-model",
                "llm.temperature": 0.7,
                "llm.openai_endpoint.api_key": "test-key",
                "llm.openai_endpoint.url": "https://custom-llm-provider.com/v1",
                "llm.max_tokens": 2048,
            }
        )

        # This test verifies config passthrough (the configured URL/api_key
        # reach the ChatOpenAI constructor), not SSRF enforcement. The
        # ``custom-llm-provider.com`` placeholder host does not resolve, so
        # the openai_base SSRF guard (assert_base_url_safe) would fail-closed
        # on its live DNS lookup. Patch the guard to a passthrough — the same
        # convention as tests/llm/test_provider_base_url_ssrf.py — so the
        # passthrough assertions stay deterministic and offline. SSRF
        # enforcement itself is covered by test_provider_base_url_ssrf.py.
        with (
            patch(
                "local_deep_research.llm.providers.openai_base.ChatOpenAI"
            ) as mock_openai,
            patch(
                "local_deep_research.llm.providers.openai_base.assert_base_url_safe",
                side_effect=lambda url, **_kwargs: url,
            ),
        ):
            mock_openai.return_value = _llm_mock()

            get_llm(settings_snapshot=settings_snapshot)

            # Verify custom URL was used
            assert (
                mock_openai.call_args.kwargs["base_url"]
                == "https://custom-llm-provider.com/v1"
            )
            assert mock_openai.call_args.kwargs["api_key"] == "test-key"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])