"""Metrics model for ATIF trajectories."""

from typing import Any

from pydantic import BaseModel, Field


class Metrics(BaseModel):
    """LLM operational and confidence data."""

    prompt_tokens: int | None = Field(
        default=None,
        description="Total input tokens including cached and non-cached",
    )
    completion_tokens: int | None = Field(
        default=None,
        description="Total tokens generated by the LLM response",
    )
    cached_tokens: int | None = Field(
        default=None,
        description="Subset of prompt_tokens that were cache hits",
    )
    cost_usd: float | None = Field(
        default=None,
        description="Monetary cost of the API call in USD",
    )
    prompt_token_ids: list[int] | None = Field(
        default=None,
        description="Token IDs for prompt (input) tokens sent to the LLM, including chat history (if applicable)",
    )
    completion_token_ids: list[int] | None = Field(
        default=None,
        description="Token IDs for completion (response) tokens, enabling RL training without retokenization drift",
    )
    logprobs: list[float] | None = Field(
        default=None,
        description="Log probability assigned to each generated token",
    )
    extra: dict[str, Any] | None = Field(
        default=None,
        description="Other metrics",
    )

    model_config = {"extra": "forbid"}
