learningcircuit--local-deep-research/tests/metrics/test_token_counter_logic.py

"""Tests for token_counter.py — model detection, context overflow, token extraction, and rate limiting metrics."""

import time
from unittest.mock import MagicMock, patch

import pytest
from langchain_core.outputs import LLMResult

from local_deep_research.metrics.token_counter import TokenCountingCallback


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


def _make_callback(**kw):
    """Create a TokenCountingCallback with sensible defaults."""
    return TokenCountingCallback(
        research_id=kw.get("research_id"),
        research_context=kw.get("research_context", {}),
    )


def _make_llm_result(llm_output=None, generations=None):
    """Build a minimal LLMResult for on_llm_end tests."""
    result = MagicMock(spec=LLMResult)
    result.llm_output = llm_output
    if generations is not None:
        result.generations = generations
    else:
        result.generations = []
    return result


# ===========================================================================
# 1. Model detection fallback chain
# ===========================================================================


class TestModelDetectionFallback:
    """Verify the on_llm_start model-name resolution chain."""

    def test_preset_model_takes_priority(self):
        """Preset model should override anything in serialized/kwargs."""
        cb = _make_callback()
        cb.preset_model = "my-preset-model"
        cb.preset_provider = "my-provider"

        serialized = {"kwargs": {"model": "should-be-ignored"}}
        cb.on_llm_start(
            serialized, ["hello"], invocation_params={"model": "also-ignored"}
        )

        assert cb.current_model == "my-preset-model"
        assert cb.current_provider == "my-provider"

    def test_model_from_invocation_params(self):
        """Model extracted from invocation_params when preset is absent."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["hello"],
            invocation_params={"model": "gpt-4"},
        )
        assert cb.current_model == "gpt-4"

    def test_model_from_kwargs_directly(self):
        """Model extracted from kwargs when invocation_params lacks it."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOpenAI"},
            ["hello"],
            model="gpt-3.5-turbo",
        )
        assert cb.current_model == "gpt-3.5-turbo"

    def test_model_from_serialized_kwargs(self):
        """Model extracted from serialized['kwargs'] as next fallback."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOpenAI", "kwargs": {"model": "from-serialized"}},
            ["hello"],
        )
        assert cb.current_model == "from-serialized"

    def test_model_from_serialized_name(self):
        """serialized['name'] used when kwargs has no model."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOpenAI", "name": "my-name", "kwargs": {}},
            ["hello"],
        )
        assert cb.current_model == "my-name"

    def test_ollama_specific_extraction(self):
        """ChatOllama type triggers Ollama-specific model extraction."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama3"}},
            ["hello"],
        )
        assert cb.current_model == "llama3"

    def test_ollama_fallback_to_ollama_string(self):
        """ChatOllama without model in kwargs falls back to 'ollama'."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {}},
            ["hello"],
        )
        assert cb.current_model == "ollama"

    def test_final_fallback_to_type(self):
        """When no model name found, _type string is used."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "SomeCustomLLM", "kwargs": {}},
            ["hello"],
        )
        assert cb.current_model == "SomeCustomLLM"

    def test_final_fallback_to_unknown(self):
        """When nothing at all, model is 'unknown'."""
        cb = _make_callback()
        cb.on_llm_start({}, ["hello"])
        assert cb.current_model == "unknown"


# ===========================================================================
# 2. Provider detection
# ===========================================================================


class TestProviderDetection:
    """Verify provider extraction from serialized type strings."""

    @pytest.mark.parametrize(
        "type_str, expected",
        [
            ("ChatOllama", "ollama"),
            ("ChatOpenAI", "openai"),
            ("ChatAnthropic", "anthropic"),
        ],
    )
    def test_known_providers(self, type_str, expected):
        """Known provider types are mapped correctly."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": type_str, "kwargs": {"model": "test"}},
            ["hello"],
        )
        assert cb.current_provider == expected

    def test_unknown_provider_from_kwargs(self):
        """Unknown type falls back to provider kwarg."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "CustomLLM", "kwargs": {"model": "test"}},
            ["hello"],
            provider="custom-prov",
        )
        assert cb.current_provider == "custom-prov"

    def test_unknown_provider_no_kwarg(self):
        """No _type and no provider kwarg yields 'unknown'."""
        cb = _make_callback()
        cb.on_llm_start({"kwargs": {"model": "test"}}, ["hello"])
        assert cb.current_provider == "unknown"


# ===========================================================================
# 3. Token extraction from different response formats
# ===========================================================================


class TestTokenExtraction:
    """Verify on_llm_end token extraction from various LLMResult shapes."""

    def test_tokens_from_llm_output_token_usage(self):
        """Standard token_usage dict in llm_output."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}}, ["hi"]
        )

        result = _make_llm_result(
            llm_output={
                "token_usage": {
                    "prompt_tokens": 10,
                    "completion_tokens": 20,
                    "total_tokens": 30,
                }
            },
        )
        cb.on_llm_end(result)

        assert cb.counts["total_prompt_tokens"] == 10
        assert cb.counts["total_completion_tokens"] == 20
        assert cb.counts["total_tokens"] == 30

    def test_tokens_from_llm_output_usage_key(self):
        """Alternative 'usage' key in llm_output."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}}, ["hi"]
        )

        result = _make_llm_result(
            llm_output={
                "usage": {
                    "prompt_tokens": 5,
                    "completion_tokens": 15,
                    "total_tokens": 20,
                }
            },
        )
        cb.on_llm_end(result)

        assert cb.counts["total_tokens"] == 20

    def test_tokens_from_usage_metadata_in_generations(self):
        """Ollama-style usage_metadata in generation messages."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama3"}}, ["hi"]
        )

        msg = MagicMock()
        msg.usage_metadata = {
            "input_tokens": 8,
            "output_tokens": 12,
            "total_tokens": 20,
        }
        msg.response_metadata = {}
        gen = MagicMock()
        gen.message = msg
        result = _make_llm_result(generations=[[gen]])

        cb.on_llm_end(result)
        assert cb.counts["total_tokens"] == 20

    def test_tokens_from_response_metadata_ollama(self):
        """Ollama response_metadata with prompt_eval_count / eval_count."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama3"}}, ["hi"]
        )

        msg = MagicMock()
        msg.usage_metadata = None
        msg.response_metadata = {
            "prompt_eval_count": 100,
            "eval_count": 50,
            "total_duration": 1000,
        }
        gen = MagicMock()
        gen.message = msg
        result = _make_llm_result(generations=[[gen]])

        cb.on_llm_end(result)
        assert cb.counts["total_prompt_tokens"] == 100
        assert cb.counts["total_completion_tokens"] == 50
        assert cb.counts["total_tokens"] == 150

    def test_missing_usage_entirely(self):
        """No usage info at all — counts stay zero."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}}, ["hi"]
        )

        result = _make_llm_result(llm_output=None, generations=[])
        cb.on_llm_end(result)

        assert cb.counts["total_tokens"] == 0

    def test_empty_generations_list(self):
        """generations present but empty list."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama3"}}, ["hi"]
        )

        result = _make_llm_result(llm_output=None, generations=[[]])
        cb.on_llm_end(result)

        assert cb.counts["total_tokens"] == 0

    def test_usage_metadata_present_but_response_metadata_absent(self):
        """Generation has usage_metadata but no response_metadata attr."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama3"}}, ["hi"]
        )

        msg = MagicMock(spec=["usage_metadata"])  # no response_metadata
        msg.usage_metadata = {
            "input_tokens": 3,
            "output_tokens": 7,
            "total_tokens": 10,
        }
        gen = MagicMock()
        gen.message = msg
        result = _make_llm_result(generations=[[gen]])

        cb.on_llm_end(result)
        assert cb.counts["total_tokens"] == 10

    def test_by_model_accumulation(self):
        """Multiple calls accumulate per-model stats."""
        cb = _make_callback()

        for _ in range(3):
            cb.on_llm_start(
                {"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}}, ["hi"]
            )
            result = _make_llm_result(
                llm_output={
                    "token_usage": {
                        "prompt_tokens": 10,
                        "completion_tokens": 5,
                        "total_tokens": 15,
                    }
                },
            )
            cb.on_llm_end(result)

        assert cb.counts["by_model"]["gpt-4"]["calls"] == 3
        assert cb.counts["by_model"]["gpt-4"]["total_tokens"] == 45


# ===========================================================================
# 4. Context overflow detection
# ===========================================================================


class TestContextOverflowDetection:
    """Verify context overflow detection in on_llm_end via Ollama metrics."""

    def _trigger_overflow(
        self, context_limit, prompt_eval_count, original_prompt_estimate
    ):
        """Helper to set up and trigger context overflow path."""
        cb = _make_callback(research_context={"context_limit": context_limit})
        cb.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama3"}},
            ["x" * (original_prompt_estimate * 4)],  # chars = tokens * 4
        )

        msg = MagicMock()
        msg.usage_metadata = None
        msg.response_metadata = {
            "prompt_eval_count": prompt_eval_count,
            "eval_count": 10,
        }
        gen = MagicMock()
        gen.message = msg
        result = _make_llm_result(generations=[[gen]])
        cb.on_llm_end(result)
        return cb

    def test_overflow_detected_at_95_percent(self):
        """Context overflow flagged when prompt_eval_count >= 95% of limit."""
        cb = self._trigger_overflow(
            context_limit=1000,
            prompt_eval_count=960,  # 96% > 95%
            original_prompt_estimate=1200,
        )
        assert cb.context_truncated is True
        assert cb.tokens_truncated == 1200 - 960

    def test_no_overflow_below_threshold(self):
        """No overflow when below 80% threshold AND estimate fits limit.

        Both detection paths must NOT fire: prompt_eval_count is 70% of
        the limit (below the 80% provider-confirmed threshold), and
        original_prompt_estimate is below context_limit (so the
        estimation-based path added in PR #3791 also stays quiet).
        """
        cb = self._trigger_overflow(
            context_limit=1000,
            prompt_eval_count=700,  # 70% < 80%
            original_prompt_estimate=900,  # below limit, no estimation path
        )
        assert cb.context_truncated is False

    def test_exact_95_boundary(self):
        """Exact 95% threshold should trigger overflow."""
        cb = self._trigger_overflow(
            context_limit=1000,
            prompt_eval_count=950,  # exactly 95%
            original_prompt_estimate=1200,
        )
        assert cb.context_truncated is True

    def test_no_overflow_when_context_limit_none(self):
        """No overflow detection when context_limit is not set."""
        cb = _make_callback(research_context={})
        cb.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama3"}},
            ["hello"],
        )

        msg = MagicMock()
        msg.usage_metadata = None
        msg.response_metadata = {"prompt_eval_count": 9999, "eval_count": 10}
        gen = MagicMock()
        gen.message = msg
        result = _make_llm_result(generations=[[gen]])
        cb.on_llm_end(result)

        assert cb.context_truncated is False

    def test_truncation_ratio_zero_prompt_estimate(self):
        """When original_prompt_estimate <= prompt_eval_count, tokens_truncated is 0."""
        cb = self._trigger_overflow(
            context_limit=100,
            prompt_eval_count=96,  # 96% > 95%
            original_prompt_estimate=90,  # less than eval count
        )
        # context_truncated is set, but tokens_truncated stays 0
        # because original_prompt_estimate <= prompt_eval_count
        assert cb.context_truncated is True
        assert cb.tokens_truncated == 0

    def test_prompt_eval_count_zero_skips_overflow(self):
        """prompt_eval_count == 0 should not trigger overflow check."""
        cb = _make_callback(research_context={"context_limit": 1000})
        cb.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "llama3"}},
            ["hello"],
        )

        msg = MagicMock()
        msg.usage_metadata = None
        msg.response_metadata = {"prompt_eval_count": 0, "eval_count": 10}
        gen = MagicMock()
        gen.message = msg
        result = _make_llm_result(generations=[[gen]])
        cb.on_llm_end(result)

        assert cb.context_truncated is False


