delegate-tool 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,7 @@
1
+ Metadata-Version: 2.4
2
+ Name: delegate_tool
3
+ Version: 0.1.0
4
+ Summary: Tool to delegating tasks to other agent.
5
+ Requires-Python: >=3.13
6
+ Description-Content-Type: text/markdown
7
+ Requires-Dist: openchadpy>=0.1.23
File without changes
@@ -0,0 +1,912 @@
1
+ """
2
+ delegate Tool
3
+ =============
4
+ Delegates one or more tasks to a child agent node.
5
+
6
+ Parameters:
7
+ tasks - List of tasks/queries to delegate (required).
8
+ """
9
+
10
+ from openchadpy.context import pipeline_ctx
11
+ from openchad.chat.main import Chat
12
+ import json
13
+ import logging
14
+ import os
15
+ import re
16
+ import sys
17
+ import asyncio
18
+ from typing import Any, Dict, List, Optional, Tuple, Union
19
+
20
+ from openchadpy.context import agent_ctx, model_id_ctx, fields_ctx, additional_args_ctx
21
+ from openchadpy.pipeline_base import object_to_text_tree, PipelineBase
22
+ from openchadpy.tool_base import ToolBase, ToolRegistry
23
+
24
+ logger = logging.getLogger(__name__)
25
+
26
+
27
+ # ---------------------------------------------------------------------------
28
+ # Module-level helpers (mirrors logic from openchadpy/main.py and
29
+ # Pipeline/openchad/chat/main.py so the tool is self-contained)
30
+ # ---------------------------------------------------------------------------
31
+
32
+ _RE_BACKTICK_FENCE_PB = re.compile(
33
+ r"```(?:[a-zA-Z0-9_+\-]*)?\r?\n(.*?)```",
34
+ re.DOTALL,
35
+ )
36
+ _RE_THINK_PB = re.compile(r"<think>.*?</think>|<thinking>.*?</thinking>", re.DOTALL | re.IGNORECASE)
37
+ _RE_TOOL_PB = re.compile(r"<ToolCall\b[^>]*/>", re.DOTALL)
38
+ _RE_CBLOCK_PB = re.compile(r"<CodeBlock\b[^>]*>.*?</CodeBlock>", re.DOTALL)
39
+
40
+ def _extract_code(text: str) -> str:
41
+ """Strip MDX render tags then return the first fenced code block, or the
42
+ whole cleaned text when no fence is present.
43
+
44
+ NOTE: end() runs before finalize(), so the parser's pending buffer may
45
+ still hold the closing ``` when this is called. We handle both the
46
+ complete-fence case (regex match) and the incomplete-fence case (opening
47
+ fence only) so a missing closing ``` never breaks compilation.
48
+ """
49
+ import textwrap
50
+ text = _RE_THINK_PB.sub("", text)
51
+ text = _RE_TOOL_PB.sub("", text)
52
+ text = _RE_CBLOCK_PB.sub(
53
+ lambda m: re.sub(r"<CodeBlock\b[^>]*>", "", m.group(0)).replace("</CodeBlock>", ""),
54
+ text,
55
+ )
56
+ text = text.strip("\r\n")
57
+ matches = _RE_BACKTICK_FENCE_PB.findall(text)
58
+ if matches:
59
+ raw_code = matches[0]
60
+ else:
61
+ # Fallback: closing ``` may still be in the parser's pending buffer.
62
+ # Strip just the opening fence and treat everything after as the body.
63
+ open_match = re.match(r"```(?:[a-zA-Z0-9_+\-]*)?\r?\n(.*)", text, re.DOTALL)
64
+ if open_match:
65
+ raw_code = open_match.group(1)
66
+ # Remove any trailing incomplete ``` that made it in
67
+ raw_code = re.sub(r"\n?```\s*$", "", raw_code, flags=re.DOTALL)
68
+ else:
69
+ raw_code = text
70
+ return textwrap.dedent(raw_code.strip("\r\n")).strip()
71
+
72
+
73
+ def _build_agent_tree(flat: dict) -> dict:
74
+ """Normalise a flat agent dict and return a nested tree.
75
+
76
+ Handles two storage formats for ``children``:
77
+
78
+ * **List of IDs** ``["child_id", ...]`` – each child is a separate top-level
79
+ key in *flat* (legacy format).
80
+ * **Inline dict** ``{"child_id": {<node_data>}, ...}`` – child node data is
81
+ embedded directly inside the parent (current DB format).
82
+
83
+ Both formats are normalised into a flat ``normalized`` registry first so
84
+ that ``build_node`` can always do a simple lookup.
85
+ """
86
+ normalized: dict = {}
87
+
88
+ def _normalize_node(node_id: str, raw_node: Any) -> None:
89
+ """Register *node_id* and recursively register any inline children."""
90
+ if node_id in normalized or not isinstance(raw_node, dict):
91
+ return
92
+ node = dict(raw_node)
93
+
94
+ # Parse JSON-string fields.
95
+ for key in ("tools", "children"):
96
+ val = node.get(key)
97
+ if isinstance(val, str):
98
+ try:
99
+ node[key] = json.loads(val)
100
+ except Exception:
101
+ pass
102
+
103
+ # Normalise tools to a list.
104
+ if not isinstance(node.get("tools"), list):
105
+ node["tools"] = []
106
+
107
+ # Normalise children: convert inline dict → list of IDs and register
108
+ # each child node so build_node can look them up.
109
+ children_val = node.get("children")
110
+ if isinstance(children_val, dict):
111
+ # Inline format: {"child_id": {<node>}, ...}
112
+ node["children"] = list(children_val.keys())
113
+ for cid, cdata in children_val.items():
114
+ _normalize_node(cid, cdata)
115
+ elif isinstance(children_val, list):
116
+ # ID-list format: ["child_id", ...]
117
+ node["children"] = children_val
118
+ else:
119
+ node["children"] = []
120
+
121
+ normalized[node_id] = node
122
+
123
+ for node_id, raw_node in flat.items():
124
+ _normalize_node(node_id, raw_node)
125
+
126
+ all_child_ids = {
127
+ cid
128
+ for n in normalized.values()
129
+ for cid in n.get("children", [])
130
+ }
131
+ root_ids = [nid for nid in normalized if nid not in all_child_ids]
132
+
133
+ def build_node(nid: str) -> dict:
134
+ if nid not in normalized:
135
+ logger.warning("[_build_agent_tree] node '%s' not found – skipping", nid)
136
+ return {}
137
+ n = normalized[nid]
138
+ return {
139
+ **{k: v for k, v in n.items() if k != "children"},
140
+ "children": {
141
+ cid: build_node(cid)
142
+ for cid in n.get("children", [])
143
+ },
144
+ }
145
+
146
+ return {nid: build_node(nid) for nid in root_ids}
147
+
148
+
149
+
150
+ def _find_in_tree(tree: dict, target_id: str) -> Optional[dict]:
151
+ """Recursively search a nested agent tree for *target_id*.
152
+ Returns the matching node dict, or None if not found.
153
+ """
154
+ for nid, node in tree.items():
155
+ if nid == target_id:
156
+ return node
157
+ children = node.get("children", {})
158
+ if isinstance(children, dict) and children:
159
+ found = _find_in_tree(children, target_id)
160
+ if found is not None:
161
+ return found
162
+ return None
163
+
164
+
165
+ async def _find_agent_root(
166
+ workspace: str,
167
+ agent_id: str,
168
+ ) -> Optional[Tuple[str, dict]]:
169
+ """Find which root agent tab owns *agent_id* (which may be a child node).
170
+
171
+ Strategy:
172
+ 1. Query the raw ``agents`` SQLite table for all root tab IDs.
173
+ 2. For each root, load its flat agent list via ``Database`` and build the
174
+ nested tree with ``_build_agent_tree``.
175
+ 3. Recursively search the tree for *agent_id*.
176
+
177
+ Returns ``(root_tab_id, agent_node)`` on success, or ``None`` if not found
178
+ in any root.
179
+ """
180
+ from openchadpy.sqlite import sqlite
181
+ from openchadpy.database import Database
182
+
183
+ # 1. All root tab IDs live in the raw `agents` table (inserted by the frontend).
184
+ result = await sqlite({
185
+ "db": workspace,
186
+ "command": "query",
187
+ "sql": "SELECT id FROM agents",
188
+ })
189
+ rows = result.get("data", [])
190
+ root_ids: List[str] = [row["id"] for row in rows if row.get("id")]
191
+
192
+ logger.debug(f"[delegate] scanning {len(root_ids)} root agent(s) for '{agent_id}'")
193
+
194
+ # 2+3. Load each root's tree and search for agent_id.
195
+ for root_id in root_ids:
196
+ try:
197
+ agent_db = Database(workspace=workspace, tab_id=root_id)
198
+ flat_agents = await agent_db.get("agents")
199
+ if not flat_agents:
200
+ continue
201
+ agent_tree = _build_agent_tree(flat_agents)
202
+ node = _find_in_tree(agent_tree, agent_id)
203
+ if node is not None:
204
+ logger.info(f"[delegate] found '{agent_id}' under root '{root_id}'")
205
+ return root_id, node
206
+ except Exception as exc:
207
+ logger.warning(f"[delegate] error scanning root '{root_id}': {exc}")
208
+ continue
209
+
210
+ return None
211
+
212
+
213
+ def _parse_action_result(raw: Any) -> Optional[Dict[str, Any]]:
214
+ """Normalise a sandbox return value into a plain ActionResult-shaped dict.
215
+
216
+ Accepts:
217
+ - An ActionResult dataclass instance (has ``.result`` attribute)
218
+ - A plain dict with 'result', 'next_tasks', 'next_branch' keys
219
+ - None or anything else → returns None
220
+ """
221
+ if raw is None:
222
+ return None
223
+ if hasattr(raw, "result"):
224
+ return {
225
+ "result": raw.result,
226
+ "next_tasks": list(getattr(raw, "next_tasks", []) or []),
227
+ "next_branch": getattr(raw, "next_branch", None),
228
+ "next_branches": dict(getattr(raw, "next_branches", {}) or {}),
229
+ }
230
+ if isinstance(raw, dict) and "result" in raw:
231
+ return {
232
+ "result": raw.get("result", {}),
233
+ "next_tasks": list(raw.get("next_tasks") or []),
234
+ "next_branch": raw.get("next_branch"),
235
+ "next_branches": dict(raw.get("next_branches") or {}),
236
+ }
237
+ return None
238
+
239
+
240
+ def _serialize_ans(ans: Any) -> Any:
241
+ """Ensure that the dict returned by ``pipeline.llm_tool`` / code_sandbox is
242
+ fully JSON-serialisable.
243
+
244
+ ``code_sandbox.execute()`` returns ``{"result": <ActionResult dataclass>, ...}``.
245
+ That raw dataclass is not JSON-serialisable, which causes the downstream
246
+ ``json.dumps(tool_result)`` call in the pipeline to raise a ``TypeError``.
247
+
248
+ We convert any dataclass stored under the ``"result"`` key to a plain dict
249
+ via ``dataclasses.asdict``, leaving every other value untouched.
250
+ """
251
+ if not isinstance(ans, dict):
252
+ return ans
253
+ result_val = ans.get("result")
254
+ if result_val is not None and hasattr(result_val, "__dataclass_fields__"):
255
+ from dataclasses import asdict as _dc_asdict
256
+ try:
257
+ ans = dict(ans)
258
+ ans["result"] = _dc_asdict(result_val)
259
+ except Exception:
260
+ # Last-resort: convert each field individually.
261
+ ans = dict(ans)
262
+ serialized: Dict[str, Any] = {}
263
+ for field_name in result_val.__dataclass_fields__:
264
+ try:
265
+ serialized[field_name] = getattr(result_val, field_name)
266
+ except Exception:
267
+ serialized[field_name] = "<unreadable>"
268
+ ans["result"] = serialized
269
+ return ans
270
+
271
+
272
+ def _extract_response(ans: Any) -> Any:
273
+ """Flatten the raw ``exec_result`` / ``code_sandbox`` wrapper so the parent
274
+ agent sees only the *meaningful* payload — not the internal bookkeeping.
275
+
276
+ ``pipeline.llm_tool`` in programmatic mode returns the raw sandbox dict::
277
+
278
+ {"success": True, "result": <ActionResult>, "output": "", "error": None}
279
+
280
+ Returning this verbatim means the parent has to dig through
281
+ ``response.result.result`` (or deeper) to find actual data. This helper
282
+ unwraps it so the parent sees:
283
+
284
+ * On success → the value stored in ``ActionResult.result`` (the sandbox
285
+ code's actual return payload, e.g. ``{"path": "..."}``)
286
+ * On error → ``{"error": "<message>"}``
287
+ * Non-sandbox → the value unchanged
288
+ """
289
+ if not isinstance(ans, dict):
290
+ return ans
291
+
292
+ # Not a sandbox exec_result — pass through as-is.
293
+ if "success" not in ans:
294
+ return ans
295
+
296
+ if not ans.get("success"):
297
+ return {"error": ans.get("error") or "Unknown error"}
298
+
299
+ result = ans.get("result")
300
+ if result is None:
301
+ output = ans.get("output", "").strip()
302
+ return {"output": output} if output else {}
303
+
304
+ # Unwrap ActionResult dataclass.
305
+ if hasattr(result, "__dataclass_fields__"):
306
+ from dataclasses import asdict as _dc_asdict
307
+ try:
308
+ result = _dc_asdict(result)
309
+ except Exception:
310
+ pass
311
+
312
+ # Unwrap serialized ActionResult dict: extract only the .result payload
313
+ # so callers don't need to know about next_branch / next_tasks / etc.
314
+ if isinstance(result, dict) and "result" in result and (
315
+ "next_branch" in result or "next_branches" in result
316
+ ):
317
+ return result.get("result") or {}
318
+
319
+ return result
320
+
321
+
322
+ def _parse_bool(val: Any) -> bool:
323
+ if isinstance(val, bool):
324
+ return val
325
+ if isinstance(val, str):
326
+ return val.lower() == "true"
327
+ return False
328
+
329
+ def _create_agent_query_tool(pipeline: "PipelineBase", agent_node: dict) -> ToolRegistry:
330
+ allow_multiple = _parse_bool(agent_node.get("allowMultiple"))
331
+ if allow_multiple:
332
+ schema = {
333
+ "type": "function",
334
+ "function": {
335
+ "name": "agent_query",
336
+ "description": "Send queries to one or more agents and retrieve their responses. Use this to delegate subtasks or gather information from specialized agents.",
337
+ "parameters": {
338
+ "type": "object",
339
+ "properties": {
340
+ "queries": {
341
+ "type": "array",
342
+ "description": "List of queries to send to agents. Multiple entries allow querying several agents in a single call.",
343
+ "items": {
344
+ "type": "object",
345
+ "properties": {
346
+ "agent_id": {
347
+ "type": "string",
348
+ "description": "The unique identifier of the target agent."
349
+ },
350
+ "tasks": {
351
+ "type": "array",
352
+ "description": "List of tasks or questions to delegate to the target agent.",
353
+ "items": {
354
+ "type": "string",
355
+ "description": "A specific question or instruction for the agent."
356
+ }
357
+ }
358
+ },
359
+ "required": ["tasks"]
360
+ },
361
+ "minItems": 1
362
+ }
363
+ },
364
+ "required": ["queries"]
365
+ }
366
+ }
367
+ }
368
+ else:
369
+ schema = {
370
+ "type": "function",
371
+ "function": {
372
+ "name": "agent_query",
373
+ "description": "Send a list of tasks to a specific agent and retrieve its response. Use this to delegate subtasks or gather information from specialized agents.",
374
+ "parameters": {
375
+ "type": "object",
376
+ "properties": {
377
+ "agent_id": {
378
+ "type": "string",
379
+ "description": "The unique identifier of the target agent."
380
+ },
381
+ "tasks": {
382
+ "type": "array",
383
+ "description": "List of tasks or questions to delegate to the target agent.",
384
+ "items": {
385
+ "type": "string",
386
+ "description": "A specific question or instruction for the agent."
387
+ }
388
+ }
389
+ },
390
+ "required": ["agent_id", "tasks"]
391
+ }
392
+ }
393
+ }
394
+
395
+ async def execute_tool(**kwargs) -> Dict[str, Any]:
396
+ queries = kwargs.get("queries")
397
+ if not queries:
398
+ query_val = kwargs.get("query") or kwargs.get("tasks")
399
+ queries = [{"agent_id": kwargs.get("agent_id"), "tasks": query_val}]
400
+
401
+ async def run_agent_queries(q_item: dict) -> List[dict]:
402
+ sub_agent_id = q_item.get("agent_id")
403
+ sub_tasks = q_item.get("tasks") or q_item.get("query")
404
+ if not sub_agent_id or not sub_tasks:
405
+ return []
406
+ if isinstance(sub_tasks, str):
407
+ sub_tasks = [sub_tasks]
408
+
409
+ agent_results = []
410
+ sub_agent_tree = agent_node.get("children", {}).get(sub_agent_id)
411
+ if sub_agent_tree:
412
+ old_agent = agent_ctx.get()
413
+ old_model = model_id_ctx.get()
414
+ old_fields = fields_ctx.get()
415
+ old_additional_args = additional_args_ctx.get()
416
+
417
+ agent_ctx.set({sub_agent_id: sub_agent_tree})
418
+
419
+ sub_model = sub_agent_tree.get("model")
420
+ if sub_model:
421
+ model_id_ctx.set(sub_model)
422
+
423
+ fields_ctx.set(json.loads(sub_agent_tree.get("toolValues", "{}")))
424
+ additional_args_ctx.set(json.loads(sub_agent_tree.get("additionalArgs", "{}")))
425
+
426
+ try:
427
+ sub_agent_query = _create_agent_query_tool(pipeline, sub_agent_tree) if sub_agent_tree.get("children") else None
428
+
429
+ for task in sub_tasks:
430
+ if pipeline.cancel_event and pipeline.cancel_event.is_set():
431
+ logger.info(
432
+ "[%s] cancel_event set before agent_query sub-task '%s' – aborting",
433
+ pipeline.__class__.__name__, task,
434
+ )
435
+ break
436
+
437
+ tool_registry = {'agent_query': sub_agent_query} if sub_agent_query else None
438
+ try:
439
+ llm_task = asyncio.get_event_loop().create_task(
440
+ pipeline.llm_tool(
441
+ history_id=sub_agent_id,
442
+ query=task,
443
+ tool_registry=tool_registry,
444
+ ),
445
+ name=f"agent_query_llm_tool_{id(pipeline)}",
446
+ )
447
+ cancel_sentinel = None
448
+ if pipeline.cancel_event:
449
+ cancel_sentinel = asyncio.get_event_loop().create_task(
450
+ pipeline.cancel_event.wait(),
451
+ name=f"cancel_sentinel_aq_{id(pipeline)}",
452
+ )
453
+ try:
454
+ if cancel_sentinel is not None:
455
+ done, _ = await asyncio.wait(
456
+ {llm_task, cancel_sentinel},
457
+ return_when=asyncio.FIRST_COMPLETED,
458
+ )
459
+ if cancel_sentinel in done:
460
+ logger.info(
461
+ "[%s] agent_query: cancel_event fired – cancelling sub-task",
462
+ pipeline.__class__.__name__,
463
+ )
464
+ llm_task.cancel()
465
+ try:
466
+ await llm_task
467
+ except (asyncio.CancelledError, Exception):
468
+ pass
469
+ break
470
+ ans = llm_task.result()
471
+ else:
472
+ ans = await llm_task
473
+ finally:
474
+ if cancel_sentinel is not None and not cancel_sentinel.done():
475
+ cancel_sentinel.cancel()
476
+ try:
477
+ await cancel_sentinel
478
+ except (asyncio.CancelledError, Exception):
479
+ pass
480
+
481
+ entry = {
482
+ "agent_id": sub_agent_id,
483
+ "response": _extract_response(ans)
484
+ }
485
+ # Recursive fan-out check for programmatic child
486
+ if isinstance(ans, dict) and ans.get("success") and ans.get("result") is not None:
487
+ child_action = _parse_action_result(ans["result"])
488
+ if child_action:
489
+ has_single = child_action.get("next_branch") and child_action.get("next_tasks")
490
+ has_multi = bool(child_action.get("next_branches"))
491
+ if has_single or has_multi:
492
+ sub_results = await _fan_out_branch(pipeline, child_action)
493
+ entry["sub_branch_results"] = sub_results
494
+ agent_results.append(entry)
495
+ except Exception as e:
496
+ logger.exception("Error executing sub-agent task: %s", task)
497
+ agent_results.append({
498
+ "agent_id": sub_agent_id,
499
+ "response": f"Error: {str(e)}"
500
+ })
501
+ finally:
502
+ agent_ctx.set(old_agent)
503
+ model_id_ctx.set(old_model)
504
+ fields_ctx.set(old_fields)
505
+ additional_args_ctx.set(old_additional_args)
506
+ else:
507
+ agent_results.append({
508
+ "agent_id": sub_agent_id,
509
+ "response": f"Error: Agent '{sub_agent_id}' not found."
510
+ })
511
+ return agent_results
512
+
513
+ results = await asyncio.gather(*(run_agent_queries(q) for q in queries))
514
+ flat_results = []
515
+ for r in results:
516
+ flat_results.extend(r)
517
+ return {"results": flat_results} if allow_multiple else (flat_results[0] if flat_results else {})
518
+
519
+ return ToolRegistry(call=execute_tool, schema=schema)
520
+
521
+
522
+ async def _fan_out_single_branch(
523
+ pipeline: "PipelineBase",
524
+ child_id: str,
525
+ tasks: List[str],
526
+ cancel_event: Any,
527
+ ) -> List[Dict[str, Any]]:
528
+ """Run *tasks* against a single child branch (``child_id``).
529
+
530
+ Expects ``agent_ctx`` to already be set to the parent node so the child
531
+ node can be resolved from its ``children`` dict. The context is swapped
532
+ to the child for the duration of the call and restored before returning.
533
+
534
+ Returns a list of per-task result dicts.
535
+ """
536
+ agent = agent_ctx.get()
537
+ if not agent or not isinstance(agent, dict):
538
+ logger.warning("[_fan_out_single_branch] no agent_ctx, skipping branch '%s'", child_id)
539
+ return []
540
+
541
+ current_id = next(iter(agent))
542
+ current_node = agent[current_id]
543
+ child_node = current_node.get("children", {}).get(child_id)
544
+ if not child_node:
545
+ logger.warning(
546
+ "[_fan_out_single_branch] child node '%s' not found in '%s'",
547
+ child_id, current_id,
548
+ )
549
+ return []
550
+
551
+ # Swap agent_ctx / model_id_ctx to the child node for the duration.
552
+ old_agent = agent_ctx.get()
553
+ old_model = model_id_ctx.get()
554
+ old_fields = fields_ctx.get()
555
+ old_additional_args = additional_args_ctx.get()
556
+
557
+ agent_ctx.set({child_id: child_node})
558
+ child_model = child_node.get("model")
559
+ if child_model:
560
+ model_id_ctx.set(child_model)
561
+
562
+ fields_ctx.set(json.loads(child_node.get("toolValues", "{}")))
563
+ additional_args_ctx.set(json.loads(child_node.get("additionalArgs", "{}")))
564
+
565
+ branch_results: List[Dict[str, Any]] = []
566
+ try:
567
+ is_child_prog = _parse_bool(child_node.get("enableProgrammaticToolCalling"))
568
+ has_children = bool(child_node.get("children"))
569
+ child_agent_query = None
570
+ if not is_child_prog and has_children:
571
+ child_agent_query = _create_agent_query_tool(pipeline, child_node)
572
+
573
+ import inspect
574
+ try:
575
+ sig = inspect.signature(pipeline.llm_tool)
576
+ has_tool_registry = "tool_registry" in sig.parameters or any(
577
+ p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
578
+ )
579
+ except Exception:
580
+ has_tool_registry = False
581
+
582
+ for task in tasks:
583
+ # Pre-task cancel check.
584
+ if cancel_event and cancel_event.is_set():
585
+ logger.info(
586
+ "[_fan_out_single_branch] cancel_event set before task '%s' – aborting remaining tasks for '%s'",
587
+ task, child_id,
588
+ )
589
+ break
590
+ try:
591
+ # ── Cancel-aware LLM call (mirrors tool sentinel pattern) ──────
592
+ _llm_sentinel: Optional[asyncio.Task] = None
593
+ if cancel_event:
594
+ _llm_sentinel = asyncio.get_event_loop().create_task(
595
+ cancel_event.wait(),
596
+ name=f"cancel_sentinel_fan_out_{child_id}_{id(pipeline)}",
597
+ )
598
+ try:
599
+ llm_kwargs: Dict[str, Any] = {
600
+ "query": task,
601
+ "history_id": child_id,
602
+ }
603
+ if child_agent_query and has_tool_registry:
604
+ llm_kwargs["tool_registry"] = {"agent_query": child_agent_query}
605
+
606
+ _llm_task: asyncio.Task = asyncio.get_event_loop().create_task(
607
+ pipeline.llm_tool(**llm_kwargs),
608
+ name=f"llm_tool_fan_out_{child_id}_{id(pipeline)}",
609
+ )
610
+ try:
611
+ if _llm_sentinel is not None:
612
+ # Race: whichever finishes first wins.
613
+ _done, _ = await asyncio.wait(
614
+ {_llm_task, _llm_sentinel},
615
+ return_when=asyncio.FIRST_COMPLETED,
616
+ )
617
+ if _llm_sentinel in _done:
618
+ # Cancel was requested while the LLM call was in flight.
619
+ logger.info(
620
+ "[_fan_out_single_branch] cancel_event fired during llm_tool for task '%s' in branch '%s' – cancelling",
621
+ task, child_id,
622
+ )
623
+ _llm_task.cancel()
624
+ try:
625
+ await _llm_task
626
+ except (asyncio.CancelledError, Exception):
627
+ pass
628
+ break # exit the task loop
629
+ # LLM call finished normally – retrieve result.
630
+ ans = _llm_task.result()
631
+ else:
632
+ ans = await _llm_task
633
+ except asyncio.CancelledError:
634
+ _llm_task.cancel()
635
+ raise
636
+ except Exception:
637
+ raise
638
+ finally:
639
+ # Always clean up the sentinel so it never leaks.
640
+ if _llm_sentinel is not None and not _llm_sentinel.done():
641
+ _llm_sentinel.cancel()
642
+ try:
643
+ await _llm_sentinel
644
+ except (asyncio.CancelledError, Exception):
645
+ pass
646
+ # ─────────────────────────────────────────────────────────────
647
+ entry: Dict[str, Any] = {"agent_id": child_id, "task": task, "response": _extract_response(ans)}
648
+
649
+ # Recursive fan-out: if the child returned an ActionResult with
650
+ # next_branch+next_tasks OR next_branches, recurse.
651
+ if isinstance(ans, dict) and ans.get("success") and ans.get("result") is not None:
652
+ child_action = _parse_action_result(ans["result"])
653
+ if child_action:
654
+ has_single = child_action.get("next_branch") and child_action.get("next_tasks")
655
+ has_multi = bool(child_action.get("next_branches"))
656
+ if has_single or has_multi:
657
+ # Pre-recursion cancel check.
658
+ if cancel_event and cancel_event.is_set():
659
+ logger.info(
660
+ "[_fan_out_single_branch] cancel_event set before recursive fan-out from '%s' – skipping",
661
+ child_id,
662
+ )
663
+ else:
664
+ logger.info(
665
+ "[_fan_out_single_branch] recursive fan-out from '%s': next_branch=%s next_branches=%s tasks=%s",
666
+ child_id,
667
+ child_action.get("next_branch"),
668
+ list((child_action.get("next_branches") or {}).keys()),
669
+ child_action.get("next_tasks"),
670
+ )
671
+ # agent_ctx is currently {child_id: child_node},
672
+ # so the recursive call resolves grandchildren
673
+ # from child_node["children"].
674
+ sub_results = await _fan_out_branch(pipeline, child_action)
675
+ entry["sub_branch_results"] = sub_results
676
+
677
+ branch_results.append(entry)
678
+ except Exception as exc:
679
+ logger.exception("[_fan_out_single_branch] task failed | task=%s child=%s", task, child_id)
680
+ branch_results.append({"agent_id": child_id, "task": task, "error": str(exc)})
681
+ finally:
682
+ agent_ctx.set(old_agent)
683
+ model_id_ctx.set(old_model)
684
+ fields_ctx.set(old_fields)
685
+ additional_args_ctx.set(old_additional_args)
686
+
687
+ logger.info(
688
+ "[_fan_out_single_branch] completed %d branch tasks for '%s'",
689
+ len(branch_results),
690
+ child_id,
691
+ )
692
+ return branch_results
693
+
694
+
695
+ async def _fan_out_branch(pipeline: "PipelineBase", action_result: Dict[str, Any]) -> List[Dict[str, Any]]:
696
+ """Delegate tasks to one or more child nodes via ``pipeline.llm_tool``.
697
+
698
+ Supports two routing shapes in *action_result*:
699
+
700
+ * **Single-branch**: ``{"next_branch": str, "next_tasks": [str, ...]}``
701
+ – all tasks are sent to the one named child.
702
+
703
+ * **Multi-branch**: ``{"next_branches": {branch_id: [str, ...], ...}}``
704
+ – tasks are grouped per branch and each group is dispatched in order.
705
+
706
+ Both shapes may coexist; ``next_branches`` is processed first, then
707
+ ``next_branch`` / ``next_tasks`` if present.
708
+
709
+ Returns a flat list of branch result dicts. If a child's result is itself
710
+ an ActionResult with ``next_branch`` / ``next_tasks`` or ``next_branches``,
711
+ this function recurses automatically.
712
+
713
+ Cancellation: checks ``pipeline.cancel_event`` before every branch and
714
+ every task. On cancel, returns whatever results were already collected so
715
+ no work is orphaned and the caller can propagate the signal.
716
+ """
717
+ cancel_event = getattr(pipeline, "cancel_event", None)
718
+ all_results: List[Dict[str, Any]] = []
719
+
720
+ # ── Multi-branch path (concurrent) ───────────────────────────────────────
721
+ next_branches: Dict[str, List[str]] = action_result.get("next_branches") or {}
722
+ if next_branches:
723
+ if cancel_event and cancel_event.is_set():
724
+ logger.info(
725
+ "[_fan_out_branch] cancel_event set before multi-branch dispatch – aborting",
726
+ )
727
+ return all_results
728
+
729
+ async def _run_branch(br_id: str, br_tasks: List[str]) -> List[Dict[str, Any]]:
730
+ if cancel_event and cancel_event.is_set():
731
+ logger.info(
732
+ "[_fan_out_branch] cancel_event set before branch '%s' – skipping",
733
+ br_id,
734
+ )
735
+ return []
736
+ logger.info(
737
+ "[_fan_out_branch] multi-branch dispatch (concurrent): branch='%s', tasks=%s",
738
+ br_id, br_tasks,
739
+ )
740
+ return await _fan_out_single_branch(pipeline, br_id, br_tasks, cancel_event)
741
+
742
+ branch_results = await asyncio.gather(
743
+ *(_run_branch(br_id, br_tasks) for br_id, br_tasks in next_branches.items()),
744
+ return_exceptions=True,
745
+ )
746
+ for item in branch_results:
747
+ if isinstance(item, BaseException):
748
+ logger.error("[_fan_out_branch] branch raised an exception: %s", item)
749
+ else:
750
+ all_results.extend(item)
751
+
752
+ # ── Single-branch path ───────────────────────────────────────────────────
753
+ single_id: str | None = action_result.get("next_branch")
754
+ single_tasks: List[str] = list(action_result.get("next_tasks") or [])
755
+ if single_id and single_tasks:
756
+ if cancel_event and cancel_event.is_set():
757
+ logger.info(
758
+ "[_fan_out_branch] cancel_event set before single branch '%s' – aborting",
759
+ single_id,
760
+ )
761
+ return all_results
762
+ logger.info(
763
+ "[_fan_out_branch] single-branch dispatch: branch='%s', tasks=%s",
764
+ single_id, single_tasks,
765
+ )
766
+ results = await _fan_out_single_branch(pipeline, single_id, single_tasks, cancel_event)
767
+ all_results.extend(results)
768
+
769
+ logger.info("[_fan_out_branch] completed total %d results", len(all_results))
770
+ return all_results
771
+
772
+
773
+ class Tool(ToolBase):
774
+ """Delegate tasks to a agent node."""
775
+
776
+ name = "delegate_tool"
777
+ description = (
778
+ "Delegate a list of tasks or questions to a specialized agent node. "
779
+ "The agent will execute these tasks and return their results."
780
+ )
781
+ input_schema = {
782
+ "type": "object",
783
+ "properties": {
784
+ "tasks": {
785
+ "type": "array",
786
+ "items": {
787
+ "type": "string",
788
+ },
789
+ "description": "List of tasks or queries to run on the delegated agent.",
790
+ },
791
+ },
792
+ "required": ["tasks"],
793
+ }
794
+
795
+ allowed_callers = ["direct", "code_execution", "mcp_client"]
796
+
797
+ fields = [
798
+ {
799
+ 'name': 'Target Agent',
800
+ 'value': {
801
+ 'type': 'string',
802
+ 'placeholder': 'Agent ID...'
803
+ }
804
+ },
805
+ ]
806
+
807
+ async def run_code(self, code: str, task: str) -> Dict[str, Any]:
808
+ try:
809
+ if self.code_sandbox and self.workspace and self.tab_id:
810
+ return await self.code_sandbox.execute(code, task, workspace=self.workspace, tab_id=self.tab_id)
811
+ else:
812
+ return {"error": "Setup code_sandbox failed"}
813
+ except Exception as e:
814
+ logger.error(f"Failed to run code: {e}")
815
+ return {"error": str(e)}
816
+
817
+ async def execute(self, **kwargs) -> Dict[str, Any]: # noqa: C901
818
+ from openchadpy.database import Database
819
+
820
+ agent_id: str = self.get_field("Target Agent")
821
+ tasks: List[str] = kwargs.get("tasks", [])
822
+ pipeline: "PipelineBase | None" = pipeline_ctx.get()
823
+
824
+ if pipeline is None:
825
+ return {"error": "pipeline is required"}
826
+
827
+
828
+ logger.info(f"[delegate] agent_id={agent_id!r} tasks={tasks!r}")
829
+
830
+ if not agent_id:
831
+ return {"error": "agent_id is required and must not be empty."}
832
+ if not tasks:
833
+ return {"error": "tasks array is required and must not be empty."}
834
+
835
+ workspace = self.workspace or "global"
836
+
837
+ # ── Load agent tree ───────────────────────────────────────────────────
838
+ # agent_id may be a root tab ID or a child node ID. We try the fast
839
+ # path first (treat agent_id as a root); if that yields nothing we scan
840
+ # all roots in the workspace to find the owning tree.
841
+ try:
842
+ agent_node: Optional[dict] = None
843
+
844
+ # Fast path: agent_id is itself a root tab ID.
845
+ agent_db = Database(workspace=workspace, tab_id=agent_id)
846
+ flat_agents = await agent_db.get("agents")
847
+ if flat_agents:
848
+ agent_tree = _build_agent_tree(flat_agents)
849
+ agent_node = _find_in_tree(agent_tree, agent_id)
850
+
851
+ # Slow path: agent_id is a child node — scan all root trees.
852
+ if agent_node is None:
853
+ logger.info(
854
+ f"[delegate] '{agent_id}' not found as root; scanning all roots in workspace '{workspace}'"
855
+ )
856
+ found = await _find_agent_root(workspace, agent_id)
857
+ if found is None:
858
+ return {"error": f"Agent '{agent_id}' not found in any agent tree in workspace '{workspace}'."}
859
+ _, agent_node = found
860
+
861
+ except Exception as e:
862
+ logger.error(f"[delegate] failed to load agent tree for '{agent_id}': {e}", exc_info=True)
863
+ return {"error": f"Failed to load agent tree: {e}"}
864
+
865
+ # ── Swap context vars for the duration of delegation ─────────────────
866
+ old_agent = agent_ctx.get()
867
+ old_model = model_id_ctx.get()
868
+ old_fields = fields_ctx.get()
869
+ old_additional_args = additional_args_ctx.get()
870
+
871
+ try:
872
+ agent_ctx.set({agent_id: agent_node})
873
+
874
+ sub_model = agent_node.get("model")
875
+ if sub_model:
876
+ model_id_ctx.set(sub_model)
877
+
878
+ fields_ctx.set(json.loads(agent_node.get("toolValues", "{}")))
879
+ additional_args_ctx.set(json.loads(agent_node.get("additionalArgs", "{}")))
880
+
881
+ results: List[Dict[str, Any]] = []
882
+ for task in tasks:
883
+ try:
884
+ ans = await pipeline.llm_tool(query=task, history_id=agent_id)
885
+ # Fan-out: if the answer is an ActionResult-like dict, recurse
886
+ entry: Dict[str, Any] = {
887
+ "agent_id": agent_id,
888
+ "task": task,
889
+ "response": _extract_response(ans),
890
+ }
891
+ if isinstance(ans, dict) and ans.get("success") and ans.get("result") is not None:
892
+ child_action = _parse_action_result(ans["result"])
893
+ if child_action:
894
+ has_single = child_action.get("next_branch") and child_action.get("next_tasks")
895
+ has_multi = bool(child_action.get("next_branches"))
896
+ if has_single or has_multi:
897
+ sub_results = await _fan_out_branch(pipeline, child_action)
898
+ entry["sub_branch_results"] = sub_results
899
+ results.append(entry)
900
+ except Exception as e:
901
+ logger.error(f"[delegate] task failed for agent '{agent_id}': {e}", exc_info=True)
902
+ results.append({"agent_id": agent_id, "task": task, "error": str(e)})
903
+
904
+ return {"results": results}
905
+ finally:
906
+ agent_ctx.set(old_agent)
907
+ model_id_ctx.set(old_model)
908
+ fields_ctx.set(old_fields)
909
+ additional_args_ctx.set(old_additional_args)
910
+
911
+
912
+
@@ -0,0 +1,7 @@
1
+ Metadata-Version: 2.4
2
+ Name: delegate_tool
3
+ Version: 0.1.0
4
+ Summary: Tool to delegating tasks to other agent.
5
+ Requires-Python: >=3.13
6
+ Description-Content-Type: text/markdown
7
+ Requires-Dist: openchadpy>=0.1.23
@@ -0,0 +1,8 @@
1
+ pyproject.toml
2
+ delegate_tool/__init__.py
3
+ delegate_tool/main.py
4
+ delegate_tool.egg-info/PKG-INFO
5
+ delegate_tool.egg-info/SOURCES.txt
6
+ delegate_tool.egg-info/dependency_links.txt
7
+ delegate_tool.egg-info/requires.txt
8
+ delegate_tool.egg-info/top_level.txt
@@ -0,0 +1 @@
1
+ openchadpy>=0.1.23
@@ -0,0 +1 @@
1
+ delegate_tool
@@ -0,0 +1,9 @@
1
+ [project]
2
+ name = "delegate_tool"
3
+ version = "0.1.0"
4
+ description = "Tool to delegating tasks to other agent."
5
+ readme = "README.md"
6
+ requires-python = ">=3.13"
7
+ dependencies = [
8
+ "openchadpy>=0.1.23",
9
+ ]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+