/tests/metrics/test_token_counter_deep_coverage.py

"""Deep coverage tests for token_counter.py – targeting uncovered branches.

Focuses on:
- on_llm_start model/provider extraction paths
- on_llm_end token usage extraction (usage_metadata, response_metadata, llm_output)
- on_llm_error tracking
- _get_context_overflow_fields
- cost calculation helpers / tiktoken mocking
- TokenCounter.create_callback
- TokenCounter edge cases (no research_id, missing counts)
"""

import time
from unittest.mock import MagicMock, patch

from langchain_core.outputs import LLMResult

from local_deep_research.metrics.token_counter import (
    TokenCounter,
    TokenCountingCallback,
)

MODULE = "local_deep_research.metrics.token_counter"


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


def _make_llm_result(llm_output=None, generations=None):
    result = MagicMock(spec=LLMResult)
    result.llm_output = llm_output
    result.generations = generations or []
    return result


def _make_generation(usage_metadata=None, response_metadata=None):
    gen = MagicMock()
    msg = MagicMock()
    msg.usage_metadata = usage_metadata
    msg.response_metadata = response_metadata or {}
    gen.message = msg
    return gen


def _make_callback(**overrides):
    ctx = overrides.pop("research_context", {"research_query": "q"})
    cb = TokenCountingCallback(
        research_id=overrides.pop("research_id", "rid-1"),
        research_context=ctx,
    )
    for k, v in overrides.items():
        setattr(cb, k, v)
    return cb


# ---------------------------------------------------------------------------
# on_llm_start: model name extraction
# ---------------------------------------------------------------------------


