181 lines
7.2 KiB
Python
181 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
|
|
|
|
|
|
from ...utils.timezone import now_cn_naive
|
|
def _utc_now_naive() -> datetime:
|
|
return now_cn_naive()
|
|
|
|
|
|
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],
|
|
execution_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=execution_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,
|
|
execution_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 execution_status is not None:
|
|
row.openclaw_status = str(execution_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 ""),
|
|
"execution_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,
|
|
}
|