state-hub/api/schemas/token_event.py

import uuid
from datetime import datetime

from pydantic import BaseModel, ConfigDict, computed_field


class TokenEventCreate(BaseModel):
    tokens_in: int
    tokens_out: int
    task_id: uuid.UUID | None = None
    workstream_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


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

    id: uuid.UUID
    tokens_in: int
    tokens_out: int
    task_id: uuid.UUID | None = None
    workstream_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
    created_at: datetime

    @computed_field
    @property
    def tokens_total(self) -> int:
        return 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]


class TokenEventPatch(BaseModel):
    tokens_in: int | None = None
    tokens_out: int | None = None
    note: str | None = None
    model: str | None = None
    agent: str | 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]