worldquant-alpha-system/backend/app/models.py

"""Platform snapshots and local research are intentionally stored separately."""

from datetime import datetime, timezone

from sqlalchemy import (
    JSON,
    Boolean,
    DateTime,
    Float,
    ForeignKey,
    Index,
    Integer,
    String,
    Text,
    UniqueConstraint,
)
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column


def now() -> datetime:
    return datetime.now(timezone.utc)


class Base(DeclarativeBase):
    pass


class Admin(Base):
    __tablename__ = "admins"
    id: Mapped[int] = mapped_column(primary_key=True, default=1)
    username: Mapped[str] = mapped_column(String(100), unique=True)
    password_hash: Mapped[str] = mapped_column(Text)


class LoginSession(Base):
    __tablename__ = "login_sessions"
    token_hash: Mapped[str] = mapped_column(String(64), primary_key=True)
    expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), index=True)


class Account(Base):
    __tablename__ = "accounts"
    id: Mapped[int] = mapped_column(primary_key=True, default=1)
    email: Mapped[str | None] = mapped_column(String(254))
    password_encrypted: Mapped[str | None] = mapped_column(Text)
    wq_user_id: Mapped[str | None] = mapped_column(String(100))
    profile: Mapped[dict] = mapped_column(JSON, default=dict)
    connection_status: Mapped[str] = mapped_column(String(30), default="disconnected")
    connection_error: Mapped[str | None] = mapped_column(Text)
    verification_url: Mapped[str | None] = mapped_column(Text)
    last_synced_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    display_name: Mapped[str] = mapped_column(String(100), default="研究员")
    theme: Mapped[str] = mapped_column(String(10), default="light")
    timezone: Mapped[str] = mapped_column(String(64), default="Asia/Shanghai")
    page_size: Mapped[int] = mapped_column(Integer, default=25)


class Alpha(Base):
    __tablename__ = "alphas"
    id: Mapped[str] = mapped_column(String(100), primary_key=True)
    name: Mapped[str | None] = mapped_column(Text)
    expression: Mapped[str | None] = mapped_column(Text)
    selection: Mapped[str | None] = mapped_column(Text)
    combo: Mapped[str | None] = mapped_column(Text)
    alpha_type: Mapped[str | None] = mapped_column(String(50), index=True)
    language: Mapped[str | None] = mapped_column(String(50))
    stage: Mapped[str | None] = mapped_column(String(50), index=True)
    status: Mapped[str | None] = mapped_column(String(80), index=True)
    hidden: Mapped[bool] = mapped_column(Boolean, default=False)
    region: Mapped[str | None] = mapped_column(String(50), index=True)
    universe: Mapped[str | None] = mapped_column(String(100))
    settings: Mapped[dict] = mapped_column(JSON, default=dict)
    is_metrics: Mapped[dict] = mapped_column(JSON, default=dict)
    os_metrics: Mapped[dict] = mapped_column(JSON, default=dict)
    checks: Mapped[list] = mapped_column(JSON, default=list)
    sharpe: Mapped[float | None] = mapped_column(Float, index=True)
    fitness: Mapped[float | None] = mapped_column(Float, index=True)
    returns: Mapped[float | None] = mapped_column(Float)
    turnover: Mapped[float | None] = mapped_column(Float)
    margin: Mapped[float | None] = mapped_column(Float)
    drawdown: Mapped[float | None] = mapped_column(Float)
    date_created: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), index=True)
    date_submitted: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    synced_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
    raw: Mapped[dict] = mapped_column(JSON, default=dict)
    __table_args__ = (Index("ix_alphas_region_status", "region", "status"),)


class Research(Base):
    __tablename__ = "research"
    alpha_id: Mapped[str] = mapped_column(ForeignKey("alphas.id"), primary_key=True)
    note: Mapped[str] = mapped_column(Text, default="")
    tags: Mapped[list] = mapped_column(JSON, default=list)
    favorite: Mapped[bool] = mapped_column(Boolean, default=False)
    state: Mapped[str] = mapped_column(String(30), default="inbox", index=True)
    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
    version: Mapped[int] = mapped_column(Integer, default=1, server_default="1")


class Pnl(Base):
    __tablename__ = "pnl_cache"
    alpha_id: Mapped[str] = mapped_column(ForeignKey("alphas.id"), primary_key=True)
    raw: Mapped[dict] = mapped_column(JSON)
    points: Mapped[list] = mapped_column(JSON)
    fetched_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)


