backend/core/agent_subagents.py

"""
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,
    )