File size: 5,677 Bytes
969891d | 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 | """
Sub-agent fan-out.
Some questions decompose cleanly into independent investigations — "compare these
five species", "which of these regions has the highest abundance". Running those
in one agent means a long serial chain of tool calls and a context that fills
with intermediate rows.
`spawn_subagents` runs them concurrently instead. Each child gets its own scoped
question and a read-only toolset, and returns a short text finding. The parent
sees only those findings, so its context stays small and it can focus on
synthesis.
Children are deliberately limited:
* read-only tools — only the parent draws on the map, so N children cannot
fight over the single map layer or produce N conflicting charts;
* a smaller iteration budget, since each has a narrow question;
* a hard cap on how many run at once.
"""
from __future__ import annotations
import asyncio
import logging
from typing import Any, Dict, List
from backend.core.agent_tools import AgentContext, Tool, build_tools
logger = logging.getLogger(__name__)
# Parallel children per call. Above this the token cost grows faster than the
# insight, and the model starts inventing work to fill the slots.
MAX_SUBAGENTS = 6
# Children answer one narrow question; they do not need the parent's budget.
SUBAGENT_MAX_ITERATIONS = 6
# Tools a child may use: discovery and querying, no output-producing ones.
SUBAGENT_TOOL_NAMES = ("search_datasets", "describe_table", "sample_values",
"run_sql", "compute_stats")
SUBAGENT_PROMPT_SUFFIX = (
"\n\nYou are a sub-agent handling ONE part of a larger question. Investigate "
"only what you were asked. You cannot draw on the map or make charts — the "
"parent agent does that. Reply with a short, factual finding including the "
"concrete numbers you found and the table(s) you used. No preamble."
)
def build_subagent_tool(client, model: str, parent_ctx: AgentContext) -> Tool:
"""Build the `spawn_subagents` tool bound to a parent run."""
async def _run_child(task: str, index: int) -> Dict[str, Any]:
# Imported here to avoid a circular import at module load: agent_loop
# imports agent_tools, and this module is wired into agent_loop's toolset.
from backend.core.agent_loop import GeoAgent, AGENT_SYSTEM_PROMPT
# A child gets its own context so its queries cannot mutate the parent's
# map layer or chart, but inherits the dataset scope.
child_ctx = AgentContext(allowed_datasets=parent_ctx.allowed_datasets)
full_tools = build_tools(child_ctx)
child_tools = {n: full_tools[n] for n in SUBAGENT_TOOL_NAMES if n in full_tools}
agent = GeoAgent(client, model)
# Replace the toolset with the read-only subset.
agent.extra_tools = child_tools
finding = ""
try:
async for event in agent.run(
question=task + SUBAGENT_PROMPT_SUFFIX,
history=[],
ctx=child_ctx,
max_iterations=SUBAGENT_MAX_ITERATIONS,
):
if event.get("type") == "final":
finding = event.get("text", "")
elif event.get("type") == "error":
finding = f"(failed: {event.get('message')})"
except Exception as e: # noqa: BLE001 - one child must not sink the batch
logger.warning(f"Sub-agent {index} failed: {e}", exc_info=True)
finding = f"(failed: {e})"
# Roll the child's SQL up so the parent can cite what was actually run.
parent_ctx.sql_statements.extend(child_ctx.sql_statements)
return {"task": task, "finding": finding or "(no finding)"}
def spawn_subagents(tasks: List[str]) -> Dict[str, Any]:
if not tasks:
return {"findings": [], "note": "No tasks given."}
capped = tasks[:MAX_SUBAGENTS]
dropped = len(tasks) - len(capped)
async def _gather():
return await asyncio.gather(
*(_run_child(t, i) for i, t in enumerate(capped))
)
# build_tools' implementations are sync and executed via asyncio.to_thread,
# so this runs on a worker thread with no running loop of its own.
findings = asyncio.run(_gather())
out: Dict[str, Any] = {"findings": findings}
if dropped:
# Never silently drop work the model asked for.
out["note"] = f"Ran {len(capped)} of {len(tasks)} tasks (max {MAX_SUBAGENTS})."
return out
return Tool(
name="spawn_subagents",
description=(
"Investigate several INDEPENDENT sub-questions in parallel and get back one "
"finding each. Use when a question naturally splits — one species per task, "
"one region per task. Each sub-agent can query data but cannot draw on the "
"map or make charts, so do that yourself afterwards using their findings. "
"Do not use for a single question, or for steps that depend on each other."
),
parameters={
"type": "object",
"properties": {
"tasks": {
"type": "array",
"items": {"type": "string"},
"description": (
f"Self-contained questions, max {MAX_SUBAGENTS}. Each must name its "
"own subject explicitly, e.g. 'What is the peak weekly abundance of "
"the Barn Swallow and in which week?'"
),
}
},
"required": ["tasks"],
},
run=spawn_subagents,
)
|