tests/test_telemetry.py

from __future__ import annotations

import unittest
from uuid import uuid4

from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langchain_core.outputs import ChatGeneration, LLMResult

from app.chat_service import CLEARED_TOOL_OUTPUT_PLACEHOLDER
from app import telemetry
from app.telemetry import (
    TurnUsageHandler,
    context_window_stats,
    estimate_cost_usd,
    aggregate_cost_breakdown,
    pop_turn_events,
    pop_turn_signals,
    pricing_for_model,
    record_turn_signal,
    record_turn_event,
    record_turn_signal_max,
    reset_turn_signals,
    usage_totals,
)


class UsageTotalsTests(unittest.TestCase):
    def test_sums_across_models_including_cache_details(self) -> None:
        usage = {
            "gemini-3.5-flash": {
                "input_tokens": 1_000,
                "output_tokens": 200,
                "total_tokens": 1_200,
                "input_token_details": {"cache_read": 400},
            },
            "claude-haiku-4-5-20251001": {
                "input_tokens": 500,
                "output_tokens": 100,
                "total_tokens": 600,
                "input_token_details": {"cache_read": 50, "cache_creation": 25},
            },
        }
        totals = usage_totals(usage)
        self.assertEqual(totals["input_tokens"], 1_500)
        self.assertEqual(totals["output_tokens"], 300)
        self.assertEqual(totals["total_tokens"], 1_800)
        self.assertEqual(totals["cache_read_tokens"], 450)
        self.assertEqual(totals["cache_creation_tokens"], 25)

    def test_missing_details_default_to_zero(self) -> None:
        totals = usage_totals({"m": {"input_tokens": 10, "output_tokens": 5}})
        self.assertEqual(totals["cache_read_tokens"], 0)
        self.assertEqual(totals["total_tokens"], 0)


class CostEstimateTests(unittest.TestCase):
    def test_dated_model_name_matches_family_pricing(self) -> None:
        self.assertIsNotNone(pricing_for_model("claude-haiku-4-5-20251001"))
        self.assertIsNone(pricing_for_model("some-unknown-model"))

    def test_cache_tokens_priced_separately(self) -> None:
        # claude-haiku-4-5: input $1, output $5, cache read $0.10, write $1.25
        # per MTok. input_tokens includes the cached buckets, so the plain
        # bucket is 1M - 400k - 100k = 500k.
        usage = {
            "claude-haiku-4-5-20251001": {
                "input_tokens": 1_000_000,
                "output_tokens": 200_000,
                "input_token_details": {
                    "cache_read": 400_000,
                    "cache_creation": 100_000,
                },
            }
        }
        expected = 0.5 * 1.00 + 0.4 * 0.10 + 0.1 * 1.25 + 0.2 * 5.00
        self.assertAlmostEqual(estimate_cost_usd(usage), expected)

    def test_cache_creation_without_write_rate_bills_as_input(self) -> None:
        # gemini-3.5-flash: input $1.50, cache creation defaults to input rate
        # because Gemini's hourly explicit-cache storage is not represented in
        # per-turn usage metadata.
        usage = {
            "gemini-3.5-flash": {
                "input_tokens": 1_000_000,
                "output_tokens": 0,
                "input_token_details": {"cache_creation": 200_000},
            }
        }
        expected = 0.8 * 1.50 + 0.2 * 1.50
        self.assertAlmostEqual(estimate_cost_usd(usage), expected)

    def test_unknown_model_yields_none_not_zero(self) -> None:
        self.assertIsNone(estimate_cost_usd({"mystery-model": {"input_tokens": 10}}))
        mixed = {
            "gemini-3.5-flash": {"input_tokens": 10, "output_tokens": 1},
            "mystery-model": {"input_tokens": 10, "output_tokens": 1},
        }
        self.assertIsNone(estimate_cost_usd(mixed))

    def test_deepseek_cost_breakdown_is_mutually_exclusive(self) -> None:
        usage = {
            "deepseek-v4-flash": {
                "input_tokens": 1_000_000,
                "output_tokens": 100_000,
                "input_token_details": {"cache_read": 900_000},
            }
        }
        breakdown = aggregate_cost_breakdown(usage)
        self.assertIsNotNone(breakdown)
        self.assertAlmostEqual(breakdown["cache_miss_input_usd"], 0.014)
        self.assertAlmostEqual(breakdown["cache_read_input_usd"], 0.00252)
        self.assertAlmostEqual(breakdown["output_usd"], 0.028)
        self.assertAlmostEqual(
            breakdown["total_usd"],
            breakdown["cache_miss_input_usd"]
            + breakdown["cache_read_input_usd"]
            + breakdown["cache_creation_input_usd"]
            + breakdown["output_usd"],
        )


