Files
2026-06-11 17:11:28 +08:00

218 lines
8.4 KiB
Python

"""SQLAlchemy-backed roundtable-planning draft + recommendation-history store.
JSON columns (``step1`` / ``step2`` / ``picks`` / ``candidates``) are stored as
serialized strings in ``PortableLongText`` and round-tripped here with
``json.dumps`` / ``json.loads`` — same approach as ``ai_writing_sessions`` — so
large Step 2 transcripts never hit MySQL's 64 KB ``TEXT`` cap.
All write/read paths are scoped by ``user_id`` for ownership isolation.
"""
from __future__ import annotations
import json
from datetime import UTC, datetime
from typing import Any
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from deerflow.persistence.roundtable_drafts.model import (
RoundtableDraftRow,
RoundtableRecommendHistoryRow,
)
_UNSET = object()
def _loads(value: Any) -> Any:
if not isinstance(value, str):
return value
if not value:
return None
try:
return json.loads(value)
except Exception:
return None
def _dumps(value: Any) -> str | None:
if value is None:
return None
return json.dumps(value, ensure_ascii=False)
class RoundtableDraftRepository:
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
self._sf = session_factory
# ── drafts ────────────────────────────────────────────────────────────
@staticmethod
def _draft_to_meta(row: RoundtableDraftRow) -> dict[str, Any]:
"""Lightweight projection for the history dropdown (no big blobs)."""
return {
"id": row.id,
"task_id": row.task_id,
"title": row.title,
"furthest_step": row.furthest_step,
"created_at": row.created_at.isoformat() if isinstance(row.created_at, datetime) else row.created_at,
"updated_at": row.updated_at.isoformat() if isinstance(row.updated_at, datetime) else row.updated_at,
}
@classmethod
def _draft_to_dict(cls, row: RoundtableDraftRow) -> dict[str, Any]:
d = cls._draft_to_meta(row)
d["step1"] = _loads(row.step1)
d["step2"] = _loads(row.step2)
d["step3"] = _loads(row.step3)
return d
async def list_drafts(self, user_id: str, *, limit: int = 100) -> list[dict[str, Any]]:
stmt = (
select(RoundtableDraftRow)
.where(RoundtableDraftRow.user_id == user_id)
.order_by(RoundtableDraftRow.updated_at.desc())
.limit(limit)
)
async with self._sf() as session:
result = await session.execute(stmt)
return [self._draft_to_meta(row) for row in result.scalars()]
async def get_draft(self, draft_id: str, user_id: str) -> dict[str, Any] | None:
async with self._sf() as session:
row = await session.get(RoundtableDraftRow, draft_id)
# task 绑定的草稿共享:放开 user 归属校验(不分权);非 task 草稿仍按 user 隔离。
if row is None or (row.user_id != user_id and not row.task_id):
return None
return self._draft_to_dict(row)
async def get_latest_draft_by_task(self, task_id: str) -> dict[str, Any] | None:
"""按外部任务 id 取最新一条草稿(**忽略 user_id**,不分权)。无则 None。"""
stmt = (
select(RoundtableDraftRow)
.where(RoundtableDraftRow.task_id == task_id)
.order_by(RoundtableDraftRow.updated_at.desc())
.limit(1)
)
async with self._sf() as session:
result = await session.execute(stmt)
row = result.scalars().first()
return self._draft_to_dict(row) if row is not None else None
async def create_draft(self, user_id: str, data: dict[str, Any]) -> dict[str, Any]:
now = datetime.now(UTC)
row = RoundtableDraftRow(
id=data["id"],
user_id=user_id,
task_id=(data.get("task_id") or None),
title=(data.get("title") or "").strip(),
furthest_step=int(data.get("furthest_step") or 1),
step1=_dumps(data.get("step1")),
step2=_dumps(data.get("step2")),
step3=_dumps(data.get("step3")),
created_at=now,
updated_at=now,
)
async with self._sf() as session:
session.add(row)
await session.commit()
await session.refresh(row)
return self._draft_to_dict(row)
async def update_draft(
self,
draft_id: str,
user_id: str,
*,
title=_UNSET,
furthest_step=_UNSET,
step1=_UNSET,
step2=_UNSET,
step3=_UNSET,
task_id=_UNSET,
) -> dict[str, Any] | None:
async with self._sf() as session:
row = await session.get(RoundtableDraftRow, draft_id)
# task 绑定的草稿共享:放开 user 归属校验(不分权)。
if row is None or (row.user_id != user_id and not row.task_id):
return None
if task_id is not _UNSET and task_id is not None:
row.task_id = str(task_id)
if title is not _UNSET and title is not None:
row.title = str(title).strip()
if furthest_step is not _UNSET and furthest_step is not None:
row.furthest_step = int(furthest_step)
if step1 is not _UNSET:
row.step1 = _dumps(step1)
if step2 is not _UNSET:
row.step2 = _dumps(step2)
if step3 is not _UNSET:
row.step3 = _dumps(step3)
row.updated_at = datetime.now(UTC)
await session.commit()
await session.refresh(row)
return self._draft_to_dict(row)
async def delete_draft(self, draft_id: str, user_id: str) -> bool:
async with self._sf() as session:
row = await session.get(RoundtableDraftRow, draft_id)
# task 绑定的草稿共享:放开 user 归属校验(不分权)。
if row is None or (row.user_id != user_id and not row.task_id):
return False
await session.delete(row)
await session.commit()
return True
# ── recommendation history ────────────────────────────────────────────
@staticmethod
def _history_to_dict(row: RoundtableRecommendHistoryRow) -> dict[str, Any]:
return {
"id": row.id,
"draft_id": row.draft_id,
"objective": row.objective,
"status": row.status,
"model": row.model,
"rationale": row.rationale,
"picks": _loads(row.picks) or [],
"candidates": _loads(row.candidates) or [],
"created_at": row.created_at.isoformat() if isinstance(row.created_at, datetime) else row.created_at,
}
async def list_recommendations(
self, draft_id: str, user_id: str, *, limit: int = 20
) -> list[dict[str, Any]]:
stmt = (
select(RoundtableRecommendHistoryRow)
.where(
RoundtableRecommendHistoryRow.draft_id == draft_id,
RoundtableRecommendHistoryRow.user_id == user_id,
)
.order_by(RoundtableRecommendHistoryRow.created_at.desc())
.limit(limit)
)
async with self._sf() as session:
result = await session.execute(stmt)
return [self._history_to_dict(row) for row in result.scalars()]
async def add_recommendation(self, user_id: str, draft_id: str, data: dict[str, Any]) -> dict[str, Any]:
now = datetime.now(UTC)
row = RoundtableRecommendHistoryRow(
id=data["id"],
user_id=user_id,
draft_id=draft_id,
objective=(data.get("objective") or "").strip()[:512],
status=(data.get("status") or "done")[:32],
model=(data.get("model") or None),
rationale=data.get("rationale") or None,
picks=_dumps(data.get("picks") or []),
candidates=_dumps(data.get("candidates") or []),
created_at=now,
)
async with self._sf() as session:
session.add(row)
await session.commit()
await session.refresh(row)
return self._history_to_dict(row)