feat: add AI research chatbot with confirmed business tools

2026-09-07 23:02:55 +08:00
parent 3cd280d068
commit 79ab20b4ea
44 changed files with 4647 additions and 197 deletions
@@ -0,0 +1,202 @@
+"""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 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, now
+from .schemas import AlphaDetail, AlphaPage, BulkUpdate, JobInput, JobOutput, ResearchUpdate, normalize_tags
+
+
+class Business:
+    def __init__(self, db):
+        self.db = db
+
+    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()
+        return AlphaPage(
+            items=[summary(a, r) for a, r in rows], total=total, limit=filters.limit, offset=filters.offset
+        ).model_dump(mode="json")
+
+    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)))
+        return result
+
+    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),
+            **{
+                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 not account.password_encrypted or account.connection_status in ("disconnected", "error"):
+            raise HTTPException(409, "请先连接 WorldQuant")
+        payload = {"alpha_ids": body.alpha_ids}
+        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 == "cancel_job":
+        await runner.cancel(result["job_id"])
+    if name in ("create_sync_job", "retry_job"):
+        runner.wake.set()