from __future__ import annotations from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path from typing import Any def parse_iso(value: str) -> datetime: raw = value.strip() if raw.endswith("Z"): raw = raw[:-1] + "+00:00" if "T" not in raw: raw = f"{raw}T00:00:00+00:00" parsed = datetime.fromisoformat(raw) if parsed.tzinfo is None: parsed = parsed.replace(tzinfo=timezone.utc) return parsed.astimezone(timezone.utc) @dataclass class TokenSourceRecord: source_provider: str source_id: str source_path: Path source_created_at: datetime | None session_id: str | None = None cwd: str | None = None model: str | None = None agent: str | None = None tokens_in: int = 0 tokens_out: int = 0 cached_input_tokens: int = 0 reasoning_output_tokens: int = 0 raw_total_tokens: int | None = None parser_version: str | None = None confidence: float = 1.0 raw_metadata: dict[str, Any] = field(default_factory=dict) @property def tokens_total(self) -> int: return self.tokens_in + self.tokens_out def to_token_event_payload(self, repo_id: str | None = None) -> dict[str, Any]: raw_total = self.raw_total_tokens if raw_total is None: raw_total = self.tokens_in + self.tokens_out created_at = self.source_created_at.isoformat() if self.source_created_at else None return { "tokens_in": self.tokens_in, "tokens_out": self.tokens_out, "repo_id": repo_id, "session_id": self.session_id, "model": self.model, "agent": self.agent, "ref_type": "session", "ref_id": self.source_id, "note": f"measured:{self.source_provider}", "created_at": created_at, "measurement_kind": "measured", "source_provider": self.source_provider, "source_id": self.source_id, "source_path": str(self.source_path), "source_created_at": created_at, "parser_version": self.parser_version, "confidence": self.confidence, "cached_input_tokens": self.cached_input_tokens, "reasoning_output_tokens": self.reasoning_output_tokens, "raw_total_tokens": raw_total, "raw_metadata": self.raw_metadata, }