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

"""Tests for edge cases and advanced scenarios in custom LLM integration."""

import asyncio
from typing import Any, Iterator, List, Optional, Dict
from unittest.mock import patch

import pytest
from langchain_core.callbacks import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage
from langchain_core.outputs import (
    ChatGeneration,
    ChatGenerationChunk,
    ChatResult,
)
from pydantic import Field

from local_deep_research.config.llm_config import get_llm
from local_deep_research.llm import (
    clear_llm_registry,
    get_llm_from_registry,
    register_llm,
)


class StreamingLLM(BaseChatModel):
    """LLM that supports streaming."""

    chunks: List[str] = Field(
        default_factory=lambda: [
            "Hello",
            " world",
            " from",
            " streaming",
            " LLM",
        ]
    )

    def _generate(self, messages: List[BaseMessage], **kwargs) -> ChatResult:
        """Generate non-streaming response."""
        full_response = "".join(self.chunks)
        message = AIMessage(content=full_response)
        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])

    def _stream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[ChatGenerationChunk]:
        """Stream response chunks."""
        for chunk in self.chunks:
            yield ChatGenerationChunk(message=AIMessageChunk(content=chunk))

    @property
    def _llm_type(self) -> str:
        return "streaming"


class AsyncLLM(BaseChatModel):
    """LLM that supports async operations."""

    response: str = Field(default="Async response")
    delay: float = Field(default=0.1)

    def _generate(self, messages: List[BaseMessage], **kwargs) -> ChatResult:
        """Sync generation (fallback)."""
        message = AIMessage(content=self.response)
        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])

    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        """Async generation."""
        await asyncio.sleep(self.delay)  # Simulate async work
        message = AIMessage(content=f"Async: {self.response}")
        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])

    @property
    def _llm_type(self) -> str:
        return "async"


class BrokenLLM(BaseChatModel):
    """LLM that raises errors."""

    error_message: str = Field(default="LLM is broken")

    def _generate(self, messages: List[BaseMessage], **kwargs) -> ChatResult:
        """Always raises an error."""
        raise RuntimeError(self.error_message)

    @property
    def _llm_type(self) -> str:
        return "broken"


class SlowLLM(BaseChatModel):
    """LLM that responds slowly."""

    response: str = Field(default="Slow response")
    delay: float = Field(default=2.0)

    def _generate(self, messages: List[BaseMessage], **kwargs) -> ChatResult:
        """Generate response with delay."""
        import time

        time.sleep(self.delay)
        message = AIMessage(content=self.response)
        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])

    @property
    def _llm_type(self) -> str:
        return "slow"


class MalformedResponseLLM(BaseChatModel):
    """LLM that returns malformed responses."""

    def _generate(self, messages: List[BaseMessage], **kwargs) -> ChatResult:
        """Return invalid response structure."""
        # Return a ChatResult with no generations
        return ChatResult(generations=[])

    @property
    def _llm_type(self) -> str:
        return "malformed"


@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: str | None) -> Dict[str, Any]:
        return {"value": value, "ui_element": "text"}

    def _number_param(value: int | float | None) -> Dict[str, Any]:
        return {"value": value, "ui_element": "number"}

    def _bool_param(value: bool) -> Dict[str, Any]:
        return {"value": value, "ui_element": "checkbox"}

    return {
        "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"),
    }


def test_streaming_llm_registration(full_settings_snapshot):
    """Test registering an LLM with streaming support."""
    streaming_llm = StreamingLLM()
    register_llm("streaming", streaming_llm)

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

        # Provide settings_snapshot to avoid settings context error
        llm = get_llm(
            provider="streaming", settings_snapshot=full_settings_snapshot
        )
        assert isinstance(llm, StreamingLLM)

        # Test that streaming works
        chunks = list(llm._stream([]))
        assert len(chunks) == 5
        assert all(hasattr(chunk, "message") for chunk in chunks)


@pytest.mark.asyncio
async def test_async_llm_operations():
    """Test async LLM operations."""
    async_llm = AsyncLLM(delay=0.05)
    register_llm("async", async_llm)

    # Test async generation
    result = await async_llm._agenerate([])
    assert result.generations[0].message.content == "Async: Async response"


def test_broken_llm_error_handling(full_settings_snapshot):
    """Test handling of LLMs that raise errors."""
    broken_llm = BrokenLLM(error_message="Test error")
    register_llm("broken", broken_llm)

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

        # Provide settings_snapshot to avoid settings context error
        llm = get_llm(
            provider="broken", settings_snapshot=full_settings_snapshot
        )

        # Should raise the error when trying to generate
        with pytest.raises(RuntimeError, match="Test error"):
            llm._generate([])


def test_malformed_response_handling(full_settings_snapshot):
    """Test handling of LLMs that return malformed responses."""
    malformed_llm = MalformedResponseLLM()
    register_llm("malformed", malformed_llm)

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

        # Provide settings_snapshot to avoid settings context error
        llm = get_llm(
            provider="malformed", settings_snapshot=full_settings_snapshot
        )
        result = llm._generate([])

        # Should return empty generations
        assert len(result.generations) == 0


