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

"""Tests for metrics token_counter module."""

import time
from unittest.mock import MagicMock, Mock, patch


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


class TestTokenCountingCallbackInit:
    """Tests for TokenCountingCallback initialization."""

    def test_initializes_with_no_args(self):
        """Should initialize without arguments."""
        callback = TokenCountingCallback()
        assert callback.research_id is None
        assert callback.research_context == {}

    def test_initializes_with_research_id(self):
        """Should store research_id."""
        callback = TokenCountingCallback(research_id="test-uuid")
        assert callback.research_id == "test-uuid"

    def test_initializes_with_research_context(self):
        """Should store research context."""
        context = {"query": "test", "mode": "quick"}
        callback = TokenCountingCallback(research_context=context)
        assert callback.research_context == context

    def test_initializes_counts_structure(self):
        """Should initialize counts with correct structure."""
        callback = TokenCountingCallback()

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

    def test_initializes_timing_fields(self):
        """Should initialize timing fields."""
        callback = TokenCountingCallback()

        assert callback.start_time is None
        assert callback.response_time_ms is None
        assert callback.success_status == "success"

    def test_initializes_context_overflow_fields(self):
        """Should initialize context overflow tracking fields."""
        callback = TokenCountingCallback()

        assert callback.context_limit is None
        assert callback.context_truncated is False
        assert callback.tokens_truncated == 0


class TestTokenCountingCallbackOnLlmStart:
    """Tests for on_llm_start method."""

    def test_captures_start_time(self):
        """Should capture start time."""
        callback = TokenCountingCallback()

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

        assert callback.start_time is not None
        assert callback.start_time <= time.time()

    def test_estimates_prompt_tokens(self):
        """Should estimate prompt tokens from text length."""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            serialized={"_type": "ChatOpenAI"},
            prompts=["A" * 400],  # 400 chars ~= 100 tokens
        )

        assert callback.original_prompt_estimate == 100

    def test_extracts_model_from_invocation_params(self):
        """Should extract model from invocation_params."""
        callback = TokenCountingCallback()

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

        assert callback.current_model == "gpt-4"

    def test_extracts_model_from_kwargs(self):
        """Should extract model from kwargs."""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            serialized={},
            prompts=["test"],
            model="gpt-3.5-turbo",
        )

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

    def test_extracts_model_from_serialized(self):
        """Should extract model from serialized kwargs."""
        callback = TokenCountingCallback()

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

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

    def test_uses_preset_model(self):
        """Should use preset_model if set."""
        callback = TokenCountingCallback()
        callback.preset_model = "preset-model"

        callback.on_llm_start(
            serialized={"kwargs": {"model": "other-model"}},
            prompts=["test"],
        )

        assert callback.current_model == "preset-model"

    def test_extracts_provider_from_type(self):
        """Should extract provider from serialized type."""
        callback = TokenCountingCallback()

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

        assert callback.current_provider == "ollama"

    def test_extracts_openai_provider(self):
        """Should detect OpenAI provider."""
        callback = TokenCountingCallback()

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

        assert callback.current_provider == "openai"

    def test_extracts_anthropic_provider(self):
        """Should detect Anthropic provider."""
        callback = TokenCountingCallback()

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

        assert callback.current_provider == "anthropic"

    def test_uses_preset_provider(self):
        """Should use preset_provider if set."""
        callback = TokenCountingCallback()
        callback.preset_provider = "custom-provider"

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

        assert callback.current_provider == "custom-provider"

    def test_increments_call_count(self):
        """Should increment call count for model."""
        callback = TokenCountingCallback()

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

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

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

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

    def test_initializes_model_tracking(self):
        """Should initialize tracking for new models."""
        callback = TokenCountingCallback()

        callback.on_llm_start(
            serialized={"kwargs": {"model": "new-model"}},
            prompts=["test"],
        )

        assert "new-model" in callback.counts["by_model"]
        assert callback.counts["by_model"]["new-model"]["prompt_tokens"] == 0
        assert (
            callback.counts["by_model"]["new-model"]["completion_tokens"] == 0
        )


