state-hub/api/schemas/token_event.py

import uuid
from datetime import datetime
from typing import Any

from pydantic import AliasChoices, BaseModel, ConfigDict, Field, computed_field

from api.schemas.compat import OptionalWorkplanIdCompatMixin


class TokenEventCreate(BaseModel):
    tokens_in: int
    tokens_out: int
    task_id: uuid.UUID | None = None
    workplan_id: uuid.UUID | None = Field(
        default=None,
        validation_alias=AliasChoices("workplan_id", "workstream_id"),
    )
    repo_id: uuid.UUID | None = None
    session_id: str | None = None
    model: str | None = None
    agent: str | None = None
    ref_type: str | None = None
    ref_id: str | None = None
    note: str | None = None
    created_at: datetime | None = None
    measurement_kind: str | None = None
    source_provider: str | None = None
    source_id: str | None = None
    source_path: str | None = None
    source_created_at: datetime | None = None
    parser_version: str | None = None
    confidence: float | None = None
    cached_input_tokens: int | None = None
    reasoning_output_tokens: int | None = None
    raw_total_tokens: int | None = None
    cost_estimated_usd: float | None = None
    raw_metadata: dict[str, Any] | None = None


class TokenEventRead(OptionalWorkplanIdCompatMixin, BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: uuid.UUID
    tokens_in: int
    tokens_out: int
    task_id: uuid.UUID | None = None
    repo_id: uuid.UUID | None = None
    session_id: str | None = None
    model: str | None = None
    agent: str | None = None
    ref_type: str | None = None
    ref_id: str | None = None
    note: str | None = None
    measurement_kind: str
    source_provider: str
    source_id: str | None = None
    source_path: str | None = None
    source_created_at: datetime | None = None
    ingested_at: datetime
    parser_version: str | None = None
    confidence: float
    cached_input_tokens: int
    reasoning_output_tokens: int
    raw_total_tokens: int | None = None
    cost_estimated_usd: float | None = None
    raw_metadata: dict[str, Any] = Field(default_factory=dict)
    created_at: datetime

    @computed_field
    @property
    def tokens_total(self) -> int:
        return self.tokens_in + self.tokens_out

    @computed_field
    @property
    def token_evidence_total(self) -> int:
        return (self.raw_total_tokens or self.tokens_in + self.tokens_out)


class TokenSummary(BaseModel):
    scope: str
    scope_id: str
    tokens_in: int
    tokens_out: int
    tokens_total: int
    event_count: int
    by_model: dict[str, int]
    by_agent: dict[str, int]
    by_measurement_kind: dict[str, int] = Field(default_factory=dict)
    by_source_provider: dict[str, int] = Field(default_factory=dict)


class TokenEventPatch(BaseModel):
    tokens_in: int | None = None
    tokens_out: int | None = None
    task_id: uuid.UUID | None = None
    workplan_id: uuid.UUID | None = Field(
        default=None,
        validation_alias=AliasChoices("workplan_id", "workstream_id"),
    )
    repo_id: uuid.UUID | None = None
    session_id: str | None = None
    note: str | None = None
    model: str | None = None
    agent: str | None = None
    ref_type: str | None = None
    ref_id: str | None = None
    created_at: datetime | None = None
    measurement_kind: str | None = None
    source_provider: str | None = None
    source_id: str | None = None
    source_path: str | None = None
    source_created_at: datetime | None = None
    ingested_at: datetime | None = None
    parser_version: str | None = None
    confidence: float | None = None
    cached_input_tokens: int | None = None
    reasoning_output_tokens: int | None = None
    raw_total_tokens: int | None = None
    cost_estimated_usd: float | None = None
    raw_metadata: dict[str, Any] | None = None


class RepoTokenSummary(BaseModel):
    repo_id: uuid.UUID
    repo_slug: str
    tokens_in: int
    tokens_out: int
    tokens_total: int
    event_count: int
    by_model: dict[str, int]
    by_note: dict[str, int]
    by_measurement_kind: dict[str, int] = Field(default_factory=dict)
    by_source_provider: dict[str, int] = Field(default_factory=dict)


class TokenAggregateRow(BaseModel):
    scope_id: str
    label: str | None = None
    tokens_in: int
    tokens_out: int
    tokens_total: int
    event_count: int
    by_measurement_kind: dict[str, int] = Field(default_factory=dict)
    by_source_provider: dict[str, int] = Field(default_factory=dict)


class TokenAggregateSummary(BaseModel):
    tokens_in: int
    tokens_out: int
    tokens_total: int
    event_count: int
    first_event_at: datetime | None = None
    last_event_at: datetime | None = None
    last_ingested_at: datetime | None = None
    by_repo: list[TokenAggregateRow] = Field(default_factory=list)
    by_workstream: list[TokenAggregateRow] = Field(default_factory=list)
    by_task: list[TokenAggregateRow] = Field(default_factory=list)
    by_model: list[TokenAggregateRow] = Field(default_factory=list)
    by_measurement_kind: dict[str, int] = Field(default_factory=dict)
    by_source_provider: dict[str, int] = Field(default_factory=dict)


class TokenQualitySummary(BaseModel):
    event_count: int
    measured_event_count: int
    estimated_event_count: int
    allocated_event_count: int
    superseded_event_count: int
    fallback_event_count: int
    unattributed_measured_event_count: int
    missing_provenance_event_count: int
    duplicate_source_count: int
    last_codex_ingested_at: datetime | None = None
    last_claude_ingested_at: datetime | None = None
    last_reconciliation_at: datetime | None = None
    by_measurement_kind: dict[str, int] = Field(default_factory=dict)
    by_source_provider: dict[str, int] = Field(default_factory=dict)