SupplyAgent/supply_infra/agent_logging/register.py

"""Register infra callbacks onto the supply_agent framework hooks."""

from __future__ import annotations


def register_agent_logging_hooks() -> None:
    """Wire OSS publish + LLM billing into agent hooks (safe to call repeatedly)."""
    from supply_agent.logging.publish import set_run_artifact_publisher
    from supply_agent.logging.usage import set_llm_usage_recorder
    from supply_infra.agent_logging.publish import publish_run_artifacts_to_oss
    from supply_infra.llm_billing.recorder import record_llm_usage_to_db

    set_run_artifact_publisher(publish_run_artifacts_to_oss)
    set_llm_usage_recorder(record_llm_usage_to_db)