/tests/metrics/test_token_counter_provider_coverage.py

"""Provider-level coverage tests for token_counter.py.

Targets specific uncovered lines:
- on_llm_start: model extraction fallbacks (kwargs direct, serialized kwargs model_name)
- on_llm_start: provider extraction for unrecognized _type with provider kwarg
- on_llm_end: token extraction from llm_output 'usage' key (alternative to 'token_usage')
- on_llm_end: Ollama context overflow with zero original_prompt_estimate guard
- _save_to_db: background thread success path with full token_data dict verification
"""

import threading
from unittest.mock import MagicMock, patch

from langchain_core.outputs import LLMResult

from local_deep_research.metrics.token_counter import TokenCountingCallback


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


def _make_callback(research_context=None, research_id="rid-test", **overrides):
    """Build a TokenCountingCallback with controllable state."""
    ctx = research_context if research_context is not None else {}
    cb = TokenCountingCallback(research_id=research_id, research_context=ctx)
    for k, v in overrides.items():
        setattr(cb, k, v)
    return cb


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 or []
    return result


def _make_generation(usage_metadata=None, response_metadata=None):
    """Build a mock generation with a message carrying metadata."""
    gen = MagicMock()
    msg = MagicMock()
    msg.usage_metadata = usage_metadata
    msg.response_metadata = (
        response_metadata if response_metadata is not None else {}
    )
    gen.message = msg
    return gen


def _setup_model_counts(cb, model_name="test-model", provider="openai"):
    """Register model in cb.counts so on_llm_end can update them."""
    cb.current_model = model_name
    cb.current_provider = provider
    cb.counts["by_model"][model_name] = {
        "prompt_tokens": 0,
        "completion_tokens": 0,
        "total_tokens": 0,
        "calls": 1,
        "provider": provider,
    }


def _patch_worker_thread():
    t = MagicMock()
    t.name = "WorkerThread-1"
    return patch.object(threading, "current_thread", return_value=t)


# ---------------------------------------------------------------------------
# 1. on_llm_start – model extraction fallback: kwargs["model"] when
#    invocation_params has no model key (line 163)
# ---------------------------------------------------------------------------


class TestModelFromKwargsModel:
    """Model extracted from kwargs['model'] when invocation_params is empty."""

    def test_model_from_kwargs_model_key(self):
        """kwargs['model'] is used when invocation_params has no model."""
        cb = _make_callback()
        cb.on_llm_start(
            {},
            ["hello"],
            invocation_params={},
            model="gemma-2-27b",
        )
        assert cb.current_model == "gemma-2-27b"

    def test_model_from_kwargs_model_name_key(self):
        """kwargs['model_name'] is used when invocation_params and kwargs['model'] are absent."""
        cb = _make_callback()
        cb.on_llm_start({}, ["hello"], model_name="codellama-34b")
        assert cb.current_model == "codellama-34b"


# ---------------------------------------------------------------------------
# 2. on_llm_start – model from serialized["kwargs"]["model_name"]
#    (line 167-169, specifically the model_name branch)
# ---------------------------------------------------------------------------


class TestModelFromSerializedKwargsModelName:
    """Model extracted from serialized['kwargs']['model_name'] when 'model' absent."""

    def test_serialized_kwargs_model_name(self):
        """serialized['kwargs']['model_name'] provides the model when 'model' key missing."""
        cb = _make_callback()
        cb.on_llm_start(
            {"kwargs": {"model_name": "deepseek-coder"}},
            ["prompt"],
        )
        assert cb.current_model == "deepseek-coder"


# ---------------------------------------------------------------------------
# 3. on_llm_start – provider from kwargs when _type is unrecognized
#    (line 210-211: else branch inside the if "_type" in serialized block)
# ---------------------------------------------------------------------------


