283 lines
11 KiB
Python
283 lines
11 KiB
Python
"""SQLAlchemy-backed roundtable background-job storage.
|
|
|
|
JSON columns (``dispatch_chain`` / ``thread_ids`` / ``chain``) are stored as
|
|
serialized strings in ``PortableLongText`` and round-tripped here with
|
|
``json.dumps`` / ``json.loads`` — same approach as ``ai_writing_sessions`` /
|
|
``roundtable_drafts``.
|
|
|
|
User isolation mirrors ``ai_writing_sessions``: methods take ``user_id`` defaulting
|
|
to ``AUTO`` (resolved from the run context) so backend graph nodes can write without
|
|
threading a user id through, while gateway routes pass it explicitly.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
from sqlalchemy import delete, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from deerflow.persistence.roundtable_jobs.model import RoundtableJobRow
|
|
from deerflow.runtime.user_context import AUTO, _AutoSentinel, resolve_user_id
|
|
|
|
_UNSET = object()
|
|
|
|
# 终态：列表/清理判定用。
|
|
TERMINAL_STATUSES = ("done", "error", "cancelled")
|
|
|
|
|
|
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 RoundtableJobRepository:
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._sf = session_factory
|
|
|
|
@staticmethod
|
|
def _row_to_dict(row: RoundtableJobRow) -> dict[str, Any]:
|
|
d = row.to_dict()
|
|
for k in ("created_at", "updated_at"):
|
|
if isinstance(d.get(k), datetime):
|
|
d[k] = d[k].isoformat()
|
|
for k in ("dispatch_chain", "thread_ids", "chain", "dialogues", "step3"):
|
|
d[k] = _loads(d.get(k))
|
|
return d
|
|
|
|
# ── create ────────────────────────────────────────────────────────────
|
|
|
|
async def create(
|
|
self,
|
|
*,
|
|
id: str,
|
|
draft_id: str | None = None,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
orchestration_mode: str = "recommend",
|
|
coordinator_name: str | None = None,
|
|
thread_ids: dict[str, str] | None = None,
|
|
chain: dict[str, Any] | None = None,
|
|
dispatch_chain: list[dict[str, Any]] | None = None,
|
|
status: str = "queued",
|
|
phase: str = "initializing",
|
|
) -> dict[str, Any]:
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RoundtableJobRepository.create")
|
|
now = datetime.now(UTC)
|
|
row = RoundtableJobRow(
|
|
id=id,
|
|
draft_id=draft_id,
|
|
user_id=resolved_user_id,
|
|
status=status,
|
|
phase=phase,
|
|
cycle=0,
|
|
dispatch_chain=_dumps(dispatch_chain),
|
|
active_agent_id=None,
|
|
consensus_percentage=0,
|
|
pending_clarification=None,
|
|
thread_ids=_dumps(thread_ids),
|
|
coordinator_name=coordinator_name,
|
|
orchestration_mode=orchestration_mode,
|
|
chain=_dumps(chain),
|
|
error=None,
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
async with self._sf() as session:
|
|
session.add(row)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return self._row_to_dict(row)
|
|
|
|
# ── update progress ─────────────────────────────────────────────────────
|
|
|
|
async def update_progress(
|
|
self,
|
|
job_id: str,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
status=_UNSET,
|
|
phase=_UNSET,
|
|
cycle=_UNSET,
|
|
dispatch_chain=_UNSET,
|
|
active_agent_id=_UNSET,
|
|
consensus_percentage=_UNSET,
|
|
pending_clarification=_UNSET,
|
|
thread_ids=_UNSET,
|
|
coordinator_name=_UNSET,
|
|
dialogues=_UNSET,
|
|
step3=_UNSET,
|
|
error=_UNSET,
|
|
) -> dict[str, Any] | None:
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RoundtableJobRepository.update_progress")
|
|
async with self._sf() as session:
|
|
row = await session.get(RoundtableJobRow, job_id)
|
|
if row is None:
|
|
return None
|
|
if resolved_user_id is not None and row.user_id != resolved_user_id:
|
|
return None
|
|
if status is not _UNSET:
|
|
row.status = status
|
|
if phase is not _UNSET:
|
|
row.phase = phase
|
|
if cycle is not _UNSET:
|
|
row.cycle = int(cycle)
|
|
if dispatch_chain is not _UNSET:
|
|
row.dispatch_chain = _dumps(dispatch_chain)
|
|
if active_agent_id is not _UNSET:
|
|
row.active_agent_id = active_agent_id
|
|
if consensus_percentage is not _UNSET:
|
|
row.consensus_percentage = int(consensus_percentage)
|
|
if pending_clarification is not _UNSET:
|
|
row.pending_clarification = pending_clarification
|
|
if thread_ids is not _UNSET:
|
|
row.thread_ids = _dumps(thread_ids)
|
|
if coordinator_name is not _UNSET:
|
|
row.coordinator_name = coordinator_name
|
|
if dialogues is not _UNSET:
|
|
row.dialogues = _dumps(dialogues)
|
|
if step3 is not _UNSET:
|
|
row.step3 = _dumps(step3)
|
|
if error is not _UNSET:
|
|
row.error = error
|
|
row.updated_at = datetime.now(UTC)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return self._row_to_dict(row)
|
|
|
|
# ── read ────────────────────────────────────────────────────────────────
|
|
|
|
async def get(
|
|
self,
|
|
job_id: str,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> dict[str, Any] | None:
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RoundtableJobRepository.get")
|
|
async with self._sf() as session:
|
|
row = await session.get(RoundtableJobRow, job_id)
|
|
if row is None:
|
|
return None
|
|
if resolved_user_id is not None and row.user_id != resolved_user_id:
|
|
return None
|
|
return self._row_to_dict(row)
|
|
|
|
async def get_by_draft(
|
|
self,
|
|
draft_id: str,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> dict[str, Any] | None:
|
|
"""Return the most-recently-updated job for a draft (or None)."""
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RoundtableJobRepository.get_by_draft")
|
|
stmt = select(RoundtableJobRow).where(RoundtableJobRow.draft_id == draft_id)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RoundtableJobRow.user_id == resolved_user_id)
|
|
stmt = stmt.order_by(RoundtableJobRow.updated_at.desc()).limit(1)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
row = result.scalars().first()
|
|
return self._row_to_dict(row) if row is not None else None
|
|
|
|
async def list_by_user(
|
|
self,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
limit: int = 50,
|
|
) -> list[dict[str, Any]]:
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RoundtableJobRepository.list_by_user")
|
|
stmt = select(RoundtableJobRow)
|
|
if resolved_user_id is not None:
|
|
stmt = stmt.where(RoundtableJobRow.user_id == resolved_user_id)
|
|
stmt = stmt.order_by(RoundtableJobRow.updated_at.desc()).limit(limit)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [self._row_to_dict(r) for r in result.scalars()]
|
|
|
|
async def status_map_by_user(
|
|
self,
|
|
user_id: str,
|
|
*,
|
|
limit: int = 200,
|
|
) -> dict[str, dict[str, str]]:
|
|
"""``{draft_id: {"jobId": ..., "status": ...}}`` for the latest job per draft.
|
|
|
|
Drives the history dropdown's per-draft status badge (Phase 6). Walks jobs
|
|
newest-first and keeps the first (latest) one seen for each draft_id.
|
|
"""
|
|
stmt = (
|
|
select(RoundtableJobRow.id, RoundtableJobRow.draft_id, RoundtableJobRow.status)
|
|
.where(RoundtableJobRow.user_id == user_id)
|
|
.where(RoundtableJobRow.draft_id.is_not(None))
|
|
.order_by(RoundtableJobRow.updated_at.desc())
|
|
.limit(limit)
|
|
)
|
|
out: dict[str, dict[str, str]] = {}
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
for job_id, draft_id, status in result.all():
|
|
if draft_id is None or draft_id in out:
|
|
continue
|
|
out[draft_id] = {"jobId": job_id, "status": status}
|
|
return out
|
|
|
|
# ── delete / cleanup ─────────────────────────────────────────────────────
|
|
|
|
async def delete(
|
|
self,
|
|
job_id: str,
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> bool:
|
|
resolved_user_id = resolve_user_id(user_id, method_name="RoundtableJobRepository.delete")
|
|
async with self._sf() as session:
|
|
row = await session.get(RoundtableJobRow, job_id)
|
|
if row is None:
|
|
return False
|
|
if resolved_user_id is not None and row.user_id != resolved_user_id:
|
|
return False
|
|
await session.delete(row)
|
|
await session.commit()
|
|
return True
|
|
|
|
async def list_older_than(
|
|
self,
|
|
cutoff: datetime,
|
|
*,
|
|
only_finished: bool = False,
|
|
limit: int = 500,
|
|
) -> list[str]:
|
|
"""列出 updated_at 早于 cutoff 的作业 id（最多 limit 个），供清理任务用。
|
|
|
|
only_finished=True 时只列终态（done/error/cancelled），避免清掉还在跑的作业。
|
|
"""
|
|
stmt = select(RoundtableJobRow.id).where(RoundtableJobRow.updated_at < cutoff)
|
|
if only_finished:
|
|
stmt = stmt.where(RoundtableJobRow.status.in_(TERMINAL_STATUSES))
|
|
stmt = stmt.order_by(RoundtableJobRow.updated_at.asc()).limit(limit)
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
return [row[0] for row in result.all()]
|
|
|
|
async def delete_by_ids(self, job_ids: list[str]) -> int:
|
|
if not job_ids:
|
|
return 0
|
|
stmt = delete(RoundtableJobRow).where(RoundtableJobRow.id.in_(job_ids))
|
|
async with self._sf() as session:
|
|
result = await session.execute(stmt)
|
|
await session.commit()
|
|
return int(result.rowcount or 0)
|