class ContextWindowStatsTests(unittest.TestCase):
    def test_counts_summaries_and_cleared_tool_outputs(self) -> None:
        messages = [
            HumanMessage(
                content="Here is a summary of the conversation to date: ...",
                additional_kwargs={"lc_source": "summarization"},
            ),
            HumanMessage(content="What is RAG?"),
            AIMessage(content="RAG is retrieval-augmented generation."),
            ToolMessage(
                content=CLEARED_TOOL_OUTPUT_PLACEHOLDER,
                tool_call_id="call_1",
            ),
            ToolMessage(content="$ ls\nwiki", tool_call_id="call_2"),
        ]
        stats = context_window_stats(messages, CLEARED_TOOL_OUTPUT_PLACEHOLDER)
        self.assertEqual(stats["context_messages"], 5)
        self.assertEqual(stats["summary_messages"], 1)
        self.assertEqual(stats["cleared_tool_outputs"], 1)
        self.assertGreater(stats["context_tokens_approx"], 0)

    def test_empty_context(self) -> None:
        stats = context_window_stats([], CLEARED_TOOL_OUTPUT_PLACEHOLDER)
        self.assertEqual(stats["context_messages"], 0)
        self.assertEqual(stats["context_tokens_approx"], 0)


class TurnUsageHandlerTests(unittest.TestCase):
    def test_counts_calls_and_aggregates_usage(self) -> None:
        handler = TurnUsageHandler()

        def result(input_tokens: int, output_tokens: int) -> LLMResult:
            message = AIMessage(
                content="ok",
                usage_metadata={
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens,
                    "total_tokens": input_tokens + output_tokens,
                },
                response_metadata={"model_name": "gemini-3.5-flash"},
            )
            return LLMResult(generations=[[ChatGeneration(message=message)]])

        handler.on_llm_end(result(100, 20))
        handler.on_llm_end(result(50, 10))

        self.assertEqual(handler.llm_calls, 2)
        usage = handler.usage_metadata["gemini-3.5-flash"]
        self.assertEqual(usage["input_tokens"], 150)
        self.assertEqual(usage["output_tokens"], 30)

    def test_records_one_explanatory_row_per_call(self) -> None:
        handler = TurnUsageHandler()
        run_id = uuid4()
        handler.on_chat_model_start(
            {},
            [[HumanMessage(content="hello")]],
            run_id=run_id,
            metadata={"lc_source": "summarization"},
        )
        message = AIMessage(
            content="summary",
            usage_metadata={
                "input_tokens": 100,
                "output_tokens": 20,
                "total_tokens": 120,
                "input_token_details": {"cache_read": 80},
            },
            response_metadata={"model_name": "deepseek-v4-flash"},
        )
        handler.on_llm_end(
            LLMResult(generations=[[ChatGeneration(message=message)]]),
            run_id=run_id,
        )
        call = handler.model_calls[0]
        self.assertEqual(call["source"], "summarization")
        self.assertEqual(call["cache_read_tokens"], 80)
        self.assertEqual(call["cache_miss_tokens"], 20)
        self.assertTrue(call["cache_details_reported"])
        self.assertGreater(call["request_context_tokens_approx"], 0)
        self.assertAlmostEqual(
            call["cost"]["total_usd"],
            estimate_cost_usd({"deepseek-v4-flash": message.usage_metadata}),
        )


