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

"""Tests for custom LLM integration with benchmarking system."""

from typing import List
from unittest.mock import 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.llm import clear_llm_registry, register_llm


class BenchmarkLLM(BaseChatModel):
    """LLM designed for benchmark testing."""

    correct_answers: dict = Field(default_factory=dict)

    def _generate(self, messages: List[BaseMessage], **kwargs) -> ChatResult:
        """Generate responses based on predefined correct answers."""
        query = messages[-1].content if messages else ""

        # Check if we have a predefined answer
        for key, answer in self.correct_answers.items():
            if key.lower() in query.lower():
                message = AIMessage(content=answer)
                generation = ChatGeneration(message=message)
                return ChatResult(generations=[generation])

        # Default response
        message = AIMessage(content="I don't know")
        generation = ChatGeneration(message=message)
        return ChatResult(generations=[generation])

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


@pytest.fixture(autouse=True)
def clear_registry():
    """Clear the registry before and after each test."""
    clear_llm_registry()
    yield
    clear_llm_registry()


def test_custom_llm_with_benchmarks():
    """Test that custom LLMs work with the benchmark system."""
    # Create a benchmark LLM with some correct answers
    benchmark_llm = BenchmarkLLM(
        correct_answers={
            "capital of France": "Paris",
            "2+2": "4",
            "meaning of life": "42",
        }
    )

    register_llm("benchmark_llm", benchmark_llm)

    # Mock the benchmark flow
    with patch("local_deep_research.config.llm_config.get_llm") as mock_get_llm:
        # Return our benchmark LLM when requested
        mock_get_llm.return_value = benchmark_llm

        # Simulate benchmark questions
        questions = [
            ("What is the capital of France?", "Paris"),
            ("What is 2+2?", "4"),
            ("What is the meaning of life?", "42"),
            ("Unknown question", "I don't know"),
        ]

        for question, expected in questions:
            from langchain_core.messages import HumanMessage

            result = benchmark_llm._generate([HumanMessage(content=question)])
            assert result.generations[0].message.content == expected


def test_benchmark_llm_metrics():
    """Test that custom LLMs properly track metrics for benchmarks."""

    class MetricsLLM(BaseChatModel):
        """LLM that tracks metrics."""

        total_tokens: int = Field(default=0)
        call_count: int = Field(default=0)

        def _generate(
            self, messages: List[BaseMessage], **kwargs
        ) -> ChatResult:
            self.call_count += 1
            response = f"Response {self.call_count}"
            self.total_tokens += len(response.split())

            message = AIMessage(content=response)
            generation = ChatGeneration(
                message=message,
                generation_info={"token_count": len(response.split())},
            )
            return ChatResult(generations=[generation])

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

    metrics_llm = MetricsLLM()
    register_llm("metrics_llm", metrics_llm)

    # Simulate multiple benchmark runs
    from langchain_core.messages import HumanMessage

    for i in range(5):
        metrics_llm._generate([HumanMessage(content=f"Query {i}")])

    assert metrics_llm.call_count == 5
    assert metrics_llm.total_tokens == 10  # "Response X" = 2 tokens each


def test_custom_llm_accuracy_scoring():
    """Test custom LLMs with accuracy scoring."""

    class ScoringLLM(BaseChatModel):
        """LLM that returns scores with answers."""

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

            # Return answer with confidence score
            if "easy" in query.lower():
                response = "Easy answer"
                confidence = 1.0
            elif "hard" in query.lower():
                response = "Hard answer"
                confidence = 0.5
            else:
                response = "Unknown"
                confidence = 0.1

            message = AIMessage(content=response)
            generation = ChatGeneration(
                message=message, generation_info={"confidence": confidence}
            )
            return ChatResult(generations=[generation])

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

    scoring_llm = ScoringLLM()
    register_llm("scoring_llm", scoring_llm)

    # Test different query types
    from langchain_core.messages import HumanMessage

    easy_result = scoring_llm._generate([HumanMessage(content="Easy question")])
    assert easy_result.generations[0].generation_info["confidence"] == 1.0

    hard_result = scoring_llm._generate([HumanMessage(content="Hard question")])
    assert hard_result.generations[0].generation_info["confidence"] == 0.5

    unknown_result = scoring_llm._generate(
        [HumanMessage(content="Random question")]
    )
    assert unknown_result.generations[0].generation_info["confidence"] == 0.1


def test_benchmark_with_custom_llm_factory():
    """Test benchmarking with LLM factories."""
    factory_calls = []

    def create_benchmark_llm(model_name=None, temperature=0.7, **kwargs):
        """Factory that creates benchmark-optimized LLMs."""
        factory_calls.append(
            {"model_name": model_name, "temperature": temperature}
        )

        # Return different LLMs based on model name
        if model_name == "accurate":
            return BenchmarkLLM(correct_answers={"test": "correct"})
        return BenchmarkLLM(correct_answers={})

    register_llm("benchmark_factory", create_benchmark_llm)

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

        # Test with accurate model
        with patch(
            "local_deep_research.llm.is_llm_registered",
            return_value=True,
        ):
            with patch(
                "local_deep_research.llm.get_llm_from_registry",
                return_value=create_benchmark_llm,
            ):
                from local_deep_research.config.llm_config import get_llm

                # A settings snapshot (default "both" scope) is required so
                # the egress-policy PEP lets a non-local registered provider
                # reach the registry dispatch this test exercises.
                accurate_llm = get_llm(
                    provider="benchmark_factory",
                    model_name="accurate",
                    temperature=0.1,
                    settings_snapshot={"search.tool": "searxng"},
                )

                # Should create the accurate version
                from langchain_core.messages import HumanMessage

                result = accurate_llm._generate(
                    [HumanMessage(content="test question")]
                )
                assert result.generations[0].message.content == "correct"


def test_benchmark_comparison_with_custom_llms():
    """Test comparing multiple custom LLMs in benchmarks."""
    # Create multiple LLMs with different characteristics
    fast_llm = BenchmarkLLM(correct_answers={"quick": "fast response"})
    accurate_llm = BenchmarkLLM(
        correct_answers={
            "quick": "fast response",
            "complex": "detailed accurate response",
        }
    )

    register_llm("fast_llm", fast_llm)
    register_llm("accurate_llm", accurate_llm)

    # Simulate benchmark comparison
    test_queries = ["quick question", "complex problem"]

    results = {}
    for llm_name, llm in [
        ("fast_llm", fast_llm),
        ("accurate_llm", accurate_llm),
    ]:
        results[llm_name] = []
        for query in test_queries:
            from langchain_core.messages import HumanMessage

            response = llm._generate([HumanMessage(content=query)])
            results[llm_name].append(response.generations[0].message.content)

    # Fast LLM should only answer one correctly
    assert results["fast_llm"][0] == "fast response"
    assert results["fast_llm"][1] == "I don't know"

    # Accurate LLM should answer both
    assert results["accurate_llm"][0] == "fast response"
    assert results["accurate_llm"][1] == "detailed accurate response"