File size: 8,421 Bytes
81e5fe7 0721bb4 81e5fe7 0721bb4 81e5fe7 0721bb4 81e5fe7 5a60e93 81e5fe7 5a60e93 81e5fe7 5a60e93 81e5fe7 5a60e93 81e5fe7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | """TaskRunner — deterministic execution of a static `TaskList`. Zero LLM.
Executes tasks in dependency order, parallelizing each ready "wave" with
`asyncio.gather`. For each task it resolves `${t<id>}` placeholders from upstream
results, does an internal `validate_args`, invokes each tool via the `ToolInvoker`
seam, and records a `TaskResult`. On failure it **degrades and continues**: the
task is marked failed, its dependents are skipped, independent branches keep
running. There is no replanning and no mid-run LLM (INV-6).
`success_criteria` is *not* machine-evaluated here (it is free text); task status
is derived from tool execution outcomes and carried to the Assembler to report.
See AGENT_ARCHITECTURE_CONTEXT_new.md §7.4.
"""
from __future__ import annotations
import asyncio
from typing import Any
from src.middlewares.logging import get_logger
from ..planner.contracts import ToolOutput, ToolRegistry
from ..planner.schemas import PLACEHOLDER_RE, Task
from ..planner.schemas import TaskList as PlanTaskList
from .invoker import ToolInvoker
from .schemas import RunState, TaskResult, TaskStatus
logger = get_logger("task_runner")
class TaskRunner:
"""Runs a `TaskList` against a `ToolInvoker`, producing a `RunState`."""
def __init__(self, invoker: ToolInvoker, registry: ToolRegistry) -> None:
self._invoker = invoker
self._registry = registry
async def run(self, task_list: PlanTaskList, business_context_id: str) -> RunState:
tasks_by_id: dict[str, Task] = {t.id: t for t in task_list.tasks}
results: dict[str, TaskResult] = {}
remaining: set[str] = set(tasks_by_id)
while remaining:
ready = [
tid
for tid in remaining
if all(dep in results for dep in tasks_by_id[tid].depends_on)
]
if not ready:
# A dependency points outside the plan (or a cycle slipped past the
# planner validator): nothing more can run. Fail the rest honestly.
for tid in list(remaining):
results[tid] = TaskResult(
task_id=tid,
stage=tasks_by_id[tid].stage,
status="failure",
objective=tasks_by_id[tid].objective,
error="unresolved dependency; task could not run",
)
remaining.discard(tid)
break
# Skip any ready task whose dependency failed (degrade-and-continue).
to_run: list[Task] = []
for tid in ready:
task = tasks_by_id[tid]
failed = [d for d in task.depends_on if results[d].status == "failure"]
if failed:
results[tid] = TaskResult(
task_id=tid,
stage=task.stage,
status="failure",
objective=task.objective,
error=f"skipped: upstream {failed} did not succeed",
)
remaining.discard(tid)
else:
to_run.append(task)
if not to_run:
continue # remaining dependents will be re-evaluated (and skipped)
wave = await asyncio.gather(
*(self._run_task(task, results) for task in to_run)
)
for tr in wave:
results[tr.task_id] = tr
remaining.discard(tr.task_id)
return RunState(
plan_id=task_list.plan_id,
business_context_id=business_context_id,
results=results,
open_questions=list(task_list.open_questions),
)
async def _run_task(self, task: Task, results: dict[str, TaskResult]) -> TaskResult:
outputs: list[ToolOutput] = []
for call in task.tool_calls:
resolved = self._resolve_args(call.args, results)
arg_error = self._validate_args(call.tool, resolved)
if arg_error is not None:
outputs.append(ToolOutput(tool=call.tool, kind="error", error=arg_error))
continue
outputs.append(await self._safe_invoke(call.tool, resolved))
status = _label(outputs)
error: str | None = None
if status == "failure":
errs = [o.error for o in outputs if o.kind == "error" and o.error]
error = errs[0] if errs else "all tool calls failed"
return TaskResult(
task_id=task.id,
stage=task.stage,
status=status,
objective=task.objective,
outputs=outputs,
error=error,
)
def _resolve_args(
self, args: dict[str, Any], results: dict[str, TaskResult]
) -> dict[str, Any]:
return {k: self._resolve_value(v, results) for k, v in args.items()}
@staticmethod
def _resolve_value(value: Any, results: dict[str, TaskResult]) -> Any:
# Resolve "${t<id>}" placeholders from upstream results, recursing into
# containers so a placeholder NESTED inside a tool arg is resolved too —
# not just a whole-value "${t2}" data arg. This is what lets a
# "${t2.customer_id}" filter value inside a retrieve_data IR become a real
# list (value-handoff): "cari-lalu-lookup" / anti-join across two steps.
if isinstance(value, str):
match = PLACEHOLDER_RE.fullmatch(value.strip())
if match:
# group(1) is "t<id>" (whole output → Pattern A, materialized to a
# DataFrame by the invoker) or "t<id>.<column>" (value-handoff →
# the list of that column's upstream values).
task_id, _, column = match.group(1).partition(".")
upstream = results.get(task_id)
if upstream is None or not upstream.outputs:
return [] if column else None
out = upstream.outputs[-1]
return _column_values(out, column) if column else out
return value
if isinstance(value, dict):
return {k: TaskRunner._resolve_value(v, results) for k, v in value.items()}
if isinstance(value, list):
return [TaskRunner._resolve_value(v, results) for v in value]
return value
def _validate_args(self, tool: str, resolved: dict[str, Any]) -> str | None:
spec = self._registry.get(tool)
if spec is None:
return f"tool {tool!r} not in registry"
required = spec.input_schema.get("required", [])
missing = [r for r in required if resolved.get(r) is None]
if missing:
return f"missing required arg(s): {sorted(missing)}"
return None
async def _safe_invoke(self, tool: str, args: dict[str, Any]) -> ToolOutput:
try:
return await self._invoker.invoke(tool, args)
except Exception as exc: # noqa: BLE001 — backstop; the invoker is never-throw (§8.4)
logger.warning("tool invoker raised", tool=tool, error=str(exc))
return ToolOutput(tool=tool, kind="error", error=f"invoker raised: {exc}")
def _column_values(out: ToolOutput, column: str) -> list[Any]:
"""Ordered, de-duplicated values of `column` from a table ToolOutput.
Backs value-handoff: a "${t<id>.<column>}" reference resolves to the SET of
that column's upstream values so an `in`/`not_in` filter can look up / anti-
join by a prior step's result (e.g. "customers who never ordered" =
`CustomerID not_in ${orders.customer_id}`). Empty list when the column is
absent or the output isn't a table — an `in` then matches nothing and a
`not_in` matches everything, the correct set semantics for an empty reference.
"""
columns = out.columns or []
if column not in columns:
return []
idx = columns.index(column)
seen: dict[Any, None] = {}
for row in out.rows or []:
if idx < len(row):
seen.setdefault(row[idx], None)
return list(seen)
def _label(outputs: list[ToolOutput]) -> TaskStatus:
if not outputs:
return "failure"
errors = sum(1 for o in outputs if o.kind == "error")
if errors == 0:
return "success"
if errors == len(outputs):
return "failure"
return "partial"
|