Fixed and improved token tracking

api/models/token_event.py

@@ -1,8 +1,10 @@
 import uuid
 from datetime import datetime
 
-from sqlalchemy import DateTime, ForeignKey, Integer, Text, func
-from sqlalchemy.dialects.postgresql import UUID
+from typing import Any
+
+from sqlalchemy import DateTime, Float, ForeignKey, Integer, Text, UniqueConstraint, func
+from sqlalchemy.dialects.postgresql import JSONB, UUID
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from api.models.base import Base, new_uuid
@@ -10,6 +12,14 @@ from api.models.base import Base, new_uuid
 
 class TokenEvent(Base):
     __tablename__ = "token_events"
+    __table_args__ = (
+        UniqueConstraint(
+            "measurement_kind",
+            "source_provider",
+            "source_id",
+            name="uq_token_events_source_identity",
+        ),
+    )
 
     id: Mapped[uuid.UUID] = mapped_column(
         UUID(as_uuid=True), primary_key=True, default=new_uuid
@@ -31,6 +41,35 @@ class TokenEvent(Base):
     ref_type: Mapped[str | None] = mapped_column(Text, nullable=True)
     ref_id: Mapped[str | None] = mapped_column(Text, nullable=True)
     note: Mapped[str | None] = mapped_column(Text, nullable=True)
+    measurement_kind: Mapped[str] = mapped_column(
+        Text, nullable=False, default="estimated", server_default="estimated", index=True
+    )
+    source_provider: Mapped[str] = mapped_column(
+        Text, nullable=False, default="manual", server_default="manual", index=True
+    )
+    source_id: Mapped[str | None] = mapped_column(Text, nullable=True, index=True)
+    source_path: Mapped[str | None] = mapped_column(Text, nullable=True)
+    source_created_at: Mapped[datetime | None] = mapped_column(
+        DateTime(timezone=True), nullable=True, index=True
+    )
+    ingested_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), server_default=func.now(), nullable=False, index=True
+    )
+    parser_version: Mapped[str | None] = mapped_column(Text, nullable=True)
+    confidence: Mapped[float] = mapped_column(
+        Float, nullable=False, default=0.35, server_default="0.35"
+    )
+    cached_input_tokens: Mapped[int] = mapped_column(
+        Integer, nullable=False, default=0, server_default="0"
+    )
+    reasoning_output_tokens: Mapped[int] = mapped_column(
+        Integer, nullable=False, default=0, server_default="0"
+    )
+    raw_total_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True)
+    cost_estimated_usd: Mapped[float | None] = mapped_column(Float, nullable=True)
+    raw_metadata: Mapped[dict[str, Any]] = mapped_column(
+        JSONB, nullable=False, default=dict, server_default="{}"
+    )
     created_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True), server_default=func.now(), nullable=False, index=True
     )