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 SelfCorrelation(Base):
    __tablename__ = "self_correlations"
    alpha_id: Mapped[str] = mapped_column(ForeignKey("alphas.id"), primary_key=True)
    region: Mapped[str | None] = mapped_column(String(50), index=True)
    result: Mapped[dict] = mapped_column(JSON)
    stale: Mapped[bool] = mapped_column(Boolean, default=False)
    calculated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now)


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"),)