class TestProviderFromKwargsWithUnrecognizedType:
    """When _type exists but is not ChatOllama/ChatOpenAI/ChatAnthropic,
    provider falls back to kwargs.get('provider', 'unknown')."""

    def test_unrecognized_type_with_provider_kwarg(self):
        """Unrecognized _type uses provider kwarg."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatGoogleGenAI", "kwargs": {"model": "gemini-pro"}},
            ["prompt"],
            provider="google",
        )
        assert cb.current_provider == "google"

    def test_unrecognized_type_without_provider_kwarg(self):
        """Unrecognized _type without provider kwarg defaults to 'unknown'."""
        cb = _make_callback()
        cb.on_llm_start(
            {"_type": "ChatGoogleGenAI", "kwargs": {"model": "gemini-pro"}},
            ["prompt"],
        )
        assert cb.current_provider == "unknown"


# ---------------------------------------------------------------------------
# 4. on_llm_end – token extraction from llm_output 'usage' key
#    (line 240-241: the or branch: response.llm_output.get("usage", {}))
# ---------------------------------------------------------------------------


class TestTokenExtractionFromUsageKey:
    """Tokens extracted from llm_output['usage'] when 'token_usage' absent."""

    def test_usage_key_provides_tokens(self):
        """llm_output['usage'] provides prompt_tokens/completion_tokens."""
        cb = _make_callback()
        _setup_model_counts(cb, "gpt-4o", "openai")

        result = _make_llm_result(
            llm_output={
                "usage": {
                    "prompt_tokens": 45,
                    "completion_tokens": 30,
                    "total_tokens": 75,
                }
            },
        )
        cb.on_llm_end(result)

        assert cb.counts["total_prompt_tokens"] == 45
        assert cb.counts["total_completion_tokens"] == 30
        assert cb.counts["total_tokens"] == 75


# ---------------------------------------------------------------------------
# 5. on_llm_end – Ollama context overflow: original_prompt_estimate == 0
#    guards against division-by-zero (line 312-314)
# ---------------------------------------------------------------------------


class TestContextOverflowZeroPromptEstimate:
    """When original_prompt_estimate is 0 the truncation_ratio is 0, not a ZeroDivisionError."""

    def test_zero_prompt_estimate_no_division_error(self):
        """truncation_ratio stays 0 when original_prompt_estimate is 0."""
        cb = _make_callback(research_context={"context_limit": 100})
        _setup_model_counts(cb, "llama3", "ollama")
        # context_limit is set during on_llm_start; set it directly here
        cb.context_limit = 100
        cb.original_prompt_estimate = 0

        gen = _make_generation(
            usage_metadata=None,
            response_metadata={
                "prompt_eval_count": 96,  # >= 95 (95% of 100)
                "eval_count": 20,
            },
        )
        result = _make_llm_result(generations=[[gen]])
        cb.on_llm_end(result)

        assert cb.context_truncated is True
        # original_prompt_estimate (0) <= prompt_eval_count (96), so
        # the inner if branch is not entered; tokens_truncated stays 0
        assert cb.tokens_truncated == 0
        assert cb.truncation_ratio == 0.0

    def test_context_overflow_calculates_truncation(self):
        """When prompt estimate > actual, truncation fields are populated."""
        cb = _make_callback(research_context={"context_limit": 1000})
        _setup_model_counts(cb, "llama3", "ollama")
        # context_limit is set during on_llm_start; set it directly here
        cb.context_limit = 1000
        cb.original_prompt_estimate = 1200

        gen = _make_generation(
            usage_metadata=None,
            response_metadata={
                "prompt_eval_count": 960,  # >= 950 (95% of 1000)
                "eval_count": 50,
            },
        )
        result = _make_llm_result(generations=[[gen]])
        cb.on_llm_end(result)

        assert cb.context_truncated is True
        assert cb.tokens_truncated == 240  # 1200 - 960
        expected_ratio = 240 / 1200
        assert abs(cb.truncation_ratio - expected_ratio) < 1e-9


# ---------------------------------------------------------------------------
# 6. _save_to_db – background thread success path
#    (lines 408-482): verifies metrics_writer is called with correct data
# ---------------------------------------------------------------------------


class TestSaveToDbBackgroundThread:
    """_save_to_db in a background thread writes via metrics_writer."""

    def test_metrics_writer_called_with_token_data(self):
        """metrics_writer.write_token_metrics receives correct token_data dict."""
        cb = _make_callback(
            research_context={
                "username": "alice",
                "user_password": "secret",
                "research_query": "quantum computing",
                "research_mode": "detailed",
                "research_phase": "search",
                "search_iteration": 2,
                "search_engines_planned": ["google", "bing"],
                "search_engine_selected": "google",
            },
            research_id="rid-42",
        )
        cb.current_model = "gpt-4o"
        cb.current_provider = "openai"
        cb.response_time_ms = 350
        cb.success_status = "success"
        cb.error_type = None
        cb.calling_file = "runner.py"
        cb.calling_function = "execute"
        cb.call_stack = "runner.py:execute:10"

        with _patch_worker_thread():
            with patch(
                "local_deep_research.database.thread_metrics.metrics_writer"
            ) as mock_writer:
                cb._save_to_db(prompt_tokens=100, completion_tokens=50)

                mock_writer.set_user_password.assert_called_once_with(
                    "alice", "secret"
                )
                mock_writer.write_token_metrics.assert_called_once()

                call_args = mock_writer.write_token_metrics.call_args
                assert call_args[0][0] == "alice"
                assert call_args[0][1] == "rid-42"

                token_data = call_args[0][2]
                assert token_data["model_name"] == "gpt-4o"
                assert token_data["provider"] == "openai"
                assert token_data["prompt_tokens"] == 100
                assert token_data["completion_tokens"] == 50
                assert token_data["research_query"] == "quantum computing"
                assert token_data["research_mode"] == "detailed"
                assert token_data["research_phase"] == "search"
                assert token_data["search_iteration"] == 2
                # search_engines_planned list should be JSON-serialized
                assert (
                    token_data["search_engines_planned"] == '["google", "bing"]'
                )
                assert token_data["search_engine_selected"] == "google"
                assert token_data["response_time_ms"] == 350
                assert token_data["calling_file"] == "runner.py"

    def test_no_username_logs_warning_and_returns(self):
        """Missing username in research_context causes early return."""
        cb = _make_callback(
            research_context={"user_password": "secret"},
            research_id="rid-99",
        )
        cb.current_model = "gpt-4"
        cb.current_provider = "openai"

        with _patch_worker_thread():
            with patch(
                "local_deep_research.database.thread_metrics.metrics_writer"
            ) as mock_writer:
                cb._save_to_db(prompt_tokens=10, completion_tokens=5)
                mock_writer.write_token_metrics.assert_not_called()

    def test_no_password_logs_warning_and_returns(self):
        """Missing password in research_context causes early return after username check."""
        cb = _make_callback(
            research_context={"username": "bob"},
            research_id="rid-88",
        )
        cb.current_model = "gpt-4"
        cb.current_provider = "openai"

        with _patch_worker_thread():
            with patch(
                "local_deep_research.database.thread_metrics.metrics_writer"
            ) as mock_writer:
                cb._save_to_db(prompt_tokens=10, completion_tokens=5)
                mock_writer.write_token_metrics.assert_not_called()