learningcircuit--local-deep-research/tests/config/test_llm_config_coverage.py

"""
Comprehensive tests for provider availability functions and related helpers
in local_deep_research/config/llm_config.py.

Focuses on gaps not covered by existing test files:
- wrap_llm_without_think_tags() context_limit injection, token counter, string responses
- _get_context_window_for_provider() edge cases
- get_selected_llm_provider() with snapshot parameter
"""

from unittest.mock import MagicMock, patch

import pytest

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

MODULE = "local_deep_research.config.llm_config"
# _get_context_window_for_provider delegates to the _helpers twin, which reads
# settings via thread_settings.get_setting_from_snapshot (function-local
# import). Context-window assertions must therefore patch this module, not
# MODULE.
THREAD_SETTINGS = "local_deep_research.config.thread_settings"


# ===================================================================
# get_selected_llm_provider
# ===================================================================


class TestGetSelectedLlmProvider:
    """Additional coverage for get_selected_llm_provider()."""

    def test_with_explicit_snapshot(self):
        from local_deep_research.config.llm_config import (
            get_selected_llm_provider,
        )

        result = get_selected_llm_provider(
            settings_snapshot={"llm.provider": "Google"}
        )
        assert result == "google"

    def test_mixed_case_normalised(self):
        from local_deep_research.config.llm_config import (
            get_selected_llm_provider,
        )

        result = get_selected_llm_provider(
            settings_snapshot={"llm.provider": "OpenAI_Endpoint"}
        )
        assert result == "openai_endpoint"

    def test_default_when_key_missing(self):
        from local_deep_research.config.llm_config import (
            get_selected_llm_provider,
        )

        result = get_selected_llm_provider(settings_snapshot={})
        assert result == "ollama"


# ===================================================================
# _get_context_window_for_provider  (additional edge cases)
# ===================================================================


class TestGetContextWindowForProvider:
    """Additional edge-case coverage for _get_context_window_for_provider."""

    def test_openrouter_treated_as_cloud(self):
        with patch(
            f"{THREAD_SETTINGS}.get_setting_from_snapshot",
            return_value=True,
        ):
            from local_deep_research.config.llm_config import (
                _get_context_window_for_provider,
            )

            assert _get_context_window_for_provider("openrouter") is None

    def test_local_provider_with_float_value(self):
        """Float value from settings is coerced to int."""
        with patch(
            f"{THREAD_SETTINGS}.get_setting_from_snapshot",
            return_value=8192.7,
        ):
            from local_deep_research.config.llm_config import (
                _get_context_window_for_provider,
            )

            result = _get_context_window_for_provider("lmstudio")
            assert result == 8192
            assert isinstance(result, int)

    def test_cloud_restricted_with_float_value(self):
        """Float cloud window is coerced to int."""
        call_num = [0]

        def fake_setting(key, default, settings_snapshot=None):
            call_num[0] += 1
            if key == "llm.context_window_unrestricted":
                return False
            if key == "llm.context_window_size":
                return 65536.9
            return default

        with patch(
            f"{THREAD_SETTINGS}.get_setting_from_snapshot",
            side_effect=fake_setting,
        ):
            from local_deep_research.config.llm_config import (
                _get_context_window_for_provider,
            )

            result = _get_context_window_for_provider("openai")
            assert result == 65536
            assert isinstance(result, int)


# ===================================================================
# wrap_llm_without_think_tags
# ===================================================================


