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

"""Tests for TokenCountingCallback real behavior - covering token extraction,
model detection, provider detection, context overflow, and error handling."""

import time
from unittest.mock import Mock, patch

from langchain_core.outputs import LLMResult

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


class TestTokenCountingCallbackModelDetection:
    """Tests for model name detection from various sources."""

    def test_preset_model_takes_priority(self):
        """Preset model name should override all other detection."""
        callback = TokenCountingCallback()
        callback.preset_model = "my-preset-model"
        callback.preset_provider = "openai"

        serialized = {"name": "different-model", "_type": "ChatOpenAI"}
        callback.on_llm_start(serialized, ["test prompt"])

        assert callback.current_model == "my-preset-model"

    def test_model_from_invocation_params(self):
        """Model should be extracted from invocation_params."""
        callback = TokenCountingCallback()

        serialized = {}
        kwargs = {"invocation_params": {"model": "gpt-4-turbo"}}
        callback.on_llm_start(serialized, ["test"], **kwargs)

        assert callback.current_model == "gpt-4-turbo"

    def test_model_from_serialized_kwargs(self):
        """Model should be extracted from serialized kwargs."""
        callback = TokenCountingCallback()

        serialized = {"kwargs": {"model": "claude-3-opus"}}
        callback.on_llm_start(serialized, ["test"])

        assert callback.current_model == "claude-3-opus"

    def test_model_from_serialized_name(self):
        """Model should be extracted from serialized name."""
        callback = TokenCountingCallback()

        serialized = {"name": "my-model"}
        callback.on_llm_start(serialized, ["test"])

        assert callback.current_model == "my-model"

    def test_ollama_model_extraction(self):
        """Ollama model should be extracted from serialized type and kwargs."""
        callback = TokenCountingCallback()

        serialized = {
            "_type": "ChatOllama",
            "kwargs": {"model": "llama3:8b"},
        }
        callback.on_llm_start(serialized, ["test"])

        assert callback.current_model == "llama3:8b"

    def test_ollama_type_without_model_name(self):
        """Ollama type without model name should default to 'ollama'."""
        callback = TokenCountingCallback()

        serialized = {"_type": "ChatOllama"}
        callback.on_llm_start(serialized, ["test"])

        assert callback.current_model == "ollama"

    def test_unknown_model_fallback(self):
        """Unknown model should fall back to 'unknown'."""
        callback = TokenCountingCallback()

        serialized = {}
        callback.on_llm_start(serialized, ["test"])

        assert callback.current_model == "unknown"

    def test_type_as_model_fallback(self):
        """_type should be used as model name if no other source found."""
        callback = TokenCountingCallback()

        serialized = {"_type": "CustomLLM"}
        callback.on_llm_start(serialized, ["test"])

        assert callback.current_model == "CustomLLM"


