life-echo/api/tests/core/test_llm_telemetry.py

"""LLM telemetry helpers (no real Collector required)."""

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import patch

import pytest

from app.core import llm_telemetry
from app.core.config import settings


class TestExtractTokenUsage:
    def test_usage_metadata_object(self) -> None:
        msg = SimpleNamespace(usage_metadata=SimpleNamespace(input_tokens=10, output_tokens=4))
        assert llm_telemetry.extract_token_usage(msg) == (10, 4)

    def test_response_metadata_dict(self) -> None:
        msg = SimpleNamespace(
            usage_metadata=None,
            response_metadata={"token_usage": {"prompt_tokens": 3, "completion_tokens": 7}},
        )
        assert llm_telemetry.extract_token_usage(msg) == (3, 7)

    def test_missing_usage_returns_zero(self) -> None:
        assert llm_telemetry.extract_token_usage(SimpleNamespace()) == (0, 0)


class TestOtelDisabledNoOp:
    def test_record_llm_completion_disabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
        monkeypatch.setattr(settings, "otel_enabled", False)
        llm_telemetry.record_llm_completion(
            agent="Test",
            provider="mock",
            model="m",
            duration_ms=1.0,
            input_tokens=5,
            output_tokens=2,
        )

    def test_langchain_invoke_span_disabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
        monkeypatch.setattr(settings, "otel_enabled", False)
        with llm_telemetry.langchain_invoke_span(
            agent="Test",
            provider="mock",
            model="m",
            call_type="chat",
        ) as ctx:
            ctx["response"] = SimpleNamespace(
                usage_metadata=SimpleNamespace(input_tokens=1, output_tokens=1)
            )
            assert ctx["outcome"] == "ok"


class TestLangchainInvokeSpanRecordsTokens:
    def test_records_completion_with_tokens(self, monkeypatch: pytest.MonkeyPatch) -> None:
        monkeypatch.setattr(settings, "otel_enabled", True)
        recorded: list[dict] = []

        def _capture(**kwargs: object) -> None:
            recorded.append(kwargs)

        with patch.object(llm_telemetry, "record_llm_completion", side_effect=_capture):
            with llm_telemetry.langchain_invoke_span(
                agent="TestAgent",
                provider="mock",
                model="m1",
                call_type="chat",
            ) as ctx:
                ctx["response"] = SimpleNamespace(
                    usage_metadata=SimpleNamespace(input_tokens=11, output_tokens=5)
                )

        assert len(recorded) == 1
        assert recorded[0]["input_tokens"] == 11
        assert recorded[0]["output_tokens"] == 5
        assert recorded[0]["agent"] == "TestAgent"


class TestObserveAinvokeExtraAttributes:
    @pytest.mark.asyncio
    async def test_response_latency_on_span(self, monkeypatch: pytest.MonkeyPatch) -> None:
        monkeypatch.setattr(settings, "otel_enabled", True)
        from opentelemetry.sdk.trace import TracerProvider
        from opentelemetry.sdk.trace.export import SimpleSpanProcessor
        from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
            InMemorySpanExporter,
        )
        from opentelemetry import trace

        exporter = InMemorySpanExporter()
        provider = TracerProvider()
        provider.add_span_processor(SimpleSpanProcessor(exporter))
        monkeypatch.setattr(
            "app.core.llm_telemetry.get_tracer",
            lambda _name: provider.get_tracer("test"),
        )

        class _LLM:
            async def ainvoke(self, messages: list) -> SimpleNamespace:
                return SimpleNamespace(
                    usage_metadata=SimpleNamespace(input_tokens=1, output_tokens=1)
                )

        await llm_telemetry.observe_ainvoke(
            _LLM(),
            [],
            agent="Test",
            provider="mock",
            model="m",
            extra_span_attributes={"llm.custom": "x"},
        )
        spans = exporter.get_finished_spans()
        assert spans
        attrs = dict(spans[-1].attributes or {})
        assert "llm.response_latency_ms" in attrs
        assert attrs.get("llm.custom") == "x"