class TestTokenCountingCallbackOnLlmEnd:
    """Tests for on_llm_end method."""

    def test_calculates_response_time(self):
        """Should calculate response time in milliseconds."""
        callback = TokenCountingCallback()
        callback.start_time = time.time() - 0.5  # 500ms ago

        mock_response = Mock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 10,
                "completion_tokens": 5,
                "total_tokens": 15,
            }
        }
        mock_response.generations = []

        callback.on_llm_end(mock_response)

        assert callback.response_time_ms is not None
        assert callback.response_time_ms >= 500

    def test_extracts_tokens_from_llm_output(self):
        """Should extract token counts from 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,
        }

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

        callback.on_llm_end(mock_response)

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

    def test_extracts_tokens_from_usage_metadata(self):
        """Should extract tokens from usage_metadata (Ollama format)."""
        callback = TokenCountingCallback()
        callback.current_model = "llama"
        callback.counts["by_model"]["llama"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
        }

        mock_message = Mock()
        mock_message.usage_metadata = {
            "input_tokens": 80,
            "output_tokens": 40,
            "total_tokens": 120,
        }
        mock_message.response_metadata = {}

        mock_generation = Mock()
        mock_generation.message = mock_message

        mock_response = Mock()
        mock_response.llm_output = None
        mock_response.generations = [[mock_generation]]

        callback.on_llm_end(mock_response)

        assert callback.counts["total_prompt_tokens"] == 80
        assert callback.counts["total_completion_tokens"] == 40

    def test_extracts_tokens_from_response_metadata_ollama(self):
        """Should extract tokens from Ollama response_metadata."""
        callback = TokenCountingCallback()
        callback.current_model = "mistral"
        callback.counts["by_model"]["mistral"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
        }

        mock_message = Mock()
        mock_message.usage_metadata = None
        mock_message.response_metadata = {
            "prompt_eval_count": 60,
            "eval_count": 30,
        }

        mock_generation = Mock()
        mock_generation.message = mock_message

        mock_response = Mock()
        mock_response.llm_output = None
        mock_response.generations = [[mock_generation]]

        callback.on_llm_end(mock_response)

        assert callback.counts["total_prompt_tokens"] == 60
        assert callback.counts["total_completion_tokens"] == 30

    def test_detects_context_overflow(self):
        """Should detect context overflow when near limit."""
        callback = TokenCountingCallback()
        callback.current_model = "model"
        callback.context_limit = 1000
        callback.original_prompt_estimate = 1100  # More than what was processed
        callback.counts["by_model"]["model"] = {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
            "calls": 1,
        }

        mock_message = Mock()
        mock_message.usage_metadata = None
        mock_message.response_metadata = {
            "prompt_eval_count": 950,  # 95% of context_limit
            "eval_count": 50,
        }

        mock_generation = Mock()
        mock_generation.message = mock_message

        mock_response = Mock()
        mock_response.llm_output = None
        mock_response.generations = [[mock_generation]]

        callback.on_llm_end(mock_response)

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

    def test_updates_model_counts(self):
        """Should update per-model token counts."""
        callback = TokenCountingCallback()
        callback.current_model = "gpt-4"
        callback.counts["by_model"]["gpt-4"] = {
            "prompt_tokens": 10,
            "completion_tokens": 5,
            "total_tokens": 15,
            "calls": 1,
        }

        mock_response = Mock()
        mock_response.llm_output = {
            "token_usage": {
                "prompt_tokens": 20,
                "completion_tokens": 10,
                "total_tokens": 30,
            }
        }
        mock_response.generations = []

        callback.on_llm_end(mock_response)

        assert callback.counts["by_model"]["gpt-4"]["prompt_tokens"] == 30
        assert callback.counts["by_model"]["gpt-4"]["completion_tokens"] == 15


class TestTokenCountingCallbackOnLlmError:
    """Tests for on_llm_error method."""

    def test_tracks_error_status(self):
        """Should set success_status to 'error'."""
        callback = TokenCountingCallback()
        callback.start_time = time.time()

        callback.on_llm_error(ValueError("Test error"))

        assert callback.success_status == "error"

    def test_tracks_error_type(self):
        """Should record error type name."""
        callback = TokenCountingCallback()
        callback.start_time = time.time()

        callback.on_llm_error(TimeoutError("Timeout"))

        assert callback.error_type == "TimeoutError"

    def test_calculates_response_time_on_error(self):
        """Should calculate response time even on error."""
        callback = TokenCountingCallback()
        callback.start_time = time.time() - 0.1  # 100ms ago

        callback.on_llm_error(Exception("Error"))

        assert callback.response_time_ms is not None
        assert callback.response_time_ms >= 100


class TestTokenCountingCallbackSaveToDb:
    """Tests for _save_to_db method."""

    def test_uses_thread_metrics_from_background(self):
        """Should use thread metrics writer from background thread."""
        callback = TokenCountingCallback(
            research_id="test-uuid",
            research_context={
                "username": "testuser",
                "user_password": "testpass",
            },
        )
        callback.current_model = "gpt-4"
        callback.current_provider = "openai"

        mock_writer = MagicMock()

        with patch("threading.current_thread") as mock_thread:
            mock_thread.return_value.name = "WorkerThread"
            with patch(
                "local_deep_research.database.thread_metrics.metrics_writer",
                mock_writer,
            ):
                callback._save_to_db(100, 50)

                mock_writer.set_user_password.assert_called_with(
                    "testuser", "testpass"
                )
                mock_writer.write_token_metrics.assert_called()


class TestTokenCountingCallbackGetCounts:
    """Tests for get_counts method."""

    def test_returns_counts_dict(self):
        """Should return the counts dictionary."""
        callback = TokenCountingCallback()
        callback.counts["total_tokens"] = 150

        result = callback.get_counts()

        assert result["total_tokens"] == 150


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

    def test_returns_overflow_fields(self):
        """Should return context overflow fields."""
        callback = TokenCountingCallback()
        callback.context_limit = 4096
        callback.context_truncated = True
        callback.tokens_truncated = 500
        callback.truncation_ratio = 0.1

        result = callback._get_context_overflow_fields()

        assert result["context_limit"] == 4096
        assert result["context_truncated"] is True
        assert result["tokens_truncated"] == 500

    def test_returns_none_for_non_truncated(self):
        """Should return None for truncation fields when not truncated."""
        callback = TokenCountingCallback()
        callback.context_truncated = False

        result = callback._get_context_overflow_fields()

        assert result["tokens_truncated"] is None
        assert result["truncation_ratio"] is None


class TestTokenCounterInit:
    """Tests for TokenCounter initialization."""

    def test_initializes(self):
        """Should initialize successfully."""
        counter = TokenCounter()
        assert hasattr(counter, "_get_metrics_from_encrypted_db")


class TestTokenCounterCreateCallback:
    """Tests for create_callback method."""

    def test_returns_callback_instance(self):
        """Should return TokenCountingCallback instance."""
        counter = TokenCounter()

        callback = counter.create_callback()

        assert isinstance(callback, TokenCountingCallback)

    def test_passes_research_id(self):
        """Should pass research_id to callback."""
        counter = TokenCounter()

        callback = counter.create_callback(research_id="test-uuid")

        assert callback.research_id == "test-uuid"

    def test_passes_research_context(self):
        """Should pass research_context to callback."""
        counter = TokenCounter()
        context = {"query": "test", "mode": "quick"}

        callback = counter.create_callback(research_context=context)

        assert callback.research_context == context


class TestTokenCounterGetResearchMetrics:
    """Tests for get_research_metrics method."""

    def test_method_exists(self):
        """TokenCounter should have get_research_metrics method."""
        counter = TokenCounter()
        assert hasattr(counter, "get_research_metrics")
        assert callable(counter.get_research_metrics)


class TestTokenCounterGetOverallMetrics:
    """Tests for get_overall_metrics method."""

    def test_returns_encrypted_db_metrics(self):
        """Should return metrics from encrypted DB directly."""
        counter = TokenCounter()

        encrypted_metrics = {
            "total_tokens": 1000,
            "total_researches": 5,
            "by_model": [
                {
                    "model": "gpt-4",
                    "tokens": 1000,
                    "calls": 5,
                    "prompt_tokens": 600,
                    "completion_tokens": 400,
                }
            ],
            "recent_researches": [],
            "token_breakdown": {
                "total_input_tokens": 600,
                "total_output_tokens": 400,
                "avg_input_tokens": 120,
                "avg_output_tokens": 80,
                "avg_total_tokens": 200,
            },
        }

        with patch.object(
            counter,
            "_get_metrics_from_encrypted_db",
            return_value=encrypted_metrics,
        ):
            result = counter.get_overall_metrics()

        assert result["total_tokens"] == 1000


class TestTokenCounterGetEmptyMetrics:
    """Tests for _get_empty_metrics method."""

    def test_returns_correct_structure(self):
        """Should return empty metrics with correct structure."""
        counter = TokenCounter()

        result = counter._get_empty_metrics()

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


class TestTokenCounterGetEnhancedMetrics:
    """Tests for get_enhanced_metrics method."""

    def test_method_exists(self):
        """TokenCounter should have get_enhanced_metrics method."""
        counter = TokenCounter()
        assert hasattr(counter, "get_enhanced_metrics")
        assert callable(counter.get_enhanced_metrics)


class TestTokenCounterGetResearchTimelineMetrics:
    """Tests for get_research_timeline_metrics method."""

    def test_method_exists(self):
        """TokenCounter should have get_research_timeline_metrics method."""
        counter = TokenCounter()
        assert hasattr(counter, "get_research_timeline_metrics")
        assert callable(counter.get_research_timeline_metrics)