deerflow-code/offline-backend-20260512/backend/packages/harness/deerflow/persistence/roundtable_drafts/sql.py

"""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)
            # 个人草稿严格按 user 隔离。taskId 深链聊天记录已迁出到独立的
            # ``roundtable_task_drafts`` 表（不分权），不再走这里。
            if row is None or row.user_id != user_id:
                return None
            return self._draft_to_dict(row)

    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)
            # 个人草稿严格按 user 隔离（task 聊天记录已迁出到独立存储）。
            if row is None or row.user_id != user_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)
            # 个人草稿严格按 user 隔离（task 聊天记录已迁出到独立存储）。
            if row is None or row.user_id != user_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)