SupplyAgent/tests/supply_agent/test_llm_usage_hook.py

"""LLM usage recording hook stays free of supply_infra at import time."""

from __future__ import annotations

import importlib
from pathlib import Path

from supply_agent.logging.logger import AgentLogger
from supply_agent.logging.usage import (
    get_llm_usage_recorder,
    record_llm_usage,
    set_llm_usage_recorder,
)
from supply_agent.types import Message, Role


def test_usage_module_does_not_import_supply_infra() -> None:
    source = importlib.util.find_spec("supply_agent.logging.usage")
    assert source is not None and source.origin is not None
    text = Path(source.origin).read_text(encoding="utf-8")
    assert "import supply_infra" not in text
    assert "from supply_infra" not in text


def test_record_llm_usage_is_noop_without_recorder() -> None:
    previous = get_llm_usage_recorder()
    try:
        set_llm_usage_recorder(None)
        record_llm_usage({"usage": {"cost": 0.01}})
    finally:
        set_llm_usage_recorder(previous)


def test_log_llm_output_invokes_usage_recorder(tmp_path: Path) -> None:
    previous = get_llm_usage_recorder()
    seen: list[dict] = []

    def _recorder(payload: dict) -> None:
        seen.append(payload)

    try:
        set_llm_usage_recorder(_recorder)
        logger = AgentLogger(tmp_path, enabled=True)
        logger.start_run("hello", model="google/gemini-2.5-flash", agent_name="demo_agent")

        class _Usage:
            def model_dump(self) -> dict:
                return {
                    "prompt_tokens": 10,
                    "completion_tokens": 5,
                    "total_tokens": 15,
                    "cost": 0.0012,
                }

        class _Response:
            model = "google/gemini-2.5-flash"
            usage = _Usage()

            def model_dump(self) -> dict:
                return {"model": self.model, "usage": self.usage.model_dump()}

        logger.log_llm_output(
            1,
            Message(role=Role.ASSISTANT, content="ok"),
            raw_response=_Response(),
            model="google/gemini-2.5-flash",
        )
        assert len(seen) == 1
        assert seen[0]["agent_name"] == "demo_agent"
        assert seen[0]["model"] == "google/gemini-2.5-flash"
        assert seen[0]["usage"]["cost"] == 0.0012
        assert seen[0]["usage"]["total_tokens"] == 15
    finally:
        set_llm_usage_recorder(previous)