class TestWrapLlmWithoutThinkTags:
    """Comprehensive tests for the ProcessingLLMWrapper created by wrap_llm_without_think_tags."""

    def _make_wrapper(self, mock_llm, **kwargs):
        """Create wrapper with rate-limiting disabled."""
        defaults = {
            "research_id": None,
            "provider": None,
            "research_context": None,
            "settings_snapshot": None,
        }
        defaults.update(kwargs)
        with patch(
            f"{MODULE}.get_setting_from_snapshot",
            return_value=False,
        ):
            from local_deep_research.config.llm_config import (
                wrap_llm_without_think_tags,
            )

            return wrap_llm_without_think_tags(mock_llm, **defaults)

    # --- basic wrapper behaviour ---

    def test_wrapper_has_base_llm(self):
        llm = MagicMock()
        w = self._make_wrapper(llm)
        assert w.base_llm is llm

    def test_invoke_calls_base_llm(self):
        llm = MagicMock()
        resp = MagicMock()
        resp.content = "hello"
        llm.invoke.return_value = resp
        w = self._make_wrapper(llm)
        w.invoke("prompt")
        llm.invoke.assert_called_once_with("prompt")

    def test_think_tags_removed_from_content(self):
        llm = MagicMock()
        resp = MagicMock()
        resp.content = "<think>internal reasoning</think>final answer"
        llm.invoke.return_value = resp
        w = self._make_wrapper(llm)
        result = w.invoke("prompt")
        assert "<think>" not in result.content
        assert "final answer" in result.content

    def test_think_tags_removed_from_string_response(self):
        llm = MagicMock()
        llm.invoke.return_value = "<think>thought</think>answer"
        w = self._make_wrapper(llm)
        result = w.invoke("prompt")
        # A bare-string return is wrapped into a message so callers can rely on
        # .content; think tags are still stripped.
        assert not isinstance(result, str)
        assert "<think>" not in result.content
        assert "answer" in result.content

    def test_response_without_content_attr_returned_as_is(self):
        """Response that is neither string nor has .content is passed through."""
        llm = MagicMock()
        resp = 42  # int has no .content
        llm.invoke.return_value = resp
        w = self._make_wrapper(llm)
        result = w.invoke("prompt")
        assert result == 42

    def test_string_response_wrapped_in_message(self):
        """A bare-string return is wrapped into an AIMessage with .content."""
        from langchain_core.messages import AIMessage

        llm = MagicMock()
        llm.invoke.return_value = "<think>t</think>final"
        w = self._make_wrapper(llm)
        result = w.invoke("prompt")
        assert isinstance(result, AIMessage)
        assert result.content == "final"

    def test_preserves_reasoning_content_and_tool_calls(self):
        """Stripping <think> from .content must NOT drop reasoning_content/tool_calls.

        Guards against worsening DeepSeek thinking-mode round-tripping (#4194):
        we only rewrite .content in place, leaving the rest of the message intact.
        """
        from langchain_core.messages import AIMessage

        llm = MagicMock()
        llm.invoke.return_value = AIMessage(
            content="<think>reasoning</think>answer",
            additional_kwargs={"reasoning_content": "R"},
            tool_calls=[
                {"name": "search", "args": {}, "id": "1", "type": "tool_call"}
            ],
        )
        w = self._make_wrapper(llm)
        result = w.invoke("prompt")
        assert result.content == "answer"
        assert result.additional_kwargs["reasoning_content"] == "R"
        assert result.tool_calls and result.tool_calls[0]["name"] == "search"

    def test_ainvoke_normalizes_string_response(self):
        """ainvoke applies the same normalization as invoke (str -> message)."""
        import asyncio
        from unittest.mock import AsyncMock

        llm = MagicMock()
        llm.ainvoke = AsyncMock(return_value="<think>t</think>async answer")
        w = self._make_wrapper(llm)
        result = asyncio.run(w.ainvoke("prompt"))
        assert not isinstance(result, str)
        assert result.content == "async answer"

    def test_invoke_exception_propagated(self):
        llm = MagicMock()
        llm.invoke.side_effect = ConnectionError("timeout")
        w = self._make_wrapper(llm)
        with pytest.raises(ConnectionError, match="timeout"):
            w.invoke("prompt")

    # --- __getattr__ delegation ---

    def test_getattr_delegates_to_base_llm(self):
        llm = MagicMock()
        llm.model_name = "gpt-4"
        llm.some_custom_attr = "custom_value"
        w = self._make_wrapper(llm)
        assert w.model_name == "gpt-4"
        assert w.some_custom_attr == "custom_value"

    # --- context_limit injection ---

    def test_context_limit_set_in_research_context(self):
        """wrap_llm sets context_limit in research_context when provider is local."""
        llm = MagicMock()
        research_ctx = {}

        def fake_setting(key, default=None, settings_snapshot=None):
            if key == "rate_limiting.llm_enabled":
                return False
            if key == "llm.local_context_window_size":
                return 4096
            return default

        with (
            patch(
                f"{MODULE}.get_setting_from_snapshot", side_effect=fake_setting
            ),
            patch(
                f"{THREAD_SETTINGS}.get_setting_from_snapshot",
                side_effect=fake_setting,
            ),
        ):
            from local_deep_research.config.llm_config import (
                wrap_llm_without_think_tags,
            )

            wrap_llm_without_think_tags(
                llm, provider="ollama", research_context=research_ctx
            )
        assert research_ctx.get("context_limit") == 4096

    def test_context_limit_not_overwritten_if_already_set(self):
        """If research_context already has context_limit, it should NOT be overwritten."""
        llm = MagicMock()
        research_ctx = {"context_limit": 9999}

        def fake_setting(key, default=None, settings_snapshot=None):
            if key == "rate_limiting.llm_enabled":
                return False
            if key == "llm.local_context_window_size":
                return 4096
            return default

        with (
            patch(
                f"{MODULE}.get_setting_from_snapshot", side_effect=fake_setting
            ),
            patch(
                f"{THREAD_SETTINGS}.get_setting_from_snapshot",
                side_effect=fake_setting,
            ),
        ):
            from local_deep_research.config.llm_config import (
                wrap_llm_without_think_tags,
            )

            wrap_llm_without_think_tags(
                llm, provider="ollama", research_context=research_ctx
            )
        assert research_ctx["context_limit"] == 9999

    def test_context_limit_not_set_for_unrestricted_cloud(self):
        """Cloud unrestricted provider returns None window, so context_limit not set."""
        llm = MagicMock()
        research_ctx = {}

        def fake_setting(key, default=None, settings_snapshot=None):
            if key == "rate_limiting.llm_enabled":
                return False
            if key == "llm.context_window_unrestricted":
                return True
            return default

        with (
            patch(
                f"{MODULE}.get_setting_from_snapshot", side_effect=fake_setting
            ),
            patch(
                f"{THREAD_SETTINGS}.get_setting_from_snapshot",
                side_effect=fake_setting,
            ),
        ):
            from local_deep_research.config.llm_config import (
                wrap_llm_without_think_tags,
            )

            wrap_llm_without_think_tags(
                llm, provider="openai", research_context=research_ctx
            )
        assert "context_limit" not in research_ctx

    def test_no_crash_when_research_context_is_none(self):
        """No crash when research_context=None."""
        llm = MagicMock()
        w = self._make_wrapper(llm, provider="openai", research_context=None)
        assert w is not None

    # --- rate limiting integration ---

    def test_rate_limiting_applied_when_enabled(self):
        llm = MagicMock()
        rate_limited_llm = MagicMock()

        with (
            patch(
                f"{MODULE}.get_setting_from_snapshot",
                return_value=True,
            ),
            patch(
                "local_deep_research.web_search_engines.rate_limiting.llm.create_rate_limited_llm_wrapper",
                return_value=rate_limited_llm,
            ) as mock_rl,
        ):
            from local_deep_research.config.llm_config import (
                wrap_llm_without_think_tags,
            )

            wrapper = wrap_llm_without_think_tags(llm, provider="openai")
            mock_rl.assert_called_once_with(llm, "openai")
            # The wrapper wraps the rate-limited LLM
            assert wrapper.base_llm is rate_limited_llm

    def test_rate_limiting_not_applied_when_disabled(self):
        llm = MagicMock()

        with (
            patch(
                f"{MODULE}.get_setting_from_snapshot",
                return_value=False,
            ),
            patch(
                "local_deep_research.web_search_engines.rate_limiting.llm.create_rate_limited_llm_wrapper",
            ) as mock_rl,
        ):
            from local_deep_research.config.llm_config import (
                wrap_llm_without_think_tags,
            )

            wrapper = wrap_llm_without_think_tags(llm, provider="openai")
            mock_rl.assert_not_called()
            assert wrapper.base_llm is llm

    # --- token counter callback ---

    def test_token_counter_attached_when_research_id_given(self):
        """When research_id is set, a token counting callback is added."""
        llm = MagicMock()
        llm.callbacks = None
        llm.model_name = "test-model"

        mock_counter = MagicMock()
        mock_callback = MagicMock()
        mock_counter.create_callback.return_value = mock_callback

        with (
            patch(
                f"{MODULE}.get_setting_from_snapshot",
                return_value=False,
            ),
            patch(
                "local_deep_research.metrics.TokenCounter",
                return_value=mock_counter,
            ),
        ):
            from local_deep_research.config.llm_config import (
                wrap_llm_without_think_tags,
            )

            wrap_llm_without_think_tags(llm, research_id=42, provider="openai")
        mock_counter.create_callback.assert_called_once_with(42, None)
        assert mock_callback.preset_provider == "openai"
        assert mock_callback.preset_model == "test-model"

    def test_token_counter_uses_model_attr_fallback(self):
        """If llm has .model but not .model_name, uses .model."""
        llm = MagicMock(spec=["invoke", "callbacks", "model"])
        llm.callbacks = None
        llm.model = "claude-3-opus"

        mock_counter = MagicMock()
        mock_callback = MagicMock()
        mock_counter.create_callback.return_value = mock_callback

        with (
            patch(
                f"{MODULE}.get_setting_from_snapshot",
                return_value=False,
            ),
            patch(
                "local_deep_research.metrics.TokenCounter",
                return_value=mock_counter,
            ),
        ):
            from local_deep_research.config.llm_config import (
                wrap_llm_without_think_tags,
            )

            wrap_llm_without_think_tags(
                llm, research_id=1, provider="anthropic"
            )
        assert mock_callback.preset_model == "claude-3-opus"

    def test_callbacks_extended_when_existing(self):
        """If llm.callbacks already has entries, new callback is appended."""
        existing_cb = MagicMock()
        llm = MagicMock()
        llm.callbacks = [existing_cb]
        llm.model_name = "m"

        mock_counter = MagicMock()
        mock_callback = MagicMock()
        mock_counter.create_callback.return_value = mock_callback

        with (
            patch(
                f"{MODULE}.get_setting_from_snapshot",
                return_value=False,
            ),
            patch(
                "local_deep_research.metrics.TokenCounter",
                return_value=mock_counter,
            ),
        ):
            from local_deep_research.config.llm_config import (
                wrap_llm_without_think_tags,
            )

            wrap_llm_without_think_tags(llm, research_id=10)
        assert mock_callback in llm.callbacks
        assert existing_cb in llm.callbacks

    def test_no_callbacks_when_no_research_id(self):
        """When research_id is None, no callbacks are attached."""
        llm = MagicMock()
        llm.callbacks = None

        with patch(
            f"{MODULE}.get_setting_from_snapshot",
            return_value=False,
        ):
            from local_deep_research.config.llm_config import (
                wrap_llm_without_think_tags,
            )

            wrap_llm_without_think_tags(llm, research_id=None)
        # callbacks should remain None (nothing to attach)
        assert llm.callbacks is None