# ===========================================================================
# 5. on_llm_error tracking
# ===========================================================================


class TestOnLLMError:
    """Verify error tracking in on_llm_error."""

    def test_error_sets_status_and_type(self):
        """Error callback records status and error type."""
        cb = _make_callback()
        cb.start_time = time.time()

        cb.on_llm_error(ValueError("boom"))

        assert cb.success_status == "error"
        assert cb.error_type == "ValueError"
        assert cb.response_time_ms is not None

    def test_error_saves_to_db_when_research_id_set(self):
        """Error with research_id triggers _save_to_db with zero tokens."""
        cb = _make_callback(research_id="test-123")
        cb.start_time = time.time()

        with patch.object(cb, "_save_to_db") as mock_save:
            cb.on_llm_error(RuntimeError("fail"))
            mock_save.assert_called_once_with(0, 0)


# ===========================================================================
# 6. _get_context_overflow_fields
# ===========================================================================


class TestContextOverflowFields:
    """Verify _get_context_overflow_fields output."""

    def test_fields_when_no_overflow(self):
        """Fields should indicate no truncation."""
        cb = _make_callback()
        fields = cb._get_context_overflow_fields()

        assert fields["context_truncated"] is False
        assert fields["tokens_truncated"] is None
        assert fields["truncation_ratio"] is None

    def test_fields_when_overflow_detected(self):
        """Fields should include truncation details."""
        cb = _make_callback()
        cb.context_limit = 1000
        cb.context_truncated = True
        cb.tokens_truncated = 200
        cb.truncation_ratio = 0.2
        cb.ollama_metrics = {"prompt_eval_count": 800, "eval_count": 50}

        fields = cb._get_context_overflow_fields()

        assert fields["context_truncated"] is True
        assert fields["tokens_truncated"] == 200
        assert fields["truncation_ratio"] == 0.2
        assert fields["ollama_prompt_eval_count"] == 800


# ===========================================================================
# 7. Prompt estimate calculation
# ===========================================================================


class TestPromptEstimate:
    """Verify original_prompt_estimate is calculated from prompt chars."""

    def test_estimate_from_multiple_prompts(self):
        """Estimate is sum of chars // 4."""
        cb = _make_callback()
        cb.on_llm_start(
            {}, ["aaaa", "bbbbbbbb"]
        )  # 4 + 8 = 12 chars => 3 tokens
        assert cb.original_prompt_estimate == 3

    def test_estimate_empty_prompts(self):
        """Empty prompts list yields 0 estimate."""
        cb = _make_callback()
        cb.on_llm_start({}, [])
        assert cb.original_prompt_estimate == 0