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

"""High-value pure logic tests for TokenCountingCallback.

Covers initialization, on_llm_start model/provider detection,
on_llm_end token extraction and accumulation, on_llm_error tracking,
and context overflow detection.
"""

import time
from unittest.mock import MagicMock

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 mock LLMResult."""
    result = MagicMock(spec=LLMResult)
    result.llm_output = llm_output
    result.generations = generations if generations is not None else []
    return result


def _setup_model(cb, model="test-model", provider="unknown"):
    """Run on_llm_start so current_model and by_model are initialised."""
    cb.on_llm_start(
        {"kwargs": {"model": model}},
        ["prompt"],
    )


# ===========================================================================
# 1. Initialisation
# ===========================================================================


class TestTokenCountingCallbackInit:
    """Verify default and custom initialisation."""

    def test_default_research_id_is_none(self):
        cb = TokenCountingCallback()
        assert cb.research_id is None

    def test_custom_research_id(self):
        cb = TokenCountingCallback(research_id="abc-123")
        assert cb.research_id == "abc-123"

    def test_default_research_context_is_empty_dict(self):
        cb = TokenCountingCallback()
        assert cb.research_context == {}

    def test_custom_research_context(self):
        ctx = {"research_query": "test", "research_mode": "deep"}
        cb = TokenCountingCallback(research_context=ctx)
        assert cb.research_context is ctx

    def test_counts_structure_has_required_keys(self):
        cb = TokenCountingCallback()
        assert cb.counts["total_tokens"] == 0
        assert cb.counts["total_prompt_tokens"] == 0
        assert cb.counts["total_completion_tokens"] == 0
        assert cb.counts["by_model"] == {}

    def test_start_time_initially_none(self):
        cb = TokenCountingCallback()
        assert cb.start_time is None

    def test_success_status_initially_success(self):
        cb = TokenCountingCallback()
        assert cb.success_status == "success"


# ===========================================================================
# 2. on_llm_start — model detection
# ===========================================================================


class TestOnLlmStartModelDetection:
    """Verify model name is extracted from the right source in priority order."""

    def test_model_from_invocation_params_model_name_key(self):
        """invocation_params.model_name should be used when model key absent."""
        cb = _make_callback()
        cb.on_llm_start(
            {},
            ["prompt"],
            invocation_params={"model_name": "gpt-4o-mini"},
        )
        assert cb.current_model == "gpt-4o-mini"

    def test_model_from_kwargs_model_name_key(self):
        """kwargs.model_name (direct) should be used as fallback."""
        cb = _make_callback()
        cb.on_llm_start({}, ["prompt"], model_name="claude-3-haiku")
        assert cb.current_model == "claude-3-haiku"

    def test_serialized_kwargs_model_name_key(self):
        """serialized['kwargs']['model_name'] should be used."""
        cb = _make_callback()
        cb.on_llm_start(
            {"kwargs": {"model_name": "gemma-7b"}},
            ["prompt"],
        )
        assert cb.current_model == "gemma-7b"

    def test_preset_model_overrides_all(self):
        """Preset model takes absolute priority."""
        cb = _make_callback()
        cb.preset_model = "preset-llm"
        cb.on_llm_start(
            {"kwargs": {"model": "ignored"}, "name": "also-ignored"},
            ["prompt"],
            invocation_params={"model": "still-ignored"},
        )
        assert cb.current_model == "preset-llm"

    def test_invocation_params_model_beats_serialized(self):
        """invocation_params.model wins over serialized.kwargs.model."""
        cb = _make_callback()
        cb.on_llm_start(
            {"kwargs": {"model": "serialized-model"}},
            ["prompt"],
            invocation_params={"model": "invocation-model"},
        )
        assert cb.current_model == "invocation-model"

    def test_ollama_without_kwargs_key_defaults_to_ollama(self):
        """ChatOllama type without kwargs dict at all defaults to 'ollama'."""
        cb = _make_callback()
        cb.on_llm_start({"_type": "ChatOllama"}, ["prompt"])
        assert cb.current_model == "ollama"

    def test_by_model_entry_created_on_first_call(self):
        """First call with a model creates the by_model entry."""
        cb = _make_callback()
        cb.on_llm_start({"kwargs": {"model": "new-model"}}, ["prompt"])
        entry = cb.counts["by_model"]["new-model"]
        assert entry["prompt_tokens"] == 0
        assert entry["completion_tokens"] == 0
        assert entry["total_tokens"] == 0
        assert entry["calls"] == 1

    def test_call_count_increments_same_model(self):
        """Repeated calls with the same model increment calls counter."""
        cb = _make_callback()
        for _ in range(4):
            cb.on_llm_start({"kwargs": {"model": "m1"}}, ["p"])
        assert cb.counts["by_model"]["m1"]["calls"] == 4

    def test_multiple_models_tracked_separately(self):
        """Different models get separate by_model entries."""
        cb = _make_callback()
        cb.on_llm_start({"kwargs": {"model": "model-a"}}, ["p"])
        cb.on_llm_start({"kwargs": {"model": "model-b"}}, ["p"])
        assert "model-a" in cb.counts["by_model"]
        assert "model-b" in cb.counts["by_model"]
        assert cb.counts["by_model"]["model-a"]["calls"] == 1
        assert cb.counts["by_model"]["model-b"]["calls"] == 1


# ===========================================================================
# 3. on_llm_start — provider detection
# ===========================================================================


class TestOnLlmStartProviderDetection:
    """Verify provider extracted from _type or kwargs."""

    def test_preset_provider_overrides(self):
        cb = _make_callback()
        cb.preset_provider = "custom-provider"
        cb.on_llm_start({"_type": "ChatOpenAI"}, ["p"])
        assert cb.current_provider == "custom-provider"

    def test_provider_kwarg_fallback(self):
        """Unknown _type falls back to provider kwarg."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "UnknownLLM"},
            ["p"],
            provider="my-custom",
        )
        assert cb.current_provider == "my-custom"

    def test_provider_stored_in_by_model(self):
        """Provider should be stored in the by_model entry."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatAnthropic", "kwargs": {"model": "claude"}},
            ["p"],
        )
        assert cb.counts["by_model"]["claude"]["provider"] == "anthropic"


# ===========================================================================
# 4. on_llm_start — prompt estimation
# ===========================================================================


class TestOnLlmStartPromptEstimation:
    """Verify prompt token estimation from prompt text."""

    def test_single_prompt_estimate(self):
        cb = _make_callback()
        cb.on_llm_start({}, ["a" * 100])  # 100 chars => 25 tokens
        assert cb.original_prompt_estimate == 25

    def test_empty_prompts_leaves_estimate_unchanged(self):
        """Empty prompts list skips estimation (guarded by `if prompts:`)."""
        cb = _make_callback()
        cb.original_prompt_estimate = 999
        cb.on_llm_start({}, [])
        # Empty prompts are skipped, so the estimate stays at its prior value
        assert cb.original_prompt_estimate == 999

    def test_start_time_set(self):
        cb = _make_callback()
        before = time.time()
        cb.on_llm_start({}, ["p"])
        assert cb.start_time >= before


# ===========================================================================
# 5. on_llm_end — token extraction and accumulation
# ===========================================================================


class TestOnLlmEnd:
    """Verify token extraction from various LLMResult shapes."""

    def test_total_tokens_calculated_when_missing(self):
        """When total_tokens is missing, prompt + completion is used."""
        cb = _make_callback()
        _setup_model(cb, "m1")

        result = _make_llm_result(
            llm_output={
                "token_usage": {
                    "prompt_tokens": 40,
                    "completion_tokens": 60,
                    # no total_tokens key
                }
            }
        )
        cb.on_llm_end(result)
        assert cb.counts["total_tokens"] == 100

    def test_by_model_updated_after_on_llm_end(self):
        """Per-model counts should be updated after on_llm_end."""
        cb = _make_callback()
        _setup_model(cb, "gpt-4")

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

        model_stats = cb.counts["by_model"]["gpt-4"]
        assert model_stats["prompt_tokens"] == 10
        assert model_stats["completion_tokens"] == 20
        assert model_stats["total_tokens"] == 30

    def test_response_time_calculated(self):
        """response_time_ms should be set when start_time is present."""
        cb = _make_callback()
        _setup_model(cb, "m1")
        cb.start_time = time.time() - 0.2  # 200ms ago

        result = _make_llm_result()
        cb.on_llm_end(result)

        assert cb.response_time_ms is not None
        assert cb.response_time_ms >= 150

    def test_response_time_none_without_start_time(self):
        """response_time_ms stays None if start_time was never set."""
        cb = _make_callback()
        cb.current_model = "m1"
        cb.counts["by_model"]["m1"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "unknown",
        }
        cb.start_time = None

        result = _make_llm_result()
        cb.on_llm_end(result)

        assert cb.response_time_ms is None

    def test_ollama_raw_metrics_captured(self):
        """Ollama response_metadata durations stored in ollama_metrics."""
        cb = _make_callback()
        _setup_model(cb, "llama3")

        msg = MagicMock()
        msg.usage_metadata = None
        msg.response_metadata = {
            "prompt_eval_count": 50,
            "eval_count": 30,
            "total_duration": 9000000,
            "load_duration": 1000000,
            "prompt_eval_duration": 4000000,
            "eval_duration": 4000000,
        }
        gen = MagicMock()
        gen.message = msg

        result = _make_llm_result(generations=[[gen]])
        cb.on_llm_end(result)

        assert cb.ollama_metrics["total_duration"] == 9000000
        assert cb.ollama_metrics["eval_count"] == 30

    def test_usage_metadata_none_skipped(self):
        """Generation with usage_metadata=None falls through to response_metadata."""
        cb = _make_callback()
        _setup_model(cb, "llama3")

        msg = MagicMock()
        msg.usage_metadata = None
        msg.response_metadata = {
            "prompt_eval_count": 15,
            "eval_count": 10,
        }
        gen = MagicMock()
        gen.message = msg

        result = _make_llm_result(generations=[[gen]])
        cb.on_llm_end(result)

        assert cb.counts["total_prompt_tokens"] == 15
        assert cb.counts["total_completion_tokens"] == 10

    def test_accumulation_across_two_different_models(self):
        """Tokens from different models accumulate in totals and separate by_model."""
        cb = _make_callback()

        # First model
        cb.on_llm_start({"kwargs": {"model": "model-a"}}, ["p"])
        result_a = _make_llm_result(
            llm_output={
                "token_usage": {
                    "prompt_tokens": 10,
                    "completion_tokens": 5,
                    "total_tokens": 15,
                }
            }
        )
        cb.on_llm_end(result_a)

        # Second model
        cb.on_llm_start({"kwargs": {"model": "model-b"}}, ["p"])
        result_b = _make_llm_result(
            llm_output={
                "token_usage": {
                    "prompt_tokens": 20,
                    "completion_tokens": 10,
                    "total_tokens": 30,
                }
            }
        )
        cb.on_llm_end(result_b)

        # Totals
        assert cb.counts["total_tokens"] == 45
        assert cb.counts["total_prompt_tokens"] == 30

        # Per-model
        assert cb.counts["by_model"]["model-a"]["total_tokens"] == 15
        assert cb.counts["by_model"]["model-b"]["total_tokens"] == 30

    def test_llm_output_usage_key_alternative(self):
        """'usage' key in llm_output (not 'token_usage') should work."""
        cb = _make_callback()
        _setup_model(cb, "m1")

        result = _make_llm_result(
            llm_output={
                "usage": {
                    "prompt_tokens": 7,
                    "completion_tokens": 3,
                    "total_tokens": 10,
                }
            }
        )
        cb.on_llm_end(result)
        assert cb.counts["total_tokens"] == 10


# ===========================================================================
# 6. on_llm_error
# ===========================================================================


class TestOnLlmError:
    """Verify error tracking behaviour."""

    def test_sets_success_status_to_error(self):
        cb = _make_callback()
        cb.on_llm_error(ValueError("bad input"))
        assert cb.success_status == "error"

    def test_captures_error_type_name(self):
        cb = _make_callback()
        cb.on_llm_error(RuntimeError("timeout"))
        assert cb.error_type == "RuntimeError"

    def test_response_time_calculated_on_error(self):
        cb = _make_callback()
        cb.start_time = time.time() - 0.3
        cb.on_llm_error(Exception("fail"))
        assert cb.response_time_ms is not None
        assert cb.response_time_ms >= 250

    def test_response_time_none_without_start_time(self):
        cb = _make_callback()
        cb.on_llm_error(Exception("fail"))
        assert cb.response_time_ms is None


# ===========================================================================
# 7. Context overflow detection
# ===========================================================================


class TestContextOverflow:
    """Verify context overflow detection and _get_context_overflow_fields."""

    def _make_cb_with_context_limit(self, limit=4096):
        """Create a callback with context_limit set via research_context."""
        cb = TokenCountingCallback(research_context={"context_limit": limit})
        _setup_model(cb, "llama3")
        return cb

    def test_context_truncated_detected_above_threshold(self):
        """prompt_eval_count >= context_limit * 0.80 sets context_truncated."""
        cb = self._make_cb_with_context_limit(4096)

        msg = MagicMock()
        msg.usage_metadata = None
        msg.response_metadata = {
            "prompt_eval_count": 3900,  # >= 4096 * 0.80 = 3276.8
            "eval_count": 50,
        }
        gen = MagicMock()
        gen.message = msg

        result = _make_llm_result(generations=[[gen]])
        cb.on_llm_end(result)

        assert cb.context_truncated is True

    def test_context_not_truncated_below_threshold(self):
        """prompt_eval_count below 80% does not set truncated."""
        cb = self._make_cb_with_context_limit(4096)

        msg = MagicMock()
        msg.usage_metadata = None
        msg.response_metadata = {
            "prompt_eval_count": 3000,  # < 4096 * 0.80 = 3276.8
            "eval_count": 50,
        }
        gen = MagicMock()
        gen.message = msg

        result = _make_llm_result(generations=[[gen]])
        cb.on_llm_end(result)

        assert cb.context_truncated is False

    def test_tokens_truncated_calculated(self):
        """tokens_truncated = original_prompt_estimate - prompt_eval_count."""
        cb = self._make_cb_with_context_limit(4096)
        cb.original_prompt_estimate = 5000

        msg = MagicMock()
        msg.usage_metadata = None
        msg.response_metadata = {
            "prompt_eval_count": 3900,
            "eval_count": 50,
        }
        gen = MagicMock()
        gen.message = msg

        result = _make_llm_result(generations=[[gen]])
        cb.on_llm_end(result)

        assert cb.tokens_truncated == 1100  # 5000 - 3900
        assert abs(cb.truncation_ratio - 0.22) < 0.01  # 1100/5000

    def test_get_context_overflow_fields_when_truncated(self):
        """_get_context_overflow_fields returns values when truncated."""
        cb = _make_callback()
        cb.context_limit = 4096
        cb.context_truncated = True
        cb.tokens_truncated = 500
        cb.truncation_ratio = 0.1

        fields = cb._get_context_overflow_fields()
        assert fields["context_limit"] == 4096
        assert fields["context_truncated"] is True
        assert fields["tokens_truncated"] == 500
        assert fields["truncation_ratio"] == 0.1

    def test_get_context_overflow_fields_when_not_truncated(self):
        """_get_context_overflow_fields returns None for truncation fields when not truncated."""
        cb = _make_callback()
        cb.context_limit = 4096
        cb.context_truncated = False

        fields = cb._get_context_overflow_fields()
        assert fields["context_truncated"] is False
        assert fields["tokens_truncated"] is None
        assert fields["truncation_ratio"] is None