feat:新增投顾agent和nl2sqlagent

repositories/advisor_draft.py

@@ -0,0 +1,62 @@
+"""投顾 Agent 草稿仓储。"""
+from __future__ import annotations
+
+from sqlalchemy import select
+
+from agent.advisor_agent.drafts import (
+    DRAFT_STATUS_DISCARDED,
+    DRAFT_STATUS_DRAFT,
+)
+from model.advisor_draft import AdvisorDraft
+from repositories.base import BaseRepository
+
+
+class AdvisorDraftRepo(BaseRepository):
+    model = AdvisorDraft
+
+    async def get_by_draft_id(self, draft_id: str) -> AdvisorDraft | None:
+        return await self.db.scalar(
+            select(AdvisorDraft).where(AdvisorDraft.draft_id == draft_id)
+        )
+
+    async def save(self, draft: AdvisorDraft) -> AdvisorDraft:
+        self.db.add(draft)
+        await self.db.commit()
+        await self.db.refresh(draft)
+        return draft
+
+    async def list_drafts(
+        self,
+        *,
+        advisor_id: int | None = None,
+        customer_id: int | None = None,
+        status: str | None = None,
+        limit: int = 20,
+        offset: int = 0,
+    ) -> tuple[int, list[AdvisorDraft]]:
+        filters = []
+        if advisor_id is not None:
+            filters.append(AdvisorDraft.advisor_id == advisor_id)
+        if customer_id is not None:
+            filters.append(AdvisorDraft.customer_id == customer_id)
+        if status is not None:
+            filters.append(AdvisorDraft.status == status)
+
+        total = await self.count(where=filters)
+        items = await self.list(
+            where=filters,
+            order_by=AdvisorDraft.update_time.desc(),
+            limit=limit,
+            offset=offset,
+        )
+        return total, items
+
+    async def discard(self, draft: AdvisorDraft) -> AdvisorDraft:
+        if draft.status == DRAFT_STATUS_DISCARDED:
+            return draft
+        if draft.status != DRAFT_STATUS_DRAFT:
+            raise ValueError("草稿状态无效")
+        draft.status = DRAFT_STATUS_DISCARDED
+        await self.db.commit()
+        await self.db.refresh(draft)
+        return draft