PMFlow/backend/app/services/ai_orchestration_run_servic...

180 lines
7.2 KiB
Python

"""AI 编排运行账本服务。"""
from __future__ import annotations
from datetime import datetime, timedelta, timezone
from typing import Any, Optional
from sqlalchemy import and_, select
from sqlalchemy.ext.asyncio import AsyncSession
from ..models import AIOrchestrationRun
def _utc_now_naive() -> datetime:
return datetime.now(timezone.utc).replace(tzinfo=None)
class AIOrchestrationRunService:
def __init__(self, db: AsyncSession) -> None:
self.db = db
async def create_run(
self,
*,
tenant_id: int,
trace_id: str,
run_type: str,
project_id: int,
signal_id: Optional[int],
response_mode: str,
status: str,
planner_status: Optional[str],
openclaw_status: Optional[str],
worker_total: int,
worker_completed: int,
worker_failed: int,
summary: str,
planner_diagnostics: Optional[dict[str, Any]],
result_payload: Optional[dict[str, Any]],
created_by_user_id: Optional[int],
started_at: Optional[datetime] = None,
completed_at: Optional[datetime] = None,
) -> AIOrchestrationRun:
row = AIOrchestrationRun(
tenant_id=tenant_id,
trace_id=trace_id,
run_type=run_type,
project_id=project_id,
signal_id=signal_id,
response_mode=response_mode,
status=status,
planner_status=planner_status,
openclaw_status=openclaw_status,
worker_total=max(0, int(worker_total)),
worker_completed=max(0, int(worker_completed)),
worker_failed=max(0, int(worker_failed)),
summary=str(summary or "").strip()[:1000] or None,
planner_diagnostics=planner_diagnostics if isinstance(planner_diagnostics, dict) else None,
result_payload=result_payload if isinstance(result_payload, dict) else None,
created_by_user_id=created_by_user_id,
started_at=started_at or _utc_now_naive(),
completed_at=completed_at,
)
self.db.add(row)
await self.db.flush()
return row
async def get_run(self, *, tenant_id: int, trace_id: str) -> Optional[AIOrchestrationRun]:
return (
await self.db.execute(
select(AIOrchestrationRun).where(
and_(
AIOrchestrationRun.tenant_id == int(tenant_id),
AIOrchestrationRun.trace_id == str(trace_id),
)
)
)
).scalar_one_or_none()
async def update_run(
self,
*,
tenant_id: int,
trace_id: str,
status: Optional[str] = None,
planner_status: Optional[str] = None,
openclaw_status: Optional[str] = None,
worker_total: Optional[int] = None,
worker_completed: Optional[int] = None,
worker_failed: Optional[int] = None,
summary: Optional[str] = None,
planner_diagnostics: Optional[dict[str, Any]] = None,
result_payload: Optional[dict[str, Any]] = None,
completed_at: Optional[datetime] = None,
) -> Optional[AIOrchestrationRun]:
row = await self.get_run(tenant_id=tenant_id, trace_id=trace_id)
if row is None:
return None
if status is not None:
row.status = str(status or "").strip()[:20]
if planner_status is not None:
row.planner_status = str(planner_status or "").strip()[:20] or None
if openclaw_status is not None:
row.openclaw_status = str(openclaw_status or "").strip()[:20] or None
if worker_total is not None:
row.worker_total = max(0, int(worker_total))
if worker_completed is not None:
row.worker_completed = max(0, int(worker_completed))
if worker_failed is not None:
row.worker_failed = max(0, int(worker_failed))
if summary is not None:
row.summary = str(summary or "").strip()[:1000] or None
if planner_diagnostics is not None:
row.planner_diagnostics = planner_diagnostics
if result_payload is not None:
row.result_payload = result_payload
if completed_at is not None:
row.completed_at = completed_at
return row
async def mark_stale_runs_interrupted(
self,
*,
older_than_minutes: int = 10,
) -> int:
threshold = _utc_now_naive() - timedelta(minutes=max(1, int(older_than_minutes)))
rows = (
await self.db.execute(
select(AIOrchestrationRun).where(
and_(
AIOrchestrationRun.status.in_(("queued", "running")),
AIOrchestrationRun.started_at.is_not(None),
AIOrchestrationRun.started_at < threshold,
)
)
)
).scalars().all()
count = 0
for row in rows:
payload = row.result_payload if isinstance(row.result_payload, dict) else {}
payload["run_status"] = "interrupted"
payload["pending_workers"] = payload.get("pending_workers") or []
row.status = "interrupted"
row.completed_at = _utc_now_naive()
row.result_payload = payload
count += 1
return count
@staticmethod
def serialize_run(row: AIOrchestrationRun) -> dict[str, Any]:
payload = row.result_payload if isinstance(row.result_payload, dict) else {}
pending_workers = payload.get("pending_workers") if isinstance(payload.get("pending_workers"), list) else []
worker_total = int(row.worker_total or 0)
worker_completed = int(row.worker_completed or 0)
worker_failed = int(row.worker_failed or 0)
worker_pending = max(0, worker_total - worker_completed - worker_failed)
if not pending_workers and worker_pending > 0:
pending_workers = [f"pending_{idx}" for idx in range(1, worker_pending + 1)]
return {
"trace_id": str(row.trace_id or ""),
"run_type": str(row.run_type or ""),
"project_id": int(row.project_id),
"signal_id": int(row.signal_id) if row.signal_id is not None else None,
"response_mode": str(row.response_mode or ""),
"run_status": str(row.status or ""),
"planner_status": str(row.planner_status or ""),
"openclaw_status": str(row.openclaw_status or ""),
"worker_total": worker_total,
"worker_completed": worker_completed,
"worker_failed": worker_failed,
"worker_pending": worker_pending,
"pending_workers": pending_workers,
"summary": str(row.summary or payload.get("summary") or "").strip(),
"planner_diagnostics": row.planner_diagnostics if isinstance(row.planner_diagnostics, dict) else None,
"result_snapshot": payload,
"created_by_user_id": int(row.created_by_user_id) if row.created_by_user_id is not None else None,
"started_at": row.started_at.isoformat() if row.started_at else None,
"completed_at": row.completed_at.isoformat() if row.completed_at else None,
}