class LangchainOpenAICacheFieldContractTests(unittest.TestCase):
    """Pin the langchain-openai usage conversion our cache accounting rides on.

    DeepSeek (and OpenAI) report cached prompt tokens as
    ``prompt_tokens_details.cached_tokens``; langchain-openai must surface that
    as ``usage_metadata.input_token_details.cache_read`` or TurnUsageHandler's
    cache buckets (and the ~50x DeepSeek cache-read discount) silently read 0.
    """

    def _convert(self, payload: dict) -> dict:
        try:
            from langchain_openai.chat_models.base import _create_usage_metadata
        except ImportError as exc:
            self.fail(
                "langchain_openai.chat_models.base._create_usage_metadata is no "
                f"longer importable ({exc}). A langchain-openai upgrade moved "
                "the usage conversion; re-pin the prompt_tokens_details."
                "cached_tokens -> input_token_details.cache_read mapping "
                "against its new location."
            )
        return _create_usage_metadata(payload)

    def test_cached_tokens_map_to_cache_read_details(self) -> None:
        usage = self._convert(
            {
                "prompt_tokens": 1_000,
                "completion_tokens": 100,
                "total_tokens": 1_100,
                "prompt_tokens_details": {"cached_tokens": 900},
            }
        )
        # LangChain convention: input_tokens INCLUDES the cached bucket; the
        # cost code carves cache_read out instead of adding it on top.
        self.assertEqual(usage["input_tokens"], 1_000)
        self.assertEqual(usage["output_tokens"], 100)
        self.assertEqual(usage["input_token_details"]["cache_read"], 900)

    def test_converted_usage_flows_through_turn_usage_handler(self) -> None:
        usage = self._convert(
            {
                "prompt_tokens": 1_000,
                "completion_tokens": 100,
                "total_tokens": 1_100,
                "prompt_tokens_details": {"cached_tokens": 900},
            }
        )
        handler = TurnUsageHandler()
        run_id = uuid4()
        handler.on_chat_model_start({}, [[HumanMessage(content="q")]], run_id=run_id)
        message = AIMessage(
            content="answer",
            usage_metadata=usage,
            response_metadata={"model_name": "deepseek-v4-flash"},
        )
        handler.on_llm_end(
            LLMResult(generations=[[ChatGeneration(message=message)]]),
            run_id=run_id,
        )
        call = handler.model_calls[0]
        self.assertEqual(call["cache_read_tokens"], 900)
        self.assertEqual(call["cache_miss_tokens"], 100)
        self.assertTrue(call["cache_details_reported"])
        # deepseek-v4-flash: $0.14 miss / $0.0028 cache-read / $0.28 output
        # per MTok, so the cache discount must show up in the billed cost.
        expected = (100 * 0.14 + 900 * 0.0028 + 100 * 0.28) / 1_000_000
        self.assertAlmostEqual(call["cost"]["total_usd"], expected)


class TurnSignalRegistryTests(unittest.TestCase):
    def test_accumulates_and_pops_per_turn(self) -> None:
        reset_turn_signals("turn-a")
        record_turn_signal("turn-a", "dropped_messages", 3)
        record_turn_signal("turn-a", "dropped_messages", 2)
        record_turn_signal("turn-a", "truncated_tool_outputs", 1)
        signals = pop_turn_signals("turn-a")
        self.assertEqual(signals["dropped_messages"], 5)
        self.assertEqual(signals["truncated_tool_outputs"], 1)
        # Popping clears the entry: a second pop is empty.
        self.assertEqual(pop_turn_signals("turn-a"), {})

    def test_structured_events_are_isolated_and_popped(self) -> None:
        reset_turn_signals("turn-events")
        record_turn_event("turn-events", {"event": "summarization", "tokens": 9})
        self.assertEqual(
            pop_turn_events("turn-events"),
            [{"event": "summarization", "tokens": 9}],
        )
        self.assertEqual(pop_turn_events("turn-events"), [])

    def test_turns_are_isolated_and_noops_are_ignored(self) -> None:
        reset_turn_signals("turn-x")
        reset_turn_signals("turn-y")
        record_turn_signal("turn-x", "dropped_messages", 4)
        record_turn_signal("turn-y", "dropped_messages", 0)  # no-op
        record_turn_signal("", "dropped_messages", 9)  # no turn id -> ignored
        self.assertEqual(pop_turn_signals("turn-x"), {"dropped_messages": 4})
        self.assertEqual(pop_turn_signals("turn-y"), {})

    def test_max_records_peak_not_sum(self) -> None:
        # A middleware fires once per model call within a turn; the max across
        # calls is the real per-turn figure, not the (overlapping) sum.
        reset_turn_signals("turn-m")
        record_turn_signal_max("turn-m", "dropped_messages", 5)
        record_turn_signal_max("turn-m", "dropped_messages", 8)
        record_turn_signal_max("turn-m", "dropped_messages", 3)
        self.assertEqual(pop_turn_signals("turn-m"), {"dropped_messages": 8})

    def test_overflow_evicts_oldest_keeps_recent(self) -> None:
        # The cap must evict the OLDEST turns, never wipe the freshest in-flight
        # ones (the concurrency-safety fix).
        reset_turn_signals("oldest-turn")
        record_turn_signal_max("oldest-turn", "dropped_messages", 1)
        for i in range(telemetry._MAX_TRACKED_TURNS + 5):
            reset_turn_signals(f"filler-{i}")
        reset_turn_signals("recent-turn")
        record_turn_signal_max("recent-turn", "dropped_messages", 7)
        self.assertEqual(pop_turn_signals("recent-turn"), {"dropped_messages": 7})
        self.assertEqual(pop_turn_signals("oldest-turn"), {})  # evicted


if __name__ == "__main__":
    unittest.main()