"""SQLAlchemy ORM models for NEXUS.""" from __future__ import annotations import uuid from datetime import datetime from typing import Optional from sqlalchemy import ( JSON, DateTime, Float, ForeignKey, Integer, String, Text, UniqueConstraint, func, ) from sqlalchemy.dialects.postgresql import JSONB, UUID from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship def _uuid() -> str: return str(uuid.uuid4()) class Base(DeclarativeBase): pass class User(Base): __tablename__ = "users" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) email: Mapped[str] = mapped_column(String(255), unique=True, index=True, nullable=False) name: Mapped[Optional[str]] = mapped_column(String(255)) google_id: Mapped[Optional[str]] = mapped_column(String(255), unique=True, index=True) api_key_hash: Mapped[Optional[str]] = mapped_column(String(255)) plan: Mapped[str] = mapped_column(String(32), default="free", nullable=False) credits_used: Mapped[int] = mapped_column(Integer, default=0, nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), nullable=False ) tasks: Mapped[list["Task"]] = relationship(back_populates="user") class Task(Base): __tablename__ = "tasks" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) user_id: Mapped[Optional[str]] = mapped_column( String(36), ForeignKey("users.id", ondelete="SET NULL"), index=True ) topic: Mapped[str] = mapped_column(Text, nullable=False) slide_count: Mapped[int] = mapped_column(Integer, default=8, nullable=False) theme: Mapped[str] = mapped_column(String(64), default="Editorial", nullable=False) search_web: Mapped[bool] = mapped_column(default=True, nullable=False) status: Mapped[str] = mapped_column(String(32), default="pending", index=True, nullable=False) progress_pct: Mapped[float] = mapped_column(Float, default=0.0, nullable=False) current_step: Mapped[str] = mapped_column(String(255), default="", nullable=False) error_msg: Mapped[Optional[str]] = mapped_column(Text) tokens_used: Mapped[int] = mapped_column(Integer, default=0, nullable=False) cost_usd: Mapped[float] = mapped_column(Float, default=0.0, nullable=False) model_used: Mapped[Optional[str]] = mapped_column(String(64)) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), nullable=False ) completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True)) user: Mapped[Optional[User]] = relationship(back_populates="tasks") slides: Mapped[Optional["SlideDeck"]] = relationship( back_populates="task", uselist=False, cascade="all, delete-orphan" ) exports: Mapped[list["Export"]] = relationship( back_populates="task", cascade="all, delete-orphan" ) share_tokens: Mapped[list["ShareToken"]] = relationship( back_populates="task", cascade="all, delete-orphan" ) class SlideDeck(Base): __tablename__ = "slides" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) task_id: Mapped[str] = mapped_column( String(36), ForeignKey("tasks.id", ondelete="CASCADE"), unique=True, index=True ) slide_data: Mapped[list] = mapped_column(JSON().with_variant(JSONB, "postgresql"), nullable=False) theme: Mapped[str] = mapped_column(String(64), default="Editorial", nullable=False) slide_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), nullable=False ) task: Mapped[Task] = relationship(back_populates="slides") class Export(Base): __tablename__ = "exports" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) task_id: Mapped[str] = mapped_column( String(36), ForeignKey("tasks.id", ondelete="CASCADE"), index=True ) format: Mapped[str] = mapped_column(String(16), nullable=False) file_url: Mapped[str] = mapped_column(Text, nullable=False) file_size: Mapped[int] = mapped_column(Integer, default=0, nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), nullable=False ) task: Mapped[Task] = relationship(back_populates="exports") class ShareToken(Base): __tablename__ = "share_tokens" __table_args__ = (UniqueConstraint("token", name="uq_share_token"),) token: Mapped[str] = mapped_column(String(64), primary_key=True) task_id: Mapped[str] = mapped_column( String(36), ForeignKey("tasks.id", ondelete="CASCADE"), index=True ) views: Mapped[int] = mapped_column(Integer, default=0, nullable=False) expires_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True)) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), nullable=False ) task: Mapped[Task] = relationship(back_populates="share_tokens") # ── Agent runtime (Phase 1 foundation) ───────────────────────────────────── # These tables back the dynamic tool-calling loop introduced in Phase 2. # They are additive: existing slide-generation flow does not use them yet. # Tables are created automatically by ``init_models`` (dev/SQLite) and via # Alembic in production. class AgentRun(Base): __tablename__ = "agent_runs" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) task_id: Mapped[Optional[str]] = mapped_column( String(36), ForeignKey("tasks.id", ondelete="CASCADE"), index=True ) user_id: Mapped[Optional[str]] = mapped_column( String(36), ForeignKey("users.id", ondelete="SET NULL"), index=True ) goal: Mapped[str] = mapped_column(Text, nullable=False) status: Mapped[str] = mapped_column( String(32), default="pending", index=True, nullable=False ) max_steps: Mapped[int] = mapped_column(Integer, default=20, nullable=False) step_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False) error_msg: Mapped[Optional[str]] = mapped_column(Text) meta: Mapped[dict] = mapped_column( JSON().with_variant(JSONB, "postgresql"), default=dict, nullable=False ) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), nullable=False ) completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True)) steps: Mapped[list["AgentStep"]] = relationship( back_populates="run", cascade="all, delete-orphan", order_by="AgentStep.step_index" ) artifacts: Mapped[list["Artifact"]] = relationship( back_populates="run", cascade="all, delete-orphan" ) class AgentStep(Base): __tablename__ = "agent_steps" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) run_id: Mapped[str] = mapped_column( String(36), ForeignKey("agent_runs.id", ondelete="CASCADE"), index=True, nullable=False ) step_index: Mapped[int] = mapped_column(Integer, nullable=False) kind: Mapped[str] = mapped_column(String(32), nullable=False) # "thought" | "action" | "observation" | "final" action: Mapped[Optional[str]] = mapped_column(String(64)) # tool name when kind=="action" status: Mapped[str] = mapped_column(String(32), default="ok", nullable=False) input_json: Mapped[dict] = mapped_column( JSON().with_variant(JSONB, "postgresql"), default=dict, nullable=False ) output_json: Mapped[dict] = mapped_column( JSON().with_variant(JSONB, "postgresql"), default=dict, nullable=False ) error: Mapped[Optional[str]] = mapped_column(Text) started_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), nullable=False ) completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True)) run: Mapped[AgentRun] = relationship(back_populates="steps") class Artifact(Base): __tablename__ = "artifacts" id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid) run_id: Mapped[Optional[str]] = mapped_column( String(36), ForeignKey("agent_runs.id", ondelete="CASCADE"), index=True ) task_id: Mapped[Optional[str]] = mapped_column( String(36), ForeignKey("tasks.id", ondelete="CASCADE"), index=True ) artifact_type: Mapped[str] = mapped_column(String(64), index=True, nullable=False) title: Mapped[Optional[str]] = mapped_column(String(255)) meta: Mapped[dict] = mapped_column( JSON().with_variant(JSONB, "postgresql"), default=dict, nullable=False ) file_url: Mapped[Optional[str]] = mapped_column(Text) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), server_default=func.now(), nullable=False ) run: Mapped[Optional[AgentRun]] = relationship(back_populates="artifacts")