generated from coulomb/repo-seed
72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
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,
|
|
}
|