def test_concurrent_llm_registration():
    """Test concurrent registration and usage of LLMs."""
    import threading

    errors = []

    def register_and_use(name: str, response: str):
        try:
            llm = SlowLLM(response=response, delay=0.01)
            register_llm(name, llm)

            # Try to use it
            retrieved = get_llm_from_registry(name)
            assert retrieved is not None
        except Exception as e:
            errors.append(e)

    # Start multiple threads
    threads = []
    for i in range(10):
        t = threading.Thread(
            target=register_and_use, args=(f"llm_{i}", f"Response {i}")
        )
        threads.append(t)
        t.start()

    # Wait for all to complete
    for t in threads:
        t.join()

    assert len(errors) == 0

    # Verify all were registered
    for i in range(10):
        assert get_llm_from_registry(f"llm_{i}") is not None


def test_provider_name_normalization(full_settings_snapshot):
    """Test that provider names are normalized correctly."""
    llm = SlowLLM()

    # Register with mixed case - the registry stores as-is
    register_llm("MixedCase", llm)

    # The registry is case-sensitive, so we need to register with lowercase too
    register_llm("mixedcase", llm)

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

        # The provider name should be normalized to lowercase
        result = get_llm(
            provider="mixedcase", settings_snapshot=full_settings_snapshot
        )
        assert isinstance(result, SlowLLM)


def test_factory_with_invalid_signature(full_settings_snapshot):
    """Test factory functions with invalid signatures."""

    def bad_factory():
        # Missing required parameters
        return SlowLLM()

    register_llm("bad_factory", bad_factory)

    # Should raise error when trying to use with parameters. A settings
    # snapshot is required so the egress-policy PEP (which fails closed for
    # snapshot-less non-local providers) lets the call reach the factory
    # dispatch, where the signature mismatch is what we are exercising.
    with pytest.raises(TypeError):
        get_llm(
            provider="bad_factory",
            model_name="test",
            temperature=0.5,
            settings_snapshot=full_settings_snapshot,
        )


def test_llm_memory_cleanup():
    """Test that LLMs are properly cleaned up."""
    import gc
    import weakref

    # Create an LLM and register it
    llm = SlowLLM()
    weak_ref = weakref.ref(llm)
    register_llm("memory_test", llm)

    # Delete the original reference
    del llm

    # Force garbage collection
    gc.collect()

    # The LLM should still exist because registry holds a reference
    assert weak_ref() is not None

    # Clear the registry
    clear_llm_registry()
    gc.collect()

    # Now it should be garbage collected
    assert weak_ref() is None


def test_llm_with_custom_attributes():
    """Test LLMs with custom attributes and methods."""

    class CustomLLM(BaseChatModel):
        custom_attr: str = Field(default="custom")

        def custom_method(self):
            return "custom_result"

        def _generate(
            self, messages: List[BaseMessage], **kwargs
        ) -> ChatResult:
            message = AIMessage(content=f"Custom: {self.custom_attr}")
            generation = ChatGeneration(message=message)
            return ChatResult(generations=[generation])

        @property
        def _llm_type(self) -> str:
            return "custom_attrs"

    custom_llm = CustomLLM(custom_attr="test_value")
    register_llm("custom_attrs", custom_llm)

    retrieved = get_llm_from_registry("custom_attrs")
    assert retrieved.custom_attr == "test_value"
    assert retrieved.custom_method() == "custom_result"


def test_llm_state_persistence(full_settings_snapshot):
    """Test that LLM state persists across calls."""

    class StatefulLLM(BaseChatModel):
        call_count: int = Field(default=0)
        history: List[str] = Field(default_factory=list)

        def _generate(
            self, messages: List[BaseMessage], **kwargs
        ) -> ChatResult:
            self.call_count += 1
            query = messages[-1].content if messages else "empty"
            self.history.append(query)

            message = AIMessage(content=f"Call {self.call_count}: {query}")
            generation = ChatGeneration(message=message)
            return ChatResult(generations=[generation])

        @property
        def _llm_type(self) -> str:
            return "stateful"

    stateful_llm = StatefulLLM()
    register_llm("stateful", stateful_llm)

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

        # Provide settings_snapshot to avoid settings context error
        llm1 = get_llm(
            provider="stateful", settings_snapshot=full_settings_snapshot
        )
        llm2 = get_llm(
            provider="stateful", settings_snapshot=full_settings_snapshot
        )

        # Should be the same instance
        assert llm1 is llm2

        # State should persist
        from langchain_core.messages import HumanMessage

        llm1._generate([HumanMessage(content="First")])
        llm2._generate([HumanMessage(content="Second")])

        assert stateful_llm.call_count == 2
        assert stateful_llm.history == ["First", "Second"]