SupplyAgent/supply_agent/logging/usage.py

"""Pluggable hook for persisting LLM usage / cost after each completion.

Keeps the agent framework free of infrastructure (DB) dependencies.
Infra registers a recorder via ``set_llm_usage_recorder``.
"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any

LlmUsageRecorder = Callable[[dict[str, Any]], None]

_recorder: LlmUsageRecorder | None = None


def set_llm_usage_recorder(recorder: LlmUsageRecorder | None) -> None:
    """Install or clear the LLM usage recorder callback."""
    global _recorder
    _recorder = recorder


def get_llm_usage_recorder() -> LlmUsageRecorder | None:
    return _recorder


def record_llm_usage(payload: dict[str, Any]) -> None:
    """Invoke the registered recorder, if any. Never raises to callers."""
    if _recorder is None:
        return
    try:
        _recorder(payload)
    except Exception:
        # Billing must not break agent runs.
        import logging

        logging.getLogger(__name__).exception("LLM usage recorder failed")