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

"""Business operations shared by HTTP and AI; callers own transactions and authorization.

Mutations never commit here, so the AI executor can atomically save their audit result.
Job runner notifications must happen after commit, using ``notify_job``.
"""

from datetime import timezone
from uuid import uuid4

from fastapi import HTTPException
from sqlalchemy import delete, func, select, update

from .alphas import list_statement, sorted_statement, summary
from .jobs import ACTIVE
from .models import Account, Alpha, Job, JobItem, Pnl, Research, ResearchTag, SelfCorrelation, now
from .research.provenance import alpha_sources, source_kinds
from .schemas import AlphaDetail, AlphaPage, BulkUpdate, JobInput, JobOutput, ResearchUpdate, normalize_tags


class Business:
    def __init__(self, db, ai_context=None):
        from .backtests.service import Backtests
        from .catalog.service import Catalog
        from .research.service import ResearchBuilder

        self.db = db
        self.backtests = Backtests(db, ai_context)
        self.catalog = Catalog(db)
        self.research_builder = ResearchBuilder(db, self.backtests)

    async def search_alphas(self, filters):
        query = list_statement(filters)
        total = await self.db.scalar(select(func.count()).select_from(query.subquery()))
        rows = (
            await self.db.execute(
                sorted_statement(query, filters.sort, filters.direction)
                .limit(filters.limit)
                .offset(filters.offset)
            )
        ).all()
        correlations = {
            row.alpha_id: self.correlation_summary(row)
            for row in (
                await self.db.scalars(
                    select(SelfCorrelation).where(SelfCorrelation.alpha_id.in_([a.id for a, _ in rows]))
                )
            ).all()
        }
        sources = await source_kinds(self.db, [a.id for a, _ in rows])
        return AlphaPage(
            items=[
                {
                    **summary(a, r),
                    "local_correlation": correlations.get(a.id),
                    "source_kinds": sources.get(a.id, []),
                }
                for a, r in rows
            ],
            total=total,
            limit=filters.limit,
            offset=filters.offset,
        ).model_dump(mode="json")

    @staticmethod
    def correlation_summary(row):
        return {
            **{
                key: row.result.get(key)
                for key in ("status", "max_correlation", "compared_count", "skipped_count")
            },
            "stale": row.stale,
            "calculated_at": row.calculated_at.replace(
                tzinfo=row.calculated_at.tzinfo or timezone.utc
            ).isoformat(),
        }

    async def get_self_correlation(self, alpha_id):
        if not await self.db.get(Alpha, alpha_id):
            raise HTTPException(404, "Alpha 尚未同步")
        row = await self.db.get(SelfCorrelation, alpha_id)
        return {
            "cached": row is not None,
            "result": {
                **row.result,
                "stale": row.stale,
                "calculated_at": row.calculated_at.replace(
                    tzinfo=row.calculated_at.tzinfo or timezone.utc
                ).isoformat(),
            }
            if row
            else None,
        }

    async def get_alpha_facets(self):
        result = {}
        for key in ("region", "universe", "alpha_type", "language", "status", "stage"):
            column = getattr(Alpha, key)
            result[key] = list(
                (
                    await self.db.scalars(
                        select(column).where(column.is_not(None)).distinct().order_by(column)
                    )
                ).all()
            )
        result["tags"] = list(
            (await self.db.scalars(select(ResearchTag.tag).distinct().order_by(ResearchTag.tag))).all()
        )
        result["total"] = await self.db.scalar(select(func.count()).select_from(Alpha))
        result["favorites"] = await self.db.scalar(
            select(func.count()).select_from(Research).where(Research.favorite.is_(True))
        )
        result["last_sync"] = await self.db.scalar(select(func.max(Alpha.synced_at)))
        result["source"] = sorted(
            {kind for kinds in (await source_kinds(self.db)).values() for kind in kinds}
        )
        return result

    async def get_alpha_sources(self, alpha_id, limit=25, offset=0):
        return await alpha_sources(self.db, alpha_id, limit, offset)

    async def get_alpha(self, alpha_id):
        a, r = await self.db.get(Alpha, alpha_id), await self.db.get(Research, alpha_id)
        if a is None or r is None:
            raise HTTPException(404, "Alpha 尚未同步")
        return AlphaDetail(
            **summary(a, r),
            source_kinds=(await source_kinds(self.db, [alpha_id])).get(alpha_id, []),
            **{
                key: getattr(a, key)
                for key in (
                    "expression",
                    "selection",
                    "combo",
                    "settings",
                    "is_metrics",
                    "os_metrics",
                    "checks",
                )
            },
        ).model_dump(mode="json")

    async def get_alpha_pnl(self, alpha_id):
        if not await self.db.get(Alpha, alpha_id):
            raise HTTPException(404, "Alpha 尚未同步")
        row = await self.db.get(Pnl, alpha_id)
        return {
            "cached": row is not None,
            "points": row.points if row else [],
            "fetched_at": row.fetched_at.isoformat() if row else None,
        }

    async def update_research(self, alpha_id, body: ResearchUpdate):
        changes = body.model_dump(exclude_unset=True, exclude={"version"})
        # Compare-and-swap also works with SQLite, whose FOR UPDATE is a no-op.
        result = await self.db.execute(
            update(Research)
            .where(Research.alpha_id == alpha_id, Research.version == body.version)
            .values(**changes, version=Research.version + 1, updated_at=now())
        )
        if result.rowcount != 1:
            if not await self.db.get(Research, alpha_id):
                raise HTTPException(404, "Alpha 尚未同步")
            raise HTTPException(409, "研究记录已被修改，请刷新数据并重新确认；当前草稿已保留")
        if "tags" in changes:
            await self.db.execute(delete(ResearchTag).where(ResearchTag.alpha_id == alpha_id))
            self.db.add_all(ResearchTag(alpha_id=alpha_id, tag=t) for t in changes["tags"])
        await self.db.flush()
        return {"ok": True, "alpha_id": alpha_id, "version": body.version + 1}

    async def bulk_update_research(self, body: BulkUpdate):
        # Validate every target before touching any row. The outer transaction rolls back conflicts.
        rows = {
            r.alpha_id: r
            for r in (
                await self.db.scalars(select(Research).where(Research.alpha_id.in_(body.alpha_ids)))
            ).all()
        }
        if len(rows) != len(body.alpha_ids):
            raise HTTPException(404, "部分 Alpha 尚未同步，本次未修改任何记录")
        changes = []
        for alpha_id in body.alpha_ids:
            row = rows[alpha_id]
            tags = normalize_tags(list((set(row.tags) | set(body.add_tags)) - set(body.remove_tags)))
            values = {"tags": tags, "version": body.versions[alpha_id]}
            if body.state:
                values["state"] = body.state
            changes.append((alpha_id, ResearchUpdate(**values)))
        for alpha_id, value in changes:
            await self.update_research(alpha_id, value)
        return {"updated": len(changes), "alpha_ids": body.alpha_ids}

    async def create_sync_job(self, body: JobInput):
        account = await self.db.scalar(select(Account).where(Account.id == 1).with_for_update())
        if body.kind != "self_correlation" and (
            not account.password_encrypted or account.connection_status in ("disconnected", "error")
        ):
            raise HTTPException(409, "请先连接 WorldQuant")
        if body.kind == "self_correlation":
            found = set((await self.db.scalars(select(Alpha.id).where(Alpha.id.in_(body.alpha_ids)))).all())
            if found != set(body.alpha_ids):
                raise HTTPException(404, "部分 Alpha 尚未同步")
        payload = body.model_dump(mode="json", exclude={"kind"}, exclude_none=True)
        for job in (
            await self.db.scalars(select(Job).where(Job.kind == body.kind, Job.status.in_(ACTIVE)))
        ).all():
            if job.payload == payload:
                return JobOutput.model_validate(job).model_dump(mode="json")
        job = Job(id=str(uuid4()), kind=body.kind, payload=payload)
        self.db.add(job)
        await self.db.flush()
        return JobOutput.model_validate(job).model_dump(mode="json")

    async def list_jobs(self):
        return [
            {
                **JobOutput.model_validate(j).model_dump(mode="json"),
                "alpha_ids": j.payload.get("alpha_ids", []),
            }
            for j in (await self.db.scalars(select(Job).order_by(Job.created_at.desc()).limit(100))).all()
        ]

    async def get_job_status(self, job_id):
        job = await self.db.get(Job, job_id)
        if not job:
            raise HTTPException(404, "任务不存在")
        result = JobOutput.model_validate(job).model_dump(mode="json")
        result["alpha_ids"] = job.payload.get("alpha_ids", [])
        result["errors"] = [
            {"alpha_id": r.alpha_id, "error": r.error}
            for r in (
                await self.db.scalars(
                    select(JobItem).where(JobItem.job_id == job_id, JobItem.error.is_not(None)).limit(100)
                )
            ).all()
        ]
        return result

    async def cancel_job(self, job_id):
        job = await self.db.scalar(select(Job).where(Job.id == job_id).with_for_update())
        if not job:
            raise HTTPException(404, "任务不存在")
        if job.status in ACTIVE:
            job.cancel_requested = True
            if job.status != "running":
                job.status = "cancelled"
            job.updated_at = now()
        await self.db.flush()
        return {"ok": True, "job_id": job_id}

    async def retry_job(self, job_id):
        job = await self.db.scalar(select(Job).where(Job.id == job_id).with_for_update())
        if not job:
            raise HTTPException(404, "任务不存在")
        if job.status not in (
            "failed",
            "cancelled",
            "completed_with_errors",
            "waiting_connection",
            "waiting_auth",
        ):
            raise HTTPException(409, "该任务当前不需要重试")
        job.status, job.error, job.cancel_requested, job.next_retry_at = "queued", None, False, None
        job.updated_at = now()
        await self.db.flush()
        return JobOutput.model_validate(job).model_dump(mode="json")


async def notify_job(runner, name, result):
    """Notify the in-process runner only after the transaction has committed."""
    if name in ("start_backtest", "control_backtest"):
        runner.backtests.wake.set()
    if name == "cancel_job":
        await runner.cancel(result["job_id"])
    if name in ("create_sync_job", "retry_job"):
        runner.wake.set()