feat(orchestration): add orchestration system with task scheduling
Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
93
backend/app/agents/orchestration/scheduler.py
Normal file
93
backend/app/agents/orchestration/scheduler.py
Normal file
@@ -0,0 +1,93 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict, deque
|
||||
from uuid import uuid4
|
||||
|
||||
from app.agents.orchestration.budget import build_subtask_budget
|
||||
from app.agents.schemas.orchestration import SubTaskSpec, TaskGraph, TaskNode
|
||||
|
||||
|
||||
class ParallelExecutionScheduler:
|
||||
def plan(self, task_graph: TaskGraph, *, query_text: str) -> list[SubTaskSpec]:
|
||||
ordered_nodes = _topological_nodes(task_graph)
|
||||
specs: list[SubTaskSpec] = []
|
||||
for node in ordered_nodes:
|
||||
budget = build_subtask_budget(
|
||||
execution_mode=node.execution_mode,
|
||||
max_parallel_tasks=max(1, task_graph.max_parallelism),
|
||||
metadata={
|
||||
"task_graph_id": task_graph.graph_id,
|
||||
"depends_on": node.depends_on,
|
||||
},
|
||||
)
|
||||
specs.append(
|
||||
SubTaskSpec(
|
||||
subtask_id=node.node_id,
|
||||
parent_run_id=task_graph.graph_id,
|
||||
title=node.title,
|
||||
role=node.role or "master",
|
||||
goal=node.goal or query_text,
|
||||
context_slice=_build_context_slice(node, query_text),
|
||||
allowed_tools=[],
|
||||
budget_tokens=1200,
|
||||
budget_tool_calls=budget.max_tool_calls or 2,
|
||||
expected_output_schema={
|
||||
"summary": "string",
|
||||
"evidence": "list",
|
||||
"status": "completed|failed|blocked",
|
||||
},
|
||||
expected_evidence=node.expected_evidence,
|
||||
dependencies=node.depends_on,
|
||||
)
|
||||
)
|
||||
return specs
|
||||
|
||||
|
||||
def build_subtask_specs(task_graph: TaskGraph, *, query_text: str) -> list[SubTaskSpec]:
|
||||
return ParallelExecutionScheduler().plan(task_graph, query_text=query_text)
|
||||
|
||||
|
||||
def _build_context_slice(node: TaskNode, query_text: str) -> dict[str, object]:
|
||||
return {
|
||||
"query": query_text,
|
||||
"role": node.role,
|
||||
"title": node.title,
|
||||
"goal": node.goal,
|
||||
"depends_on": node.depends_on,
|
||||
}
|
||||
|
||||
|
||||
def _topological_nodes(task_graph: TaskGraph) -> list[TaskNode]:
|
||||
by_id = {node.node_id: node for node in task_graph.nodes}
|
||||
indegree = {node.node_id: 0 for node in task_graph.nodes}
|
||||
edges: dict[str, list[str]] = defaultdict(list)
|
||||
|
||||
for node in task_graph.nodes:
|
||||
for dep in node.depends_on:
|
||||
if dep not in by_id:
|
||||
continue
|
||||
edges[dep].append(node.node_id)
|
||||
indegree[node.node_id] += 1
|
||||
|
||||
ready = deque(node_id for node_id, count in indegree.items() if count == 0)
|
||||
ordered: list[TaskNode] = []
|
||||
|
||||
while ready:
|
||||
node_id = ready.popleft()
|
||||
ordered.append(by_id[node_id])
|
||||
for target in edges.get(node_id, []):
|
||||
indegree[target] -= 1
|
||||
if indegree[target] == 0:
|
||||
ready.append(target)
|
||||
|
||||
if len(ordered) != len(task_graph.nodes):
|
||||
return list(task_graph.nodes)
|
||||
return ordered
|
||||
|
||||
|
||||
def ensure_child_links(specs: list[SubTaskSpec]) -> dict[str, list[str]]:
|
||||
graph: dict[str, list[str]] = defaultdict(list)
|
||||
for spec in specs:
|
||||
for dep in spec.dependencies:
|
||||
graph[dep].append(spec.subtask_id)
|
||||
return dict(graph)
|
||||
Reference in New Issue
Block a user