Spaces:
Running
Running
File size: 5,919 Bytes
69e310f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """SQLAlchemy tables for runs, briefs, approvals and archived telemetry.
Runs on Neon Postgres in production and SQLite locally; both back the same
schema, including the partial unique index that enforces **one active run per
watchlist at the database level** rather than in application code.
"""
from __future__ import annotations
from datetime import UTC, datetime
from typing import Any
from sqlalchemy import (
JSON,
DateTime,
Float,
ForeignKey,
Index,
Integer,
String,
Text,
text,
)
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
from app.models.run import ACTIVE_STATUSES
#: SQL fragment listing the statuses that occupy the single active-run slot.
_ACTIVE_SQL = ", ".join(f"'{status}'" for status in sorted(ACTIVE_STATUSES))
ACTIVE_RUN_PREDICATE = text(f"status IN ({_ACTIVE_SQL})")
def _utcnow() -> datetime:
return datetime.now(UTC)
class Base(DeclarativeBase):
"""Declarative base."""
class RunRecord(Base):
"""One brief run, from trigger to terminal status."""
__tablename__ = "runs"
id: Mapped[str] = mapped_column(String(64), primary_key=True)
watchlist_key: Mapped[str] = mapped_column(String(255), index=True)
tickers: Mapped[list[str]] = mapped_column(JSON, default=list)
mode: Mapped[str] = mapped_column(String(32), default="standard")
status: Mapped[str] = mapped_column(String(32), index=True)
engine: Mapped[str] = mapped_column(String(32), default="deterministic")
trigger: Mapped[str] = mapped_column(String(32), default="ui")
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow)
updated_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), default=_utcnow, onupdate=_utcnow
)
finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), default=None)
iterations: Mapped[int] = mapped_column(Integer, default=0)
model_calls: Mapped[int] = mapped_column(Integer, default=0)
input_tokens: Mapped[int] = mapped_column(Integer, default=0)
output_tokens: Mapped[int] = mapped_column(Integer, default=0)
cost_usd: Mapped[float] = mapped_column(Float, default=0.0)
latency_ms: Mapped[float] = mapped_column(Float, default=0.0)
tool_calls: Mapped[int] = mapped_column(Integer, default=0)
partial: Mapped[bool] = mapped_column(default=False)
verified: Mapped[bool] = mapped_column(default=False)
abort_reason: Mapped[str | None] = mapped_column(Text, default=None)
error_count: Mapped[int] = mapped_column(Integer, default=0)
briefs: Mapped[list[BriefRecord]] = relationship(
back_populates="run", cascade="all, delete-orphan"
)
approvals: Mapped[list[ApprovalRecord]] = relationship(
back_populates="run", cascade="all, delete-orphan"
)
__table_args__ = (
# One active run per watchlist — enforced by the database, not by a
# read-check-write race in the API layer. Partial unique indexes are
# supported by both PostgreSQL and SQLite.
Index(
"uq_runs_one_active_per_watchlist",
"watchlist_key",
unique=True,
postgresql_where=ACTIVE_RUN_PREDICATE,
sqlite_where=ACTIVE_RUN_PREDICATE,
),
Index("ix_runs_created_at", "created_at"),
)
class BriefRecord(Base):
"""The archived brief and its verification report."""
__tablename__ = "briefs"
id: Mapped[str] = mapped_column(String(64), primary_key=True)
run_id: Mapped[str] = mapped_column(
String(64), ForeignKey("runs.id", ondelete="CASCADE"), index=True
)
generated_for: Mapped[str] = mapped_column(String(16))
headline: Mapped[str] = mapped_column(Text)
partial: Mapped[bool] = mapped_column(default=False)
verified: Mapped[bool] = mapped_column(default=False)
claims_total: Mapped[int] = mapped_column(Integer, default=0)
claims_matched: Mapped[int] = mapped_column(Integer, default=0)
brief_json: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
verification_json: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
markdown: Mapped[str] = mapped_column(Text, default="")
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow)
run: Mapped[RunRecord] = relationship(back_populates="briefs")
class ApprovalRecord(Base):
"""The human decision that released (or blocked) a brief."""
__tablename__ = "approvals"
id: Mapped[str] = mapped_column(String(64), primary_key=True)
run_id: Mapped[str] = mapped_column(
String(64), ForeignKey("runs.id", ondelete="CASCADE"), index=True
)
action: Mapped[str] = mapped_column(String(16))
reviewer: Mapped[str] = mapped_column(String(128), default="analyst")
note: Mapped[str | None] = mapped_column(Text, default=None)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=_utcnow)
run: Mapped[RunRecord] = relationship(back_populates="approvals")
class EventRecord(Base):
"""Archived telemetry so a completed run can be replayed in the UI.
``(run_id, seq)`` is the natural key and therefore the *primary* key: archival
is retried whenever a run reaches a terminal state, so an upsert on the
natural key must be idempotent. With a surrogate id, ``session.merge()``
would insert duplicates and trip the unique index instead.
"""
__tablename__ = "run_events"
run_id: Mapped[str] = mapped_column(String(64), primary_key=True)
seq: Mapped[int] = mapped_column(Integer, primary_key=True)
ts: Mapped[str] = mapped_column(String(40))
kind: Mapped[str] = mapped_column(String(48))
message: Mapped[str] = mapped_column(Text)
payload: Mapped[dict[str, Any]] = mapped_column(JSON, default=dict)
__table_args__ = (Index("ix_run_events_run_id", "run_id"),)
|