7ef116e2f9
Define User, Project, Experiment, Run, StageResult, Score, ResponseCache, and WebhookConfig with UUID primary keys, JSON columns, enum types (ExperimentStatus, RunStatus), full relationship cascades, and indexes. Uses sqlalchemy.JSON (not JSONB) for SQLite compatibility in single-container mode. 16 tests added covering table creation, CRUD, uniqueness constraints, default values, and cascade deletes — all passing.
276 lines
8.9 KiB
Python
276 lines
8.9 KiB
Python
"""PromptLooper SQLAlchemy ORM models."""
|
|
|
|
import enum
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
|
|
from sqlalchemy import (
|
|
JSON,
|
|
Boolean,
|
|
DateTime,
|
|
Enum,
|
|
Float,
|
|
ForeignKey,
|
|
Index,
|
|
Integer,
|
|
Numeric,
|
|
String,
|
|
Text,
|
|
)
|
|
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
|
|
|
|
|
|
def _utcnow() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
def _new_uuid() -> uuid.UUID:
|
|
return uuid.uuid4()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Base
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class Base(DeclarativeBase):
|
|
"""Shared declarative base for all models."""
|
|
|
|
type_annotation_map = {
|
|
dict: JSON,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Enums
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class ExperimentStatus(str, enum.Enum):
|
|
draft = "draft"
|
|
running = "running"
|
|
paused = "paused"
|
|
completed = "completed"
|
|
|
|
|
|
class RunStatus(str, enum.Enum):
|
|
pending = "pending"
|
|
running = "running"
|
|
completed = "completed"
|
|
failed = "failed"
|
|
cached = "cached"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Models
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class User(Base):
|
|
__tablename__ = "users"
|
|
|
|
id: Mapped[uuid.UUID] = mapped_column(
|
|
primary_key=True, default=_new_uuid
|
|
)
|
|
username: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
|
|
password_hash: Mapped[str] = mapped_column(String(255), nullable=False)
|
|
is_admin: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), default=_utcnow, nullable=False
|
|
)
|
|
|
|
# Relationships
|
|
projects: Mapped[list["Project"]] = relationship(
|
|
back_populates="owner", cascade="all, delete-orphan"
|
|
)
|
|
|
|
|
|
class Project(Base):
|
|
__tablename__ = "projects"
|
|
|
|
id: Mapped[uuid.UUID] = mapped_column(
|
|
primary_key=True, default=_new_uuid
|
|
)
|
|
name: Mapped[str] = mapped_column(String(255), nullable=False)
|
|
description: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
owner_id: Mapped[uuid.UUID] = mapped_column(
|
|
ForeignKey("users.id", ondelete="CASCADE"), nullable=False
|
|
)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), default=_utcnow, nullable=False
|
|
)
|
|
updated_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), default=_utcnow, onupdate=_utcnow, nullable=False
|
|
)
|
|
|
|
# Relationships
|
|
owner: Mapped["User"] = relationship(back_populates="projects")
|
|
experiments: Mapped[list["Experiment"]] = relationship(
|
|
back_populates="project", cascade="all, delete-orphan"
|
|
)
|
|
|
|
|
|
class Experiment(Base):
|
|
__tablename__ = "experiments"
|
|
|
|
id: Mapped[uuid.UUID] = mapped_column(
|
|
primary_key=True, default=_new_uuid
|
|
)
|
|
project_id: Mapped[uuid.UUID] = mapped_column(
|
|
ForeignKey("projects.id", ondelete="CASCADE"), nullable=False
|
|
)
|
|
name: Mapped[str] = mapped_column(String(255), nullable=False)
|
|
description: Mapped[str | None] = mapped_column(Text, nullable=True)
|
|
sample_data: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
|
pipeline_stages: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
|
scoring_config: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
|
parameter_space: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
|
status: Mapped[ExperimentStatus] = mapped_column(
|
|
Enum(ExperimentStatus, name="experiment_status"),
|
|
default=ExperimentStatus.draft,
|
|
nullable=False,
|
|
)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), default=_utcnow, nullable=False
|
|
)
|
|
updated_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), default=_utcnow, onupdate=_utcnow, nullable=False
|
|
)
|
|
|
|
# Relationships
|
|
project: Mapped["Project"] = relationship(back_populates="experiments")
|
|
runs: Mapped[list["Run"]] = relationship(
|
|
back_populates="experiment", cascade="all, delete-orphan"
|
|
)
|
|
|
|
__table_args__ = (
|
|
Index("ix_experiments_project_id", "project_id"),
|
|
Index("ix_experiments_status", "status"),
|
|
)
|
|
|
|
|
|
class Run(Base):
|
|
__tablename__ = "runs"
|
|
|
|
id: Mapped[uuid.UUID] = mapped_column(
|
|
primary_key=True, default=_new_uuid
|
|
)
|
|
experiment_id: Mapped[uuid.UUID] = mapped_column(
|
|
ForeignKey("experiments.id", ondelete="CASCADE"), nullable=False
|
|
)
|
|
config_hash: Mapped[str] = mapped_column(String(64), nullable=False)
|
|
config: Mapped[dict] = mapped_column(JSON, nullable=False)
|
|
status: Mapped[RunStatus] = mapped_column(
|
|
Enum(RunStatus, name="run_status"),
|
|
default=RunStatus.pending,
|
|
nullable=False,
|
|
)
|
|
started_at: Mapped[datetime | None] = mapped_column(
|
|
DateTime(timezone=True), nullable=True
|
|
)
|
|
completed_at: Mapped[datetime | None] = mapped_column(
|
|
DateTime(timezone=True), nullable=True
|
|
)
|
|
duration_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
tokens_in: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
tokens_out: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
cost_estimate: Mapped[float | None] = mapped_column(
|
|
Numeric(precision=12, scale=6), nullable=True
|
|
)
|
|
|
|
# Relationships
|
|
experiment: Mapped["Experiment"] = relationship(back_populates="runs")
|
|
stage_results: Mapped[list["StageResult"]] = relationship(
|
|
back_populates="run", cascade="all, delete-orphan"
|
|
)
|
|
scores: Mapped[list["Score"]] = relationship(
|
|
back_populates="run", cascade="all, delete-orphan"
|
|
)
|
|
|
|
__table_args__ = (
|
|
Index("ix_runs_experiment_id", "experiment_id"),
|
|
Index("ix_runs_config_hash", "config_hash"),
|
|
Index("ix_runs_status", "status"),
|
|
)
|
|
|
|
|
|
class StageResult(Base):
|
|
__tablename__ = "stage_results"
|
|
|
|
id: Mapped[uuid.UUID] = mapped_column(
|
|
primary_key=True, default=_new_uuid
|
|
)
|
|
run_id: Mapped[uuid.UUID] = mapped_column(
|
|
ForeignKey("runs.id", ondelete="CASCADE"), nullable=False
|
|
)
|
|
stage_index: Mapped[int] = mapped_column(Integer, nullable=False)
|
|
prompt_sent: Mapped[str] = mapped_column(Text, nullable=False)
|
|
response_raw: Mapped[str] = mapped_column(Text, nullable=False)
|
|
model_used: Mapped[str] = mapped_column(String(255), nullable=False)
|
|
parameters: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
|
tokens_in: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
tokens_out: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
latency_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
|
|
# Relationships
|
|
run: Mapped["Run"] = relationship(back_populates="stage_results")
|
|
|
|
__table_args__ = (
|
|
Index("ix_stage_results_run_id", "run_id"),
|
|
)
|
|
|
|
|
|
class Score(Base):
|
|
__tablename__ = "scores"
|
|
|
|
id: Mapped[uuid.UUID] = mapped_column(
|
|
primary_key=True, default=_new_uuid
|
|
)
|
|
run_id: Mapped[uuid.UUID] = mapped_column(
|
|
ForeignKey("runs.id", ondelete="CASCADE"), nullable=False
|
|
)
|
|
scorer_name: Mapped[str] = mapped_column(String(255), nullable=False)
|
|
value: Mapped[float] = mapped_column(Float, nullable=False)
|
|
scorer_metadata: Mapped[dict | None] = mapped_column(
|
|
"metadata", JSON, nullable=True
|
|
)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), default=_utcnow, nullable=False
|
|
)
|
|
|
|
# Relationships
|
|
run: Mapped["Run"] = relationship(back_populates="scores")
|
|
|
|
__table_args__ = (
|
|
Index("ix_scores_run_id", "run_id"),
|
|
Index("ix_scores_scorer_name", "scorer_name"),
|
|
)
|
|
|
|
|
|
class ResponseCache(Base):
|
|
__tablename__ = "response_cache"
|
|
|
|
config_hash: Mapped[str] = mapped_column(
|
|
String(64), primary_key=True
|
|
)
|
|
response: Mapped[str] = mapped_column(Text, nullable=False)
|
|
model: Mapped[str] = mapped_column(String(255), nullable=False)
|
|
tokens_in: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
tokens_out: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
latency_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), default=_utcnow, nullable=False
|
|
)
|
|
|
|
|
|
class WebhookConfig(Base):
|
|
__tablename__ = "webhook_configs"
|
|
|
|
id: Mapped[uuid.UUID] = mapped_column(
|
|
primary_key=True, default=_new_uuid
|
|
)
|
|
event_type: Mapped[str] = mapped_column(String(255), nullable=False)
|
|
url: Mapped[str] = mapped_column(String(2048), nullable=False)
|
|
headers: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
|
is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
|
|
|
|
__table_args__ = (
|
|
Index("ix_webhook_configs_event_type", "event_type"),
|
|
)
|