class TestTokenCountingCallbackProviderDetection:
    """Tests for provider detection from various sources."""

    def test_preset_provider_takes_priority(self):
        """Preset provider should override detection."""
        callback = TokenCountingCallback()
        callback.preset_provider = "anthropic"

        serialized = {"_type": "ChatOpenAI"}
        callback.on_llm_start(serialized, ["test"])

        assert callback.current_provider == "anthropic"

    def test_openai_provider_detection(self):
        """ChatOpenAI type should detect 'openai' provider."""
        callback = TokenCountingCallback()

        serialized = {"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}}
        callback.on_llm_start(serialized, ["test"])

        assert callback.current_provider == "openai"

    def test_anthropic_provider_detection(self):
        """ChatAnthropic type should detect 'anthropic' provider."""
        callback = TokenCountingCallback()

        serialized = {
            "_type": "ChatAnthropic",
            "kwargs": {"model": "claude-3-opus"},
        }
        callback.on_llm_start(serialized, ["test"])

        assert callback.current_provider == "anthropic"

    def test_ollama_provider_detection(self):
        """ChatOllama type should detect 'ollama' provider."""
        callback = TokenCountingCallback()

        serialized = {
            "_type": "ChatOllama",
            "kwargs": {"model": "llama3"},
        }
        callback.on_llm_start(serialized, ["test"])

        assert callback.current_provider == "ollama"

    def test_unknown_provider_fallback(self):
        """Unknown type should fall back to 'unknown' provider."""
        callback = TokenCountingCallback()

        serialized = {}
        callback.on_llm_start(serialized, ["test"])

        assert callback.current_provider == "unknown"


class TestTokenCountingCallbackTokenExtraction:
    """Tests for token extraction from LLM responses."""

    def test_token_extraction_from_llm_output(self):
        """Tokens should be extracted from response.llm_output."""
        callback = TokenCountingCallback()
        callback.current_model = "gpt-4"
        callback.counts["by_model"]["gpt-4"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "openai",
        }

        response = Mock(spec=LLMResult)
        response.llm_output = {
            "token_usage": {
                "prompt_tokens": 100,
                "completion_tokens": 50,
                "total_tokens": 150,
            }
        }
        response.generations = []

        callback.on_llm_end(response)

        assert callback.counts["total_prompt_tokens"] == 100
        assert callback.counts["total_completion_tokens"] == 50
        assert callback.counts["total_tokens"] == 150

    def test_token_extraction_from_usage_metadata(self):
        """Tokens should be extracted from generation.message.usage_metadata."""
        callback = TokenCountingCallback()
        callback.current_model = "llama3"
        callback.counts["by_model"]["llama3"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "ollama",
        }

        # Create mock generation with usage_metadata
        mock_message = Mock()
        mock_message.usage_metadata = {
            "input_tokens": 200,
            "output_tokens": 80,
            "total_tokens": 280,
        }
        mock_message.response_metadata = {}

        mock_generation = Mock()
        mock_generation.message = mock_message

        response = Mock(spec=LLMResult)
        response.llm_output = None
        response.generations = [[mock_generation]]

        callback.on_llm_end(response)

        assert callback.counts["total_prompt_tokens"] == 200
        assert callback.counts["total_completion_tokens"] == 80
        assert callback.counts["total_tokens"] == 280

    def test_token_extraction_from_response_metadata_ollama(self):
        """Tokens should be extracted from Ollama response_metadata."""
        callback = TokenCountingCallback()
        callback.current_model = "llama3"
        callback.counts["by_model"]["llama3"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "ollama",
        }

        mock_message = Mock()
        mock_message.usage_metadata = None
        mock_message.response_metadata = {
            "prompt_eval_count": 300,
            "eval_count": 120,
            "total_duration": 5000000000,
            "load_duration": 100000000,
            "prompt_eval_duration": 2000000000,
            "eval_duration": 2900000000,
        }

        mock_generation = Mock()
        mock_generation.message = mock_message

        response = Mock(spec=LLMResult)
        response.llm_output = None
        response.generations = [[mock_generation]]

        callback.on_llm_end(response)

        assert callback.counts["total_prompt_tokens"] == 300
        assert callback.counts["total_completion_tokens"] == 120
        assert callback.counts["total_tokens"] == 420

    def test_no_token_usage_available(self):
        """No token usage data should not crash and not update counts."""
        callback = TokenCountingCallback()
        callback.current_model = "test"
        callback.counts["by_model"]["test"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "unknown",
        }

        response = Mock(spec=LLMResult)
        response.llm_output = None
        response.generations = []

        callback.on_llm_end(response)

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

    def test_cumulative_token_counting(self):
        """Multiple on_llm_end calls should accumulate token counts."""
        callback = TokenCountingCallback()
        callback.current_model = "gpt-4"
        callback.counts["by_model"]["gpt-4"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "openai",
        }

        for i in range(3):
            response = Mock(spec=LLMResult)
            response.llm_output = {
                "token_usage": {
                    "prompt_tokens": 10,
                    "completion_tokens": 5,
                    "total_tokens": 15,
                }
            }
            response.generations = []
            callback.on_llm_end(response)

        assert callback.counts["total_tokens"] == 45
        assert callback.counts["total_prompt_tokens"] == 30
        assert callback.counts["total_completion_tokens"] == 15


class TestTokenCountingCallbackContextOverflow:
    """Tests for context overflow detection."""

    def test_context_overflow_detected(self):
        """Context overflow should be detected when near context limit."""
        callback = TokenCountingCallback(
            research_context={"context_limit": 4096}
        )
        callback.current_model = "llama3"
        callback.counts["by_model"]["llama3"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "ollama",
        }

        # Estimate original prompt is much larger
        callback.original_prompt_estimate = 5000

        # Simulate start to set context limit
        callback.on_llm_start({"_type": "ChatOllama"}, ["x" * 20000])

        mock_message = Mock()
        mock_message.usage_metadata = None
        mock_message.response_metadata = {
            "prompt_eval_count": 3900,  # 95% of 4096
            "eval_count": 100,
        }
        mock_generation = Mock()
        mock_generation.message = mock_message

        response = Mock(spec=LLMResult)
        response.llm_output = None
        response.generations = [[mock_generation]]

        callback.on_llm_end(response)

        assert callback.context_truncated is True
        assert callback.tokens_truncated > 0

    def test_no_overflow_when_below_threshold(self):
        """No overflow when prompt tokens are below 95% of context limit."""
        callback = TokenCountingCallback(
            research_context={"context_limit": 4096}
        )
        callback.current_model = "llama3"
        callback.counts["by_model"]["llama3"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "ollama",
        }

        callback.on_llm_start({"_type": "ChatOllama"}, ["x" * 4000])

        mock_message = Mock()
        mock_message.usage_metadata = None
        mock_message.response_metadata = {
            "prompt_eval_count": 2000,  # Well below 95% of 4096
            "eval_count": 100,
        }
        mock_generation = Mock()
        mock_generation.message = mock_message

        response = Mock(spec=LLMResult)
        response.llm_output = None
        response.generations = [[mock_generation]]

        callback.on_llm_end(response)

        assert callback.context_truncated is False

    def test_no_overflow_without_context_limit(self):
        """No overflow detection when context_limit is not set."""
        callback = TokenCountingCallback()
        callback.current_model = "gpt-4"
        callback.counts["by_model"]["gpt-4"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "openai",
        }

        mock_message = Mock()
        mock_message.usage_metadata = None
        mock_message.response_metadata = {
            "prompt_eval_count": 100000,
            "eval_count": 100,
        }
        mock_generation = Mock()
        mock_generation.message = mock_message

        response = Mock(spec=LLMResult)
        response.llm_output = None
        response.generations = [[mock_generation]]

        callback.on_llm_end(response)

        assert callback.context_truncated is False


class TestTokenCountingCallbackMissingUsageData:
    """Tests for recording calls when the provider reports no usage (#4457)."""

    @staticmethod
    def _response_without_usage():
        response = Mock(spec=LLMResult)
        response.llm_output = None
        response.generations = []
        return response

    def test_no_usage_with_research_id_still_saves_to_db(self):
        """A call without usage data must still be recorded (zero counts)."""
        callback = TokenCountingCallback(research_id="research-123")
        callback.current_model = "test-model"
        callback._save_to_db = Mock()

        callback.on_llm_end(self._response_without_usage())

        callback._save_to_db.assert_called_once_with(0, 0)

    def test_no_usage_without_research_id_does_not_save(self):
        """Without a research_id there is nothing to persist."""
        callback = TokenCountingCallback()
        callback.current_model = "test-model"
        callback._save_to_db = Mock()

        callback.on_llm_end(self._response_without_usage())

        callback._save_to_db.assert_not_called()

    def test_no_usage_does_not_update_counts(self):
        """Zero-count recording must not inflate in-memory token counts."""
        callback = TokenCountingCallback(research_id="research-123")
        callback.current_model = "test-model"
        callback._save_to_db = Mock()

        callback.on_llm_end(self._response_without_usage())

        assert callback.counts["total_tokens"] == 0
        assert callback.counts["total_prompt_tokens"] == 0
        assert callback.counts["total_completion_tokens"] == 0

    def test_repeated_no_usage_records_every_call_warns_once(self):
        """Every no-usage call is recorded, but the warning fires once."""
        import local_deep_research.metrics.token_counter as tc_mod

        callback = TokenCountingCallback(research_id="research-123")
        callback.current_model = "test-model"
        callback._save_to_db = Mock()

        with patch.object(tc_mod.logger, "warning") as mock_warn:
            for _ in range(5):
                callback.on_llm_end(self._response_without_usage())

        assert callback._save_to_db.call_count == 5
        mock_warn.assert_called_once()

    def test_usage_present_saves_actual_counts(self):
        """Sanity check: usage data still saves the real counts."""
        callback = TokenCountingCallback(research_id="research-123")
        callback.current_model = "gpt-4"
        callback.counts["by_model"]["gpt-4"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "openai",
        }
        callback._save_to_db = Mock()

        response = Mock(spec=LLMResult)
        response.llm_output = {
            "token_usage": {
                "prompt_tokens": 10,
                "completion_tokens": 5,
                "total_tokens": 15,
            }
        }
        response.generations = []
        callback.on_llm_end(response)

        callback._save_to_db.assert_called_once_with(10, 5)


class TestTokenCountingCallbackErrorHandling:
    """Tests for error tracking in callback."""

    def test_error_sets_status_and_type(self):
        """on_llm_error should set error status and type."""
        callback = TokenCountingCallback()
        callback.start_time = time.time()

        error = ValueError("test error")
        callback.on_llm_error(error)

        assert callback.success_status == "error"
        assert callback.error_type == "ValueError"

    def test_error_calculates_response_time(self):
        """on_llm_error should calculate response time."""
        callback = TokenCountingCallback()
        callback.start_time = time.time() - 0.5  # 500ms ago

        callback.on_llm_error(RuntimeError("test"))

        assert callback.response_time_ms is not None
        assert callback.response_time_ms >= 400  # At least 400ms

    def test_error_without_start_time(self):
        """on_llm_error without start_time should not crash."""
        callback = TokenCountingCallback()

        callback.on_llm_error(RuntimeError("test"))

        assert callback.success_status == "error"
        assert callback.response_time_ms is None


class TestTokenCountingCallbackTimingAndCallStack:
    """Tests for timing and call stack tracking."""

    def test_start_time_set_on_llm_start(self):
        """on_llm_start should record start time."""
        callback = TokenCountingCallback()

        before = time.time()
        callback.on_llm_start({}, ["test"])
        after = time.time()

        assert before <= callback.start_time <= after

    def test_response_time_calculated_on_end(self):
        """on_llm_end should calculate response time in ms."""
        callback = TokenCountingCallback()
        callback.current_model = "test"
        callback.counts["by_model"]["test"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
            "provider": "unknown",
        }
        callback.start_time = time.time() - 0.1  # 100ms ago

        response = Mock(spec=LLMResult)
        response.llm_output = None
        response.generations = []

        callback.on_llm_end(response)

        assert callback.response_time_ms is not None
        assert callback.response_time_ms >= 50  # At least 50ms

    def test_prompt_estimate_from_prompts(self):
        """original_prompt_estimate should be set from prompt length."""
        callback = TokenCountingCallback()

        callback.on_llm_start({}, ["Hello world!"])  # 12 chars ~ 3 tokens

        assert callback.original_prompt_estimate == 3  # 12 // 4

    def test_prompt_estimate_multiple_prompts(self):
        """Multiple prompts should sum their character counts for estimate."""
        callback = TokenCountingCallback()

        callback.on_llm_start({}, ["aaaa", "bbbb", "cccc"])  # 12 chars total

        assert callback.original_prompt_estimate == 3  # 12 // 4

    def test_call_count_incremented(self):
        """Call count should increment on each on_llm_start."""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            {"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}}, ["test"]
        )
        callback.on_llm_start(
            {"_type": "ChatOpenAI", "kwargs": {"model": "gpt-4"}}, ["test"]
        )

        assert callback.counts["by_model"]["gpt-4"]["calls"] == 2


class TestTokenCountingCallbackGetContextOverflowFields:
    """Tests for _get_context_overflow_fields helper."""

    def test_fields_when_no_overflow(self):
        """Fields should indicate no overflow when not truncated."""
        callback = TokenCountingCallback()

        fields = callback._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(self):
        """Fields should contain overflow data when truncated."""
        callback = TokenCountingCallback()
        callback.context_limit = 4096
        callback.context_truncated = True
        callback.tokens_truncated = 500
        callback.truncation_ratio = 0.12
        callback.ollama_metrics = {
            "prompt_eval_count": 3900,
            "eval_count": 100,
        }

        fields = callback._get_context_overflow_fields()

        assert fields["context_truncated"] is True
        assert fields["tokens_truncated"] == 500
        assert fields["truncation_ratio"] == 0.12
        assert fields["context_limit"] == 4096
        assert fields["ollama_prompt_eval_count"] == 3900


class TestTokenCounterManager:
    """Tests for TokenCounter manager class."""

    def test_create_callback_returns_callback(self):
        """create_callback should return a TokenCountingCallback instance."""
        counter = TokenCounter()
        callback = counter.create_callback(
            research_id="test-123",
            research_context={"key": "value"},
        )

        assert isinstance(callback, TokenCountingCallback)
        assert callback.research_id == "test-123"
        assert callback.research_context == {"key": "value"}

    def test_create_callback_without_args(self):
        """create_callback without args should work."""
        counter = TokenCounter()
        callback = counter.create_callback()

        assert isinstance(callback, TokenCountingCallback)
        assert callback.research_id is None

    def test_empty_metrics_structure(self):
        """_get_empty_metrics should return proper structure."""
        counter = TokenCounter()
        metrics = counter._get_empty_metrics()

        assert metrics["total_tokens"] == 0
        assert metrics["total_researches"] == 0
        assert metrics["by_model"] == []
        assert metrics["recent_researches"] == []
        assert "token_breakdown" in metrics