class ResearchTag(Base):
    __tablename__ = "research_tags"
    alpha_id: Mapped[str] = mapped_column(ForeignKey("alphas.id"), primary_key=True)
    tag: Mapped[str] = mapped_column(String(60), primary_key=True, index=True)


class Job(Base):
    __tablename__ = "sync_jobs"
    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    kind: Mapped[str] = mapped_column(String(30))
    status: Mapped[str] = mapped_column(String(30), default="queued", index=True)
    payload: Mapped[dict] = mapped_column(JSON, default=dict)
    checkpoint: Mapped[dict] = mapped_column(JSON, default=dict)
    processed: Mapped[int] = mapped_column(Integer, default=0)
    failed: Mapped[int] = mapped_column(Integer, default=0)
    total: Mapped[int | None] = mapped_column(Integer)
    error: Mapped[str | None] = mapped_column(Text)
    next_retry_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    cancel_requested: Mapped[bool] = mapped_column(Boolean, default=False)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)


class JobItem(Base):
    __tablename__ = "job_items"
    job_id: Mapped[str] = mapped_column(ForeignKey("sync_jobs.id"), primary_key=True)
    alpha_id: Mapped[str] = mapped_column(String(100), primary_key=True)
    error: Mapped[str | None] = mapped_column(Text)


class AISettings(Base):
    __tablename__ = "ai_settings"
    id: Mapped[int] = mapped_column(primary_key=True, default=1)
    base_url: Mapped[str] = mapped_column(Text, default="")
    api_key_encrypted: Mapped[str | None] = mapped_column(Text)
    model: Mapped[str] = mapped_column(String(200), default="")
    protocol: Mapped[str] = mapped_column(String(30), default="chat_completions")
    enabled: Mapped[bool] = mapped_column(Boolean, default=False)
    revision: Mapped[int] = mapped_column(Integer, default=1)
    tested_revision: Mapped[int | None] = mapped_column(Integer)
    test_results: Mapped[dict] = mapped_column(JSON, default=dict)


class AIConversation(Base):
    __tablename__ = "ai_conversations"
    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    admin_id: Mapped[int] = mapped_column(ForeignKey("admins.id"), default=1)
    title: Mapped[str] = mapped_column(String(100), default="新会话")
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)


class AIRun(Base):
    __tablename__ = "ai_runs"
    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    conversation_id: Mapped[str] = mapped_column(ForeignKey("ai_conversations.id"), index=True)
    request_id: Mapped[str] = mapped_column(String(100))
    status: Mapped[str] = mapped_column(String(30), default="running")
    context: Mapped[dict] = mapped_column(JSON, default=dict)
    model: Mapped[str] = mapped_column(String(200))
    settings_revision: Mapped[int] = mapped_column(Integer)
    model_messages: Mapped[list] = mapped_column(JSON, default=list)
    history_start: Mapped[int] = mapped_column(Integer, default=0)
    usage: Mapped[dict] = mapped_column(JSON, default=dict)
    elapsed_ms: Mapped[int] = mapped_column(Integer, default=0)
    error: Mapped[str | None] = mapped_column(Text)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
    __table_args__ = (UniqueConstraint("conversation_id", "request_id"),)


class AIMessage(Base):
    __tablename__ = "ai_messages"
    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    conversation_id: Mapped[str] = mapped_column(ForeignKey("ai_conversations.id"), index=True)
    run_id: Mapped[str] = mapped_column(ForeignKey("ai_runs.id"), index=True)
    role: Mapped[str] = mapped_column(String(20))
    parts: Mapped[list] = mapped_column(JSON, default=list)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)


class AIToolCall(Base):
    __tablename__ = "ai_tool_calls"
    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    run_id: Mapped[str] = mapped_column(ForeignKey("ai_runs.id"), index=True)
    admin_id: Mapped[int] = mapped_column(ForeignKey("admins.id"), default=1)
    call_id: Mapped[str] = mapped_column(String(200))
    name: Mapped[str] = mapped_column(String(100))
    arguments: Mapped[dict] = mapped_column(JSON)
    preview: Mapped[dict] = mapped_column(JSON, default=dict)
    result: Mapped[dict | None] = mapped_column(JSON)
    status: Mapped[str] = mapped_column(String(30), default="pending")
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
    __table_args__ = (UniqueConstraint("run_id", "call_id"),)