class TestOnLlmStartModelExtraction:
    def test_preset_model_used_when_set(self):
        cb = _make_callback()
        cb.preset_model = "my-preset-model"
        cb.preset_provider = "openai"
        cb.on_llm_start({"_type": "ChatOpenAI"}, ["hello"])
        assert cb.current_model == "my-preset-model"
        assert cb.current_provider == "openai"

    def test_model_from_invocation_params(self):
        cb = _make_callback()
        cb.on_llm_start(
            {},
            ["hello"],
            invocation_params={"model": "gpt-4-turbo"},
        )
        assert cb.current_model == "gpt-4-turbo"

    def test_model_name_from_invocation_params(self):
        cb = _make_callback()
        cb.on_llm_start(
            {},
            ["hello"],
            invocation_params={"model_name": "claude-3"},
        )
        assert cb.current_model == "claude-3"

    def test_model_from_serialized_kwargs(self):
        cb = _make_callback()
        cb.on_llm_start(
            {"kwargs": {"model": "gemma3:12b"}},
            ["hello"],
        )
        assert cb.current_model == "gemma3:12b"

    def test_model_name_from_serialized_kwargs(self):
        cb = _make_callback()
        cb.on_llm_start(
            {"kwargs": {"model_name": "llama3"}},
            ["hello"],
        )
        assert cb.current_model == "llama3"

    def test_model_from_serialized_name(self):
        cb = _make_callback()
        cb.on_llm_start({"name": "SerializedModelName"}, ["hello"])
        assert cb.current_model == "SerializedModelName"

    def test_ollama_fallback_to_kwargs_model(self):
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatOllama", "kwargs": {"model": "mistral"}},
            ["hello"],
        )
        assert cb.current_model == "mistral"
        assert cb.current_provider == "ollama"

    def test_ollama_fallback_to_type_string(self):
        """When Ollama _type present but no 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_unknown_model_from_type(self):
        cb = _make_callback()
        cb.on_llm_start({"_type": "ChatSomething"}, ["hello"])
        assert cb.current_model == "ChatSomething"

    def test_unknown_model_fallback(self):
        cb = _make_callback()
        cb.on_llm_start({}, ["hello"])
        assert cb.current_model == "unknown"

    def test_provider_ollama_from_type(self):
        cb = _make_callback()
        cb.on_llm_start({"_type": "ChatOllama"}, ["hello"])
        assert cb.current_provider == "ollama"

    def test_provider_openai_from_type(self):
        cb = _make_callback()
        cb.on_llm_start({"_type": "ChatOpenAI"}, ["hello"])
        assert cb.current_provider == "openai"

    def test_provider_anthropic_from_type(self):
        cb = _make_callback()
        cb.on_llm_start({"_type": "ChatAnthropic"}, ["hello"])
        assert cb.current_provider == "anthropic"

    def test_provider_unknown_when_no_type(self):
        cb = _make_callback()
        cb.on_llm_start({}, ["hello"])
        assert cb.current_provider == "unknown"

    def test_call_count_incremented(self):
        cb = _make_callback()
        cb.on_llm_start({"_type": "ChatOpenAI"}, ["hello"])
        cb.on_llm_start({"_type": "ChatOpenAI"}, ["hello again"])
        model = cb.current_model
        assert cb.counts["by_model"][model]["calls"] == 2

    def test_start_time_recorded(self):
        cb = _make_callback()
        cb.on_llm_start({}, ["hello"])
        assert cb.start_time is not None
        assert cb.start_time <= time.time()

    def test_prompt_estimate_computed(self):
        cb = _make_callback()
        cb.on_llm_start({}, ["a" * 400])
        assert cb.original_prompt_estimate == 100  # 400 // 4


# ---------------------------------------------------------------------------
# on_llm_end: token usage paths
# ---------------------------------------------------------------------------


class TestOnLlmEndTokenUsage:
    def _run_end(self, cb, llm_result):
        with patch.object(cb, "_save_to_db") as mock_save:
            cb.on_llm_end(llm_result)
            return mock_save

    def test_token_usage_from_llm_output(self):
        cb = _make_callback()
        cb.current_model = "gpt-4"
        cb.counts["by_model"]["gpt-4"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "openai",
        }
        result = _make_llm_result(
            llm_output={
                "token_usage": {
                    "prompt_tokens": 50,
                    "completion_tokens": 30,
                    "total_tokens": 80,
                }
            }
        )
        save_mock = self._run_end(cb, result)
        assert cb.counts["total_tokens"] == 80
        save_mock.assert_called_once_with(50, 30)

    def test_token_usage_from_usage_metadata(self):
        cb = _make_callback()
        cb.current_model = "claude-3"
        cb.counts["by_model"]["claude-3"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "anthropic",
        }
        usage_meta = {
            "input_tokens": 20,
            "output_tokens": 10,
            "total_tokens": 30,
        }
        gen = _make_generation(usage_metadata=usage_meta)
        result = _make_llm_result(generations=[[gen]])
        save_mock = self._run_end(cb, result)
        assert cb.counts["total_tokens"] == 30
        save_mock.assert_called_once_with(20, 10)

    def test_token_usage_from_response_metadata_ollama(self):
        cb = _make_callback()
        cb.current_model = "mistral"
        cb.counts["by_model"]["mistral"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "ollama",
        }
        resp_meta = {"prompt_eval_count": 40, "eval_count": 20}
        gen = _make_generation(usage_metadata=None, response_metadata=resp_meta)
        result = _make_llm_result(generations=[[gen]])
        save_mock = self._run_end(cb, result)
        assert cb.counts["total_tokens"] == 60
        save_mock.assert_called_once_with(40, 20)

    def test_no_token_usage_saves_zero_counts(self):
        """No usage data from the provider still records the call (#4457)."""
        cb = _make_callback()
        result = _make_llm_result(llm_output=None, generations=[])
        save_mock = self._run_end(cb, result)
        save_mock.assert_called_once_with(0, 0)

    def test_response_time_calculated(self):
        cb = _make_callback()
        cb.start_time = time.time() - 0.5  # 500ms ago
        result = _make_llm_result(llm_output=None, generations=[])
        cb.on_llm_end(result)
        assert cb.response_time_ms is not None
        assert cb.response_time_ms >= 400  # at least 400ms

    def test_no_save_when_no_research_id(self):
        cb = TokenCountingCallback(research_id=None)
        cb.current_model = "gpt-4"
        cb.counts["by_model"]["gpt-4"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "openai",
        }
        result = _make_llm_result(
            llm_output={
                "token_usage": {"prompt_tokens": 10, "completion_tokens": 5}
            }
        )
        with patch.object(cb, "_save_to_db") as mock_save:
            cb.on_llm_end(result)
            mock_save.assert_not_called()

    def test_context_overflow_detection(self):
        """When prompt_eval_count >= 95% of context_limit, context_truncated is set."""
        cb = _make_callback()
        cb.context_limit = 1000
        cb.original_prompt_estimate = 1200  # More than actual => truncated
        cb.current_model = "llama"
        cb.counts["by_model"]["llama"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "ollama",
        }
        resp_meta = {
            "prompt_eval_count": 960,
            "eval_count": 40,
        }  # 960 >= 950 (95%)
        gen = _make_generation(usage_metadata=None, response_metadata=resp_meta)
        result = _make_llm_result(generations=[[gen]])
        with patch.object(cb, "_save_to_db"):
            cb.on_llm_end(result)
        assert cb.context_truncated is True
        assert cb.tokens_truncated > 0


# ---------------------------------------------------------------------------
# on_llm_error
# ---------------------------------------------------------------------------


class TestOnLlmError:
    def test_sets_error_status(self):
        cb = _make_callback()
        with patch.object(cb, "_save_to_db") as mock_save:
            cb.on_llm_error(ValueError("bad input"))
        assert cb.success_status == "error"
        assert cb.error_type == "ValueError"
        mock_save.assert_called_once_with(0, 0)

    def test_calculates_response_time_on_error(self):
        cb = _make_callback()
        cb.start_time = time.time() - 1.0
        with patch.object(cb, "_save_to_db"):
            cb.on_llm_error(RuntimeError("crash"))
        assert cb.response_time_ms is not None
        assert cb.response_time_ms >= 900

    def test_no_save_when_no_research_id(self):
        cb = TokenCountingCallback(research_id=None)
        with patch.object(cb, "_save_to_db") as mock_save:
            cb.on_llm_error(RuntimeError("boom"))
        mock_save.assert_not_called()


# ---------------------------------------------------------------------------
# _get_context_overflow_fields
# ---------------------------------------------------------------------------


class TestGetContextOverflowFields:
    def test_fields_when_not_truncated(self):
        cb = _make_callback()
        cb.context_limit = 4096
        cb.context_truncated = False
        cb.tokens_truncated = 0
        cb.truncation_ratio = 0.0
        fields = cb._get_context_overflow_fields()
        assert fields["context_limit"] == 4096
        assert fields["context_truncated"] is False
        assert fields["tokens_truncated"] is None
        assert fields["truncation_ratio"] is None

    def test_fields_when_truncated(self):
        cb = _make_callback()
        cb.context_limit = 2048
        cb.context_truncated = True
        cb.tokens_truncated = 300
        cb.truncation_ratio = 0.25
        fields = cb._get_context_overflow_fields()
        assert fields["tokens_truncated"] == 300
        assert fields["truncation_ratio"] == 0.25

    def test_ollama_metrics_included(self):
        cb = _make_callback()
        cb.ollama_metrics = {
            "prompt_eval_count": 500,
            "eval_count": 100,
            "total_duration": 999,
        }
        fields = cb._get_context_overflow_fields()
        assert fields["ollama_prompt_eval_count"] == 500
        assert fields["ollama_eval_count"] == 100

    def test_missing_ollama_metrics_returns_none(self):
        cb = _make_callback()
        cb.ollama_metrics = {}
        fields = cb._get_context_overflow_fields()
        assert fields["ollama_prompt_eval_count"] is None
        assert fields["ollama_total_duration"] is None


# ---------------------------------------------------------------------------
# TokenCounter.create_callback
# ---------------------------------------------------------------------------


class TestTokenCounterCreateCallback:
    def test_create_callback_returns_callback_instance(self):
        counter = TokenCounter()
        cb = counter.create_callback("res-99")
        assert isinstance(cb, TokenCountingCallback)
        assert cb.research_id == "res-99"

    def test_create_callback_with_context(self):
        counter = TokenCounter()
        ctx = {"research_query": "AI safety", "username": "bob"}
        cb = counter.create_callback("res-42", ctx)
        assert cb.research_context == ctx

    def test_create_callback_no_context(self):
        counter = TokenCounter()
        cb = counter.create_callback("res-1", None)
        assert cb.research_context == {}

    def test_multiple_callbacks_independent(self):
        counter = TokenCounter()
        cb1 = counter.create_callback("res-1")
        cb2 = counter.create_callback("res-2")
        cb1.current_model = "gpt-4"
        assert cb2.current_model is None


# ---------------------------------------------------------------------------
# Tiktoken mocking – cost calculation helpers
# ---------------------------------------------------------------------------


class TestTiktokenMocking:
    """Test that token counting works when tiktoken is mocked."""

    def test_on_llm_start_no_tiktoken_needed(self):
        """on_llm_start should work without tiktoken."""
        cb = _make_callback()
        # No patch needed – tiktoken is not used in on_llm_start
        cb.on_llm_start({"_type": "ChatOpenAI"}, ["Hello world"])
        assert cb.current_provider == "openai"

    def test_token_count_aggregates_across_calls(self):
        cb = _make_callback()
        cb.on_llm_start({"_type": "ChatOpenAI"}, ["prompt"])
        cb.current_model = "gpt-4"
        cb.counts["by_model"].setdefault(
            "gpt-4",
            {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
                "calls": 1,
                "provider": "openai",
            },
        )

        result1 = _make_llm_result(
            llm_output={
                "token_usage": {
                    "prompt_tokens": 100,
                    "completion_tokens": 50,
                    "total_tokens": 150,
                }
            }
        )
        result2 = _make_llm_result(
            llm_output={
                "token_usage": {
                    "prompt_tokens": 200,
                    "completion_tokens": 100,
                    "total_tokens": 300,
                }
            }
        )
        with patch.object(cb, "_save_to_db"):
            cb.on_llm_end(result1)
            cb.on_llm_end(result2)

        assert cb.counts["total_tokens"] == 450
        assert cb.counts["total_prompt_tokens"] == 300
        assert cb.counts["total_completion_tokens"] == 150


# ---------------------------------------------------------------------------
# TokenCounter – cost metrics and empty states
# ---------------------------------------------------------------------------


class TestTokenCounterMetrics:
    def test_initial_counts_are_zero(self):
        counter = TokenCounter()
        cb = counter.create_callback("res-1")
        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_llm_output_usage_key_fallback(self):
        """When token_usage absent, falls back to 'usage' key in llm_output."""
        cb = _make_callback()
        cb.current_model = "gpt-3.5"
        cb.counts["by_model"]["gpt-3.5"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "openai",
        }
        result = _make_llm_result(
            llm_output={
                "usage": {
                    "prompt_tokens": 10,
                    "completion_tokens": 5,
                    "total_tokens": 15,
                }
            }
        )
        with patch.object(cb, "_save_to_db"):
            cb.on_llm_end(result)
        assert cb.counts["total_tokens"] == 15