127 lines
4.2 KiB
Python
127 lines
4.2 KiB
Python
"""Project the latest run-scoped progress snapshot from append-only messages."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import Iterable
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
from sqlalchemy import select
|
|
|
|
from core.storage.models import Message
|
|
|
|
_VALID_STATUSES = {"pending", "in_progress", "completed"}
|
|
|
|
|
|
def _normalize_step(value: Any) -> dict[str, str] | None:
|
|
if not isinstance(value, dict):
|
|
return None
|
|
step_id = str(value.get("id") or "").strip()
|
|
title = str(value.get("title") or "").strip()
|
|
status = str(value.get("status") or "pending")
|
|
if not step_id or not title:
|
|
return None
|
|
if status not in _VALID_STATUSES:
|
|
status = "pending"
|
|
return {"id": step_id, "title": title, "status": status}
|
|
|
|
|
|
def _heal_monotonic(steps: list[dict[str, str]]) -> list[dict[str, str]]:
|
|
last_completed = max(
|
|
(i for i, step in enumerate(steps) if step["status"] == "completed"),
|
|
default=-1,
|
|
)
|
|
return [
|
|
{**step, "status": "completed"} if i < last_completed else dict(step)
|
|
for i, step in enumerate(steps)
|
|
]
|
|
|
|
|
|
def apply_progress_args(
|
|
current: list[dict[str, str]], args: Any,
|
|
) -> list[dict[str, str]]:
|
|
"""Apply current full snapshots plus legacy set/update/clear calls."""
|
|
if not isinstance(args, dict):
|
|
return [dict(step) for step in current]
|
|
action = args.get("action") or ""
|
|
if action == "clear":
|
|
return []
|
|
if isinstance(args.get("steps"), list):
|
|
normalized: list[dict[str, str]] = []
|
|
for raw in args["steps"]:
|
|
step = _normalize_step(raw)
|
|
if step is not None:
|
|
normalized.append(step)
|
|
return _heal_monotonic(normalized)
|
|
if action != "update_step" or not isinstance(args.get("step"), dict):
|
|
return [dict(step) for step in current]
|
|
|
|
raw = args["step"]
|
|
step_id = str(raw.get("id") or "").strip()
|
|
if not step_id:
|
|
return [dict(step) for step in current]
|
|
next_steps: list[dict[str, str]] = []
|
|
found = False
|
|
for step in current:
|
|
if step["id"] != step_id:
|
|
next_steps.append(dict(step))
|
|
continue
|
|
found = True
|
|
status = str(raw.get("status") or step["status"])
|
|
next_steps.append({
|
|
"id": step_id,
|
|
"title": str(raw.get("title") or step["title"]).strip(),
|
|
"status": status if status in _VALID_STATUSES else "pending",
|
|
})
|
|
if not found:
|
|
normalized = _normalize_step(raw)
|
|
if normalized is not None:
|
|
next_steps.append(normalized)
|
|
return _heal_monotonic(next_steps)
|
|
|
|
|
|
def project_progress_payloads(
|
|
payloads: Iterable[dict[str, Any]],
|
|
) -> tuple[list[dict[str, str]], bool]:
|
|
steps: list[dict[str, str]] = []
|
|
seen = False
|
|
for payload in payloads:
|
|
if not isinstance(payload, dict) or payload.get("role") != "assistant":
|
|
continue
|
|
for call in payload.get("tool_calls") or []:
|
|
function = call.get("function") if isinstance(call, dict) else None
|
|
if not isinstance(function, dict) or function.get("name") != "task_progress":
|
|
continue
|
|
raw_args = function.get("arguments") or "{}"
|
|
try:
|
|
args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
|
|
except (TypeError, json.JSONDecodeError):
|
|
args = {}
|
|
steps = apply_progress_args(steps, args)
|
|
seen = True
|
|
return steps, seen
|
|
|
|
|
|
def latest_run_progress(session, task_id: UUID) -> dict[str, Any] | None:
|
|
"""Return the latest user message id and that run's projected plan."""
|
|
run = session.execute(
|
|
select(Message.message_id, Message.idx)
|
|
.where(
|
|
Message.task_id == task_id,
|
|
Message.payload["role"].astext == "user",
|
|
)
|
|
.order_by(Message.idx.desc())
|
|
.limit(1)
|
|
).first()
|
|
if run is None:
|
|
return None
|
|
payloads = session.execute(
|
|
select(Message.payload)
|
|
.where(Message.task_id == task_id, Message.idx > run.idx)
|
|
.order_by(Message.idx)
|
|
).scalars().all()
|
|
steps, seen = project_progress_payloads(payloads)
|
|
if not seen:
|
|
return None
|
|
return {"run_id": str(run.message_id), "steps": steps}
|