promptlooper/backend/schemas.py

"""PromptLooper Pydantic request/response schemas."""

import uuid
from datetime import datetime

from pydantic import BaseModel, ConfigDict, Field

from models import ExperimentStatus, RunStatus


# ---------------------------------------------------------------------------
# Shared mixins
# ---------------------------------------------------------------------------

class _TimestampMixin(BaseModel):
    created_at: datetime
    updated_at: datetime


# ---------------------------------------------------------------------------
# Project
# ---------------------------------------------------------------------------

class ProjectCreate(BaseModel):
    name: str = Field(..., min_length=1, max_length=255)
    description: str | None = None


class ProjectUpdate(BaseModel):
    name: str | None = Field(None, min_length=1, max_length=255)
    description: str | None = None


class ProjectResponse(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: uuid.UUID
    name: str
    description: str | None
    owner_id: uuid.UUID
    created_at: datetime
    updated_at: datetime


class ProjectListResponse(BaseModel):
    items: list[ProjectResponse]
    total: int


# ---------------------------------------------------------------------------
# Experiment
# ---------------------------------------------------------------------------

class ExperimentCreate(BaseModel):
    name: str = Field(..., min_length=1, max_length=255)
    description: str | None = None
    sample_data: dict | None = None
    pipeline_stages: dict | None = None
    scoring_config: dict | None = None
    parameter_space: dict | None = None


class ExperimentUpdate(BaseModel):
    name: str | None = Field(None, min_length=1, max_length=255)
    description: str | None = None
    sample_data: dict | None = None
    pipeline_stages: dict | None = None
    scoring_config: dict | None = None
    parameter_space: dict | None = None
    status: ExperimentStatus | None = None


class ExperimentResponse(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: uuid.UUID
    project_id: uuid.UUID
    name: str
    description: str | None
    sample_data: dict | None
    pipeline_stages: dict | None
    scoring_config: dict | None
    parameter_space: dict | None
    status: ExperimentStatus
    created_at: datetime
    updated_at: datetime


class ExperimentListResponse(BaseModel):
    items: list[ExperimentResponse]
    total: int


# ---------------------------------------------------------------------------
# Run
# ---------------------------------------------------------------------------

class RunResponse(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: uuid.UUID
    experiment_id: uuid.UUID
    config_hash: str
    config: dict
    status: RunStatus
    started_at: datetime | None
    completed_at: datetime | None
    duration_ms: int | None
    tokens_in: int | None
    tokens_out: int | None
    cost_estimate: float | None


class RunListResponse(BaseModel):
    items: list[RunResponse]
    total: int


# ---------------------------------------------------------------------------
# StageResult (read-only, returned inside Run details)
# ---------------------------------------------------------------------------

class StageResultResponse(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: uuid.UUID
    run_id: uuid.UUID
    stage_index: int
    prompt_sent: str
    response_raw: str
    model_used: str
    parameters: dict | None
    tokens_in: int | None
    tokens_out: int | None
    latency_ms: int | None


class RunDetailResponse(RunResponse):
    """Run with nested stage results and scores."""

    stage_results: list[StageResultResponse] = []
    scores: list["ScoreResponse"] = []


# ---------------------------------------------------------------------------
# Score
# ---------------------------------------------------------------------------

class ScoreInput(BaseModel):
    scorer_name: str = Field(..., min_length=1, max_length=255)
    value: float
    metadata: dict | None = None


class ScoreResponse(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: uuid.UUID
    run_id: uuid.UUID
    scorer_name: str
    value: float
    scorer_metadata: dict | None
    created_at: datetime


# ---------------------------------------------------------------------------
# Endpoint (LLM endpoint configuration)
# ---------------------------------------------------------------------------

class EndpointCreate(BaseModel):
    name: str = Field(..., min_length=1, max_length=255)
    url: str = Field(..., min_length=1, max_length=2048)
    api_key: str | None = None
    default_model: str | None = Field(None, max_length=255)


class EndpointUpdate(BaseModel):
    name: str | None = Field(None, min_length=1, max_length=255)
    url: str | None = Field(None, min_length=1, max_length=2048)
    api_key: str | None = None
    default_model: str | None = Field(None, max_length=255)


class EndpointResponse(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: uuid.UUID
    name: str
    url: str
    default_model: str | None


class EndpointListResponse(BaseModel):
    items: list[EndpointResponse]
    total: int


# ---------------------------------------------------------------------------
# Webhook
# ---------------------------------------------------------------------------

class WebhookCreate(BaseModel):
    event_type: str = Field(..., min_length=1, max_length=255)
    url: str = Field(..., min_length=1, max_length=2048)
    headers: dict | None = None
    is_active: bool = True


class WebhookUpdate(BaseModel):
    event_type: str | None = Field(None, min_length=1, max_length=255)
    url: str | None = Field(None, min_length=1, max_length=2048)
    headers: dict | None = None
    is_active: bool | None = None


class WebhookResponse(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: uuid.UUID
    event_type: str
    url: str
    headers: dict | None
    is_active: bool


class WebhookListResponse(BaseModel):
    items: list[WebhookResponse]
    total: int


# ---------------------------------------------------------------------------
# Auth
# ---------------------------------------------------------------------------

class SetupRequest(BaseModel):
    username: str = Field(..., min_length=1, max_length=255)
    password: str = Field(..., min_length=8)


class LoginRequest(BaseModel):
    username: str
    password: str


class TokenResponse(BaseModel):
    access_token: str
    token_type: str = "bearer"


class UserResponse(BaseModel):
    model_config = ConfigDict(from_attributes=True)

    id: uuid.UUID
    username: str
    is_admin: bool
    created_at: datetime


# ---------------------------------------------------------------------------
# Export
# ---------------------------------------------------------------------------

class ExportRunRow(BaseModel):
    """Flat row for CSV/JSON export of run results."""

    run_id: uuid.UUID
    experiment_id: uuid.UUID
    config_hash: str
    config: dict
    status: RunStatus
    duration_ms: int | None = None
    tokens_in: int | None = None
    tokens_out: int | None = None
    cost_estimate: float | None = None
    scores: dict[str, float] = Field(
        default_factory=dict,
        description="Map of scorer_name → value",
    )


class ExportResponse(BaseModel):
    experiment_id: uuid.UUID
    experiment_name: str
    rows: list[ExportRunRow]


# ---------------------------------------------------------------------------
# Health
# ---------------------------------------------------------------------------

class HealthResponse(BaseModel):
    status: str = "ok"
    database: bool
    redis: bool


# Rebuild forward refs for RunDetailResponse
RunDetailResponse.model_rebuild()