397 lines
16 KiB
Python
397 lines
16 KiB
Python
"""SQL repository for position-collaboration sessions and chain nodes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
|
|
|
from deerflow.persistence.position_roundtable.model import (
|
|
PositionRoundtableNodeRow,
|
|
PositionRoundtableSessionRow,
|
|
)
|
|
|
|
_UNSET = object()
|
|
|
|
|
|
def _dumps(value: Any) -> str | None:
|
|
if value is None:
|
|
return None
|
|
return json.dumps(value, ensure_ascii=False)
|
|
|
|
|
|
def _loads(value: Any) -> Any:
|
|
if not isinstance(value, str) or not value:
|
|
return None
|
|
try:
|
|
return json.loads(value)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _time(value: Any) -> str | Any:
|
|
return value.isoformat() if isinstance(value, datetime) else value
|
|
|
|
|
|
def _has_material(row: PositionRoundtableNodeRow) -> bool:
|
|
return row.revision > 0 or bool(row.latest_answer) or bool(_loads(row.artifact_manifest))
|
|
|
|
|
|
class PositionRoundtableRepository:
|
|
def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
|
|
self._sf = session_factory
|
|
|
|
@staticmethod
|
|
def _session_dict(row: PositionRoundtableSessionRow) -> dict[str, Any]:
|
|
return {
|
|
"id": row.id,
|
|
"external_task_id": row.external_task_id,
|
|
"task_snapshot": _loads(row.task_snapshot) or {},
|
|
"chain_id": row.chain_id,
|
|
"chain_snapshot": _loads(row.chain_snapshot),
|
|
"intent_thread_id": row.intent_thread_id,
|
|
"intent_snapshot": _loads(row.intent_snapshot),
|
|
"summary_thread_id": row.summary_thread_id,
|
|
"summary_snapshot": _loads(row.summary_snapshot),
|
|
"action_plan_thread_id": row.action_plan_thread_id,
|
|
"action_plan_snapshot": _loads(row.action_plan_snapshot),
|
|
"status": row.status,
|
|
"created_at": _time(row.created_at),
|
|
"updated_at": _time(row.updated_at),
|
|
}
|
|
|
|
@staticmethod
|
|
def _node_dict(row: PositionRoundtableNodeRow) -> dict[str, Any]:
|
|
return {
|
|
"id": row.id,
|
|
"session_id": row.session_id,
|
|
"node_key": row.node_key,
|
|
"stage_index": row.stage_index,
|
|
"seat_index": row.seat_index,
|
|
"agent_id": row.agent_id,
|
|
"position_id": row.position_id,
|
|
"thread_id": row.thread_id,
|
|
"status": row.status,
|
|
"latest_answer": row.latest_answer,
|
|
"artifact_manifest": _loads(row.artifact_manifest) or [],
|
|
"revision": row.revision,
|
|
"upstream_revision": _loads(row.upstream_revision),
|
|
"rejection_count": row.rejection_count or 0,
|
|
"last_rejection": _loads(row.last_rejection),
|
|
"invalidated_by": row.invalidated_by,
|
|
"created_at": _time(row.created_at),
|
|
"updated_at": _time(row.updated_at),
|
|
}
|
|
|
|
async def create_session(
|
|
self,
|
|
user_id: str,
|
|
*,
|
|
external_task_id: str | None,
|
|
task_snapshot: dict[str, Any] | None,
|
|
) -> dict[str, Any]:
|
|
now = datetime.now(UTC)
|
|
row = PositionRoundtableSessionRow(
|
|
id=uuid4().hex,
|
|
user_id=user_id,
|
|
external_task_id=external_task_id or None,
|
|
task_snapshot=_dumps(task_snapshot or {}),
|
|
status="intent_pending",
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
async with self._sf() as session:
|
|
session.add(row)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return self._session_dict(row)
|
|
|
|
async def list_sessions(
|
|
self, user_id: str, *, external_task_id: str | None = None, limit: int = 50
|
|
) -> list[dict[str, Any]]:
|
|
stmt = select(PositionRoundtableSessionRow).where(PositionRoundtableSessionRow.user_id == user_id)
|
|
if external_task_id:
|
|
stmt = stmt.where(PositionRoundtableSessionRow.external_task_id == external_task_id)
|
|
stmt = stmt.order_by(PositionRoundtableSessionRow.updated_at.desc()).limit(limit)
|
|
async with self._sf() as session:
|
|
rows = (await session.execute(stmt)).scalars().all()
|
|
return [self._session_dict(row) for row in rows]
|
|
|
|
async def get_session(self, session_id: str, user_id: str) -> dict[str, Any] | None:
|
|
async with self._sf() as session:
|
|
row = await session.get(PositionRoundtableSessionRow, session_id)
|
|
if row is None or row.user_id != user_id:
|
|
return None
|
|
return self._session_dict(row)
|
|
|
|
async def update_session(
|
|
self,
|
|
session_id: str,
|
|
user_id: str,
|
|
*,
|
|
task_snapshot: dict[str, Any] | object = _UNSET,
|
|
intent_thread_id: str | None | object = _UNSET,
|
|
intent_snapshot: dict[str, Any] | None | object = _UNSET,
|
|
summary_thread_id: str | None | object = _UNSET,
|
|
summary_snapshot: dict[str, Any] | None | object = _UNSET,
|
|
action_plan_thread_id: str | None | object = _UNSET,
|
|
action_plan_snapshot: dict[str, Any] | None | object = _UNSET,
|
|
status: str | object = _UNSET,
|
|
) -> dict[str, Any] | None:
|
|
async with self._sf() as session:
|
|
row = await session.get(PositionRoundtableSessionRow, session_id)
|
|
if row is None or row.user_id != user_id:
|
|
return None
|
|
if task_snapshot is not _UNSET:
|
|
row.task_snapshot = _dumps(task_snapshot)
|
|
if intent_thread_id is not _UNSET:
|
|
row.intent_thread_id = intent_thread_id or None
|
|
if intent_snapshot is not _UNSET:
|
|
row.intent_snapshot = _dumps(intent_snapshot)
|
|
if summary_thread_id is not _UNSET:
|
|
row.summary_thread_id = summary_thread_id or None
|
|
if summary_snapshot is not _UNSET:
|
|
row.summary_snapshot = _dumps(summary_snapshot)
|
|
if action_plan_thread_id is not _UNSET:
|
|
row.action_plan_thread_id = action_plan_thread_id or None
|
|
if action_plan_snapshot is not _UNSET:
|
|
row.action_plan_snapshot = _dumps(action_plan_snapshot)
|
|
if status is not _UNSET:
|
|
row.status = str(status)
|
|
row.updated_at = datetime.now(UTC)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return self._session_dict(row)
|
|
|
|
async def activate_session(
|
|
self,
|
|
session_id: str,
|
|
user_id: str,
|
|
*,
|
|
chain_id: str,
|
|
chain_snapshot: dict[str, Any],
|
|
nodes: list[dict[str, Any]],
|
|
) -> tuple[dict[str, Any], list[dict[str, Any]]] | None:
|
|
"""Freeze a chain snapshot and create its nodes once.
|
|
|
|
Repeated requests are intentionally idempotent: once nodes exist the
|
|
existing snapshot and nodes are returned unchanged.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
async with self._sf() as session:
|
|
row = await session.get(PositionRoundtableSessionRow, session_id)
|
|
if row is None or row.user_id != user_id:
|
|
return None
|
|
node_rows = (
|
|
await session.execute(
|
|
select(PositionRoundtableNodeRow)
|
|
.where(PositionRoundtableNodeRow.session_id == session_id)
|
|
.order_by(PositionRoundtableNodeRow.stage_index, PositionRoundtableNodeRow.seat_index)
|
|
)
|
|
).scalars().all()
|
|
if node_rows:
|
|
return self._session_dict(row), [self._node_dict(node) for node in node_rows]
|
|
|
|
row.chain_id = chain_id
|
|
row.chain_snapshot = _dumps(chain_snapshot)
|
|
row.status = "active"
|
|
row.updated_at = now
|
|
node_rows = [
|
|
PositionRoundtableNodeRow(
|
|
id=uuid4().hex,
|
|
session_id=session_id,
|
|
node_key=str(node["node_key"]),
|
|
stage_index=int(node["stage_index"]),
|
|
seat_index=int(node["seat_index"]),
|
|
agent_id=str(node["agent_id"]),
|
|
position_id=node.get("position_id") or None,
|
|
status=str(node.get("status") or "locked"),
|
|
latest_answer=node.get("latest_answer") or None,
|
|
artifact_manifest=_dumps(node.get("artifact_manifest") or []),
|
|
revision=int(node.get("revision") or 0),
|
|
upstream_revision=_dumps(node.get("upstream_revision")) if node.get("upstream_revision") is not None else None,
|
|
created_at=now,
|
|
updated_at=now,
|
|
)
|
|
for node in nodes
|
|
]
|
|
session.add_all(node_rows)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
for node in node_rows:
|
|
await session.refresh(node)
|
|
return self._session_dict(row), [self._node_dict(node) for node in node_rows]
|
|
|
|
async def list_nodes(self, session_id: str, user_id: str) -> list[dict[str, Any]] | None:
|
|
async with self._sf() as session:
|
|
parent = await session.get(PositionRoundtableSessionRow, session_id)
|
|
if parent is None or parent.user_id != user_id:
|
|
return None
|
|
rows = (
|
|
await session.execute(
|
|
select(PositionRoundtableNodeRow)
|
|
.where(PositionRoundtableNodeRow.session_id == session_id)
|
|
.order_by(PositionRoundtableNodeRow.stage_index, PositionRoundtableNodeRow.seat_index)
|
|
)
|
|
).scalars().all()
|
|
return [self._node_dict(row) for row in rows]
|
|
|
|
async def get_node(self, session_id: str, node_key: str, user_id: str) -> dict[str, Any] | None:
|
|
async with self._sf() as session:
|
|
stmt = (
|
|
select(PositionRoundtableNodeRow)
|
|
.join(PositionRoundtableSessionRow, PositionRoundtableSessionRow.id == PositionRoundtableNodeRow.session_id)
|
|
.where(
|
|
PositionRoundtableNodeRow.session_id == session_id,
|
|
PositionRoundtableNodeRow.node_key == node_key,
|
|
PositionRoundtableSessionRow.user_id == user_id,
|
|
)
|
|
.limit(1)
|
|
)
|
|
row = (await session.execute(stmt)).scalars().first()
|
|
return self._node_dict(row) if row is not None else None
|
|
|
|
async def update_node(
|
|
self,
|
|
session_id: str,
|
|
node_key: str,
|
|
user_id: str,
|
|
*,
|
|
thread_id: str | None | object = _UNSET,
|
|
status: str | object = _UNSET,
|
|
latest_answer: str | None | object = _UNSET,
|
|
artifact_manifest: list[dict[str, Any]] | object = _UNSET,
|
|
upstream_revision: dict[str, Any] | object = _UNSET,
|
|
last_rejection: dict[str, Any] | None | object = _UNSET,
|
|
invalidated_by: str | None | object = _UNSET,
|
|
increment_revision: bool = False,
|
|
increment_rejection: bool = False,
|
|
) -> dict[str, Any] | None:
|
|
async with self._sf() as session:
|
|
stmt = (
|
|
select(PositionRoundtableNodeRow)
|
|
.join(PositionRoundtableSessionRow, PositionRoundtableSessionRow.id == PositionRoundtableNodeRow.session_id)
|
|
.where(
|
|
PositionRoundtableNodeRow.session_id == session_id,
|
|
PositionRoundtableNodeRow.node_key == node_key,
|
|
PositionRoundtableSessionRow.user_id == user_id,
|
|
)
|
|
.limit(1)
|
|
)
|
|
row = (await session.execute(stmt)).scalars().first()
|
|
if row is None:
|
|
return None
|
|
if thread_id is not _UNSET:
|
|
row.thread_id = thread_id or None
|
|
if status is not _UNSET:
|
|
row.status = str(status)
|
|
if latest_answer is not _UNSET:
|
|
row.latest_answer = latest_answer
|
|
if artifact_manifest is not _UNSET:
|
|
row.artifact_manifest = _dumps(artifact_manifest)
|
|
if upstream_revision is not _UNSET:
|
|
row.upstream_revision = _dumps(upstream_revision)
|
|
if last_rejection is not _UNSET:
|
|
row.last_rejection = _dumps(last_rejection)
|
|
if invalidated_by is not _UNSET:
|
|
row.invalidated_by = invalidated_by or None
|
|
if increment_revision:
|
|
row.revision += 1
|
|
if increment_rejection:
|
|
row.rejection_count += 1
|
|
row.updated_at = datetime.now(UTC)
|
|
await session.commit()
|
|
await session.refresh(row)
|
|
return self._node_dict(row)
|
|
|
|
async def reject_node(
|
|
self,
|
|
session_id: str,
|
|
node_key: str,
|
|
user_id: str,
|
|
*,
|
|
title: str,
|
|
reason: str,
|
|
rejected_by: str,
|
|
) -> tuple[dict[str, Any], list[dict[str, Any]]] | None:
|
|
"""Reject one delivered node and invalidate every later stage.
|
|
|
|
Historical answers and artifact manifests are deliberately preserved.
|
|
``invalidated_by`` gates downstream rework until the rejected upstream
|
|
stage is completed again.
|
|
"""
|
|
now = datetime.now(UTC)
|
|
async with self._sf() as session:
|
|
parent = await session.get(PositionRoundtableSessionRow, session_id)
|
|
if parent is None or parent.user_id != user_id:
|
|
return None
|
|
rows = (
|
|
await session.execute(
|
|
select(PositionRoundtableNodeRow)
|
|
.where(PositionRoundtableNodeRow.session_id == session_id)
|
|
.order_by(PositionRoundtableNodeRow.stage_index, PositionRoundtableNodeRow.seat_index)
|
|
)
|
|
).scalars().all()
|
|
target = next((row for row in rows if row.node_key == node_key), None)
|
|
if target is None:
|
|
return None
|
|
|
|
rejection = {
|
|
"title": title,
|
|
"reason": reason,
|
|
"rejected_by": rejected_by,
|
|
"rejected_at": now.isoformat(),
|
|
}
|
|
summary_snapshot = _loads(parent.summary_snapshot)
|
|
if isinstance(summary_snapshot, dict) and summary_snapshot.get("status") == "done":
|
|
summary_snapshot = {
|
|
**summary_snapshot,
|
|
"status": "stale",
|
|
"invalidated_by": target.node_key,
|
|
"invalidated_at": now.isoformat(),
|
|
}
|
|
parent.summary_snapshot = _dumps(summary_snapshot)
|
|
action_plan_snapshot = _loads(parent.action_plan_snapshot)
|
|
if isinstance(action_plan_snapshot, dict) and action_plan_snapshot.get("status") == "done":
|
|
action_plan_snapshot = {
|
|
**action_plan_snapshot,
|
|
"status": "stale",
|
|
"invalidated_by": target.node_key,
|
|
"invalidated_at": now.isoformat(),
|
|
}
|
|
parent.action_plan_snapshot = _dumps(action_plan_snapshot)
|
|
if parent.status == "completed":
|
|
parent.status = "active"
|
|
target.status = "rejected"
|
|
target.last_rejection = _dumps(rejection)
|
|
target.rejection_count += 1
|
|
target.invalidated_by = None
|
|
target.updated_at = now
|
|
|
|
for row in rows:
|
|
if row.stage_index <= target.stage_index:
|
|
continue
|
|
if _has_material(row) or row.status in {"running", "done", "stale", "rejected", "error"}:
|
|
row.status = "stale"
|
|
else:
|
|
row.status = "locked"
|
|
row.invalidated_by = target.node_key
|
|
row.updated_at = now
|
|
|
|
parent.updated_at = now
|
|
await session.commit()
|
|
await session.refresh(parent)
|
|
for row in rows:
|
|
await session.refresh(row)
|
|
return self._session_dict(parent), [self._node_dict(row) for row in rows]
|
|
|
|
async def archive_session(self, session_id: str, user_id: str) -> dict[str, Any] | None:
|
|
return await self.update_session(session_id, user_id, status="archived")
|
|
|
|
async def list_results(self, session_id: str, user_id: str) -> list[dict[str, Any]] | None:
|
|
return await self.list_nodes(session_id, user_id)
|