class BacktestConfig(Base):
    __tablename__ = "backtest_config"
    id: Mapped[int] = mapped_column(primary_key=True, default=1)
    concurrency: Mapped[int] = mapped_column(Integer, default=3)
    batch_size: Mapped[int] = mapped_column(Integer, default=8)
    version: Mapped[int] = mapped_column(Integer, default=1)
    blocked_reason: Mapped[str | None] = mapped_column(Text)
    blocked_until: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))


class BacktestDraft(Base):
    __tablename__ = "backtest_drafts"
    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    version: Mapped[int] = mapped_column(Integer, default=1)
    name: Mapped[str] = mapped_column(String(200))
    source: Mapped[dict] = mapped_column(JSON)
    candidates: Mapped[list] = mapped_column(JSON)
    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)


class BacktestPreview(Base):
    __tablename__ = "backtest_previews"
    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    version: Mapped[int] = mapped_column(Integer, default=1)
    name: Mapped[str] = mapped_column(String(200))
    source: Mapped[dict] = mapped_column(JSON)
    candidates: Mapped[list] = mapped_column(JSON)
    batches: Mapped[list] = mapped_column(JSON)
    batch_size: Mapped[int] = mapped_column(Integer)
    digest: Mapped[str] = mapped_column(String(64))
    duplicates: Mapped[list] = mapped_column(JSON)
    ai_context: Mapped[dict] = mapped_column(JSON, default=dict)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)


class BacktestRun(Base):
    __tablename__ = "backtest_runs"
    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    preview_id: Mapped[str] = mapped_column(ForeignKey("backtest_previews.id"), unique=True)
    idempotency_key: Mapped[str] = mapped_column(String(100), unique=True)
    name: Mapped[str] = mapped_column(String(200))
    source: Mapped[dict] = mapped_column(JSON)
    ai_context: Mapped[dict] = mapped_column(JSON, default=dict)
    control: Mapped[str] = mapped_column(String(20), default="active")
    status: Mapped[str] = mapped_column(String(30), default="queued", index=True)
    version: Mapped[int] = mapped_column(Integer, default=1)
    event_seq: Mapped[int] = mapped_column(Integer, default=0)
    total: Mapped[int] = mapped_column(Integer)
    batch_size: Mapped[int] = mapped_column(Integer)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)


class SimulationAttempt(Base):
    __tablename__ = "simulation_attempts"
    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    run_id: Mapped[str] = mapped_column(ForeignKey("backtest_runs.id"), index=True)
    ordinal: Mapped[int] = mapped_column(Integer)
    state: Mapped[str] = mapped_column(String(30), default="queued", index=True)
    payload: Mapped[list] = mapped_column(JSON)
    progress_url: Mapped[str | None] = mapped_column(Text)
    remote_complete: Mapped[bool] = mapped_column(Boolean, default=False)
    children: Mapped[list] = mapped_column(JSON, default=list)
    receipts: Mapped[dict] = mapped_column(JSON, default=dict)
    poll_count: Mapped[int] = mapped_column(Integer, default=0)
    submit_count: Mapped[int] = mapped_column(Integer, default=0)
    next_poll_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    error: Mapped[str | None] = mapped_column(Text)
    error_code: Mapped[str | None] = mapped_column(String(50))
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
    __table_args__ = (UniqueConstraint("run_id", "ordinal"),)


class BacktestItem(Base):
    __tablename__ = "backtest_items"
    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    run_id: Mapped[str] = mapped_column(ForeignKey("backtest_runs.id"), index=True)
    attempt_id: Mapped[str] = mapped_column(ForeignKey("simulation_attempts.id"), index=True)
    client_item_id: Mapped[str] = mapped_column(String(100))
    ordinal: Mapped[int] = mapped_column(Integer)
    expression: Mapped[str] = mapped_column(Text)
    settings: Mapped[dict] = mapped_column(JSON)
    fingerprint: Mapped[str] = mapped_column(String(64), index=True)
    platform_status: Mapped[str] = mapped_column(String(30), default="pending")
    collection_status: Mapped[str] = mapped_column(String(30), default="pending")
    persistence_status: Mapped[str] = mapped_column(String(30), default="pending")
    simulation_id: Mapped[str | None] = mapped_column(String(100))
    alpha_id: Mapped[str | None] = mapped_column(String(100))
    error: Mapped[str | None] = mapped_column(Text)
    __table_args__ = (UniqueConstraint("run_id", "client_item_id"),)


