SupplyAgent/supply_infra/db/models/llm_usage_event.py

from __future__ import annotations

from datetime import datetime
from decimal import Decimal

from sqlalchemy import (
    BigInteger,
    DateTime,
    Index,
    Integer,
    Numeric,
    String,
    func,
)
from sqlalchemy.orm import Mapped, mapped_column

from supply_infra.db.base import Base


class LlmUsageEvent(Base):
    """单次 LLM 调用的用量与费用明细。"""

    __tablename__ = "llm_usage_event"
    __table_args__ = (
        Index("idx_llm_usage_biz_dt", "biz_dt"),
        Index("idx_llm_usage_biz_agent", "biz_dt", "agent_name"),
        Index("idx_llm_usage_biz_model", "biz_dt", "model"),
        Index("idx_llm_usage_run_id", "run_id"),
        Index("idx_llm_usage_created", "created_at"),
    )

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True, autoincrement=True)
    biz_dt: Mapped[str] = mapped_column(
        String(8),
        nullable=False,
        comment="业务日 YYYYMMDD（Asia/Shanghai）",
    )
    provider: Mapped[str] = mapped_column(
        String(32),
        nullable=False,
        default="openrouter",
        comment="提供商 openrouter / dashscope 等",
    )
    model: Mapped[str | None] = mapped_column(String(128), nullable=True, comment="模型名")
    agent_name: Mapped[str | None] = mapped_column(
        String(128),
        nullable=True,
        comment="Agent 名称",
    )
    run_id: Mapped[str | None] = mapped_column(String(128), nullable=True, comment="Agent run id")
    iteration: Mapped[int | None] = mapped_column(Integer, nullable=True, comment="ReAct 迭代序号")
    prompt_tokens: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        default=0,
        comment="输入 token 数",
    )
    completion_tokens: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        default=0,
        comment="输出 token 数",
    )
    total_tokens: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        default=0,
        comment="总 token 数",
    )
    cost_usd: Mapped[Decimal | None] = mapped_column(
        Numeric(14, 8),
        nullable=True,
        comment="本次调用费用（USD，优先取 provider 返回的 usage.cost）",
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime,
        nullable=False,
        server_default=func.now(),
        comment="记录时间",
    )