class BacktestResult(Base):
    __tablename__ = "backtest_results"
    item_id: Mapped[str] = mapped_column(ForeignKey("backtest_items.id"), primary_key=True)
    attempt_id: Mapped[str] = mapped_column(ForeignKey("simulation_attempts.id"))
    alpha_id: Mapped[str] = mapped_column(ForeignKey("alphas.id"), index=True)
    snapshot: Mapped[dict] = mapped_column(JSON)
    observed_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)
    complete: Mapped[bool] = mapped_column(Boolean, default=True)


class BacktestEvent(Base):
    __tablename__ = "backtest_events"
    run_id: Mapped[str] = mapped_column(ForeignKey("backtest_runs.id"), primary_key=True)
    seq: Mapped[int] = mapped_column(Integer, primary_key=True)
    kind: Mapped[str] = mapped_column(String(50))
    payload: Mapped[dict] = mapped_column(JSON)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)


class CatalogScope(Base):
    __tablename__ = "catalog_scopes"
    key: Mapped[str] = mapped_column(String(200), primary_key=True)
    scope: Mapped[dict] = mapped_column(JSON)
    catalog_version: Mapped[str | None] = mapped_column(String(36))
    synced_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))


class CatalogBatch(Base):
    __tablename__ = "catalog_batches"
    id: Mapped[str] = mapped_column(ForeignKey("sync_jobs.id"), primary_key=True)
    scope_key: Mapped[str] = mapped_column(ForeignKey("catalog_scopes.key"), index=True)
    dataset_id: Mapped[str | None] = mapped_column(String(200))
    complete: Mapped[bool] = mapped_column(Boolean, default=False)
    count: Mapped[int] = mapped_column(Integer, default=0)
    completed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))


class CatalogDataset(Base):
    __tablename__ = "catalog_datasets"
    scope_key: Mapped[str] = mapped_column(ForeignKey("catalog_scopes.key"), primary_key=True)
    id: Mapped[str] = mapped_column(String(200), primary_key=True)
    field_version: Mapped[str | None] = mapped_column(ForeignKey("catalog_batches.id"))


class CatalogEntry(Base):
    """Immutable published snapshots; staging rows remain invisible until batch completion."""
    __tablename__ = "catalog_entries"
    batch_id: Mapped[str] = mapped_column(ForeignKey("catalog_batches.id"), primary_key=True)
    id: Mapped[str] = mapped_column(String(200), primary_key=True)
    name: Mapped[str | None] = mapped_column(Text)
    category: Mapped[str | None] = mapped_column(String(200))
    subcategory: Mapped[str | None] = mapped_column(String(200))
    field_type: Mapped[str | None] = mapped_column(String(100))
    coverage: Mapped[float | None] = mapped_column(Float)
    user_count: Mapped[int | None] = mapped_column(Integer)
    alpha_count: Mapped[int | None] = mapped_column(Integer)
    field_count: Mapped[int | None] = mapped_column(Integer)
    description: Mapped[str | None] = mapped_column(Text)
    unit: Mapped[str | None] = mapped_column(Text)
    synced_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)


class CatalogNote(Base):
    __tablename__ = "catalog_notes"
    scope_key: Mapped[str] = mapped_column(ForeignKey("catalog_scopes.key"), primary_key=True)
    dataset_id: Mapped[str] = mapped_column(String(200), primary_key=True)
    # Empty field_id denotes the dataset; platform identifiers cannot be empty.
    field_id: Mapped[str] = mapped_column(String(200), primary_key=True, default="")
    note: Mapped[str] = mapped_column(Text, default="")
    version: Mapped[int] = mapped_column(Integer, default=1)
    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)


class TemplateInput(Base):
    __tablename__ = "template_inputs"
    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    scope_key: Mapped[str] = mapped_column(ForeignKey("catalog_scopes.key"), index=True)
    dataset_id: Mapped[str] = mapped_column(String(200))
    collection_version: Mapped[str] = mapped_column(ForeignKey("catalog_batches.id"))
    selection: Mapped[str] = mapped_column(String(20))
    field_ids: Mapped[list] = mapped_column(JSON)
    field_types: Mapped[dict] = mapped_column(JSON)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)