codi-api-agent 0.3.1__py3-none-any.whl

api_agent/supervisor.py

@@ -0,0 +1,178 @@
+"""Supervisor: route each query to ONE of several Agents (Phase 2 of the multi-backend plan).
+
+Each Agent owns a catalog + scoped credentials (e.g. the HTTP agent over loaded API specs,
+the SQL agent over a warehouse DSN). The supervisor picks the agent whose tools can answer,
+runs it, and — if that agent's *data source* turns out to be unreachable — falls back to the
+next-best agent, so a down warehouse degrades to the API wrapper instead of an abstain.
+
+Deliberately boring, per the plan:
+  * ONE agent per query — no cross-agent fan-out/merging (it would collide with the
+    one-source-per-table anti-stitching rules).
+  * Lexical-first routing (the same recall machinery as the op-router, at agent
+    granularity); the LLM is consulted ONLY on ambiguity, so routing usually costs zero
+    tokens.
+  * Ties prefer the agents' construction order — put the preferred backend (e.g. the
+    warehouse: fewer hops) first.
+
+`Supervisor.run(...)` mirrors `Agent.run(...)` exactly (query/history/enabled_tools/
+progress/cancel_check → AgentResult), so callers treat it as a drop-in agent.
+"""
+from __future__ import annotations
+
+import math
+from collections import Counter
+
+from .agent import Agent
+from .config import Config
+from .llm import LLMClient, extract_json
+from .log import get_logger
+from .router import _query_tokens, _tokens
+from .schemas import AgentResult
+
+_LOG = get_logger("api_agent.supervisor")
+
+# Second-place score within this fraction of first place = ambiguous → LLM tiebreak.
+_AMBIGUITY_RATIO = 0.75
+# How many of an agent's best-matching tools make up its routing score / LLM summary.
+_TOP_TOOLS = 3
+_SUMMARY_TOOLS = 6
+
+SUPERVISOR_SYSTEM = (
+    "You route a user's question to exactly ONE data agent. Each AGENT below lists a sample "
+    "of the operations it can call. Pick the agent whose operations can actually answer the "
+    "question — match by MEANING, not exact words.\n"
+    'Return ONLY JSON: {"agent": "<label>"}.'
+)
+
+
+class Supervisor:
+    def __init__(self, config: Config, agents: list[tuple[str, Agent]],
+                 llm: LLMClient | None = None):
+        """``agents`` is an ordered list of (label, Agent) — order breaks ties, so put the
+        preferred backend first. ``llm`` (for the ambiguity tiebreak only) defaults to the
+        first agent's client."""
+        if not agents:
+            raise ValueError("Supervisor needs at least one agent")
+        self.config = config
+        self.agents = agents
+        self.llm = llm or agents[0][1].llm
+
+    # ------------------------------------------------------------------ #
+    def run(self, query: str, history: list[dict] | None = None,
+            enabled_tools: list[str] | None = None,
+            progress=None, cancel_check=None) -> AgentResult:
+        """Route to one agent and run it; on a source-down abstain, try the next-best agent.
+        Same contract as Agent.run — always returns an AgentResult, never raises."""
+        if len(self.agents) == 1:
+            label, agent = self.agents[0]
+            res = agent.run(query, history=history, enabled_tools=enabled_tools,
+                            progress=progress, cancel_check=cancel_check)
+            res.answered_by = label
+            return res
+
+        order = self.route(query)
+        label, agent = order[0]
+        self._emit(progress, f"Supervisor → {label} agent")
+        res = agent.run(query, history=history, enabled_tools=enabled_tools,
+                        progress=progress, cancel_check=cancel_check)
+        res.answered_by = label
+
+        # The chosen agent's SOURCE is down (not a data/answer problem) → one fallback step.
+        if self._source_down(res) and not (cancel_check and cancel_check()):
+            fb_label, fb_agent = order[1]
+            self._emit(progress, f"{label} source unreachable — falling back to {fb_label}")
+            fb = fb_agent.run(query, history=history, enabled_tools=enabled_tools,
+                              progress=progress, cancel_check=cancel_check)
+            if not self._source_down(fb):
+                fb.answered_by = f"{fb_label} (fallback — {label} unreachable)"
+                return fb
+        return res
+
+    # ------------------------------------------------------------------ #
+    def route(self, query: str) -> list[tuple[str, Agent]]:
+        """Agents best-first. Lexical scores over the union corpus decide; construction
+        order breaks ties; the LLM is consulted only when the top two are ambiguous."""
+        scored = self._lexical_scores(query)  # [(label, agent, score)] in construction order
+        ranked = sorted(scored, key=lambda x: -x[2])  # stable → ties keep preference order
+        top, second = ranked[0], ranked[1]
+        ambiguous = top[2] <= 0 or second[2] >= _AMBIGUITY_RATIO * top[2]
+        _LOG.info("supervisor route: query=%r scores={%s}%s", query[:200],
+                  ", ".join(f"{label}: {score:.2f}" for label, _, score in scored),
+                  " — AMBIGUOUS, asking LLM" if ambiguous else f" → {top[0]}")
+        if ambiguous:
+            pick = self._llm_route(query)
+            _LOG.info("supervisor route: LLM tiebreak → %s", pick or "(failed — lexical order stands)")
+            if pick is not None:
+                ranked.sort(key=lambda x: (x[0] != pick,))  # chosen first, rest keep order
+        return [(label, agent) for label, agent, _ in ranked]
+
+    def _lexical_scores(self, query: str) -> list[tuple[str, Agent, float]]:
+        """Score each agent = the sum of its best-matching tools' idf-weighted overlap with
+        the query, with document frequencies computed over the UNION of all agents' tools so
+        scores are comparable across agents (same math as Router._lexical_rank)."""
+        q = set(_query_tokens(query))
+        docs: list[tuple[int, set]] = []  # (agent_index, token set) per tool
+        for i, (_label, agent) in enumerate(self.agents):
+            for name, tool in agent.catalog.tools.items():
+                docs.append((i, set(_tokens(f"{name} {tool.description}"))))
+        df: Counter = Counter()
+        for _, toks in docs:
+            df.update(toks)
+        n_docs = max(len(docs), 1)
+        per_agent: dict[int, list[float]] = {i: [] for i in range(len(self.agents))}
+        for i, toks in docs:
+            per_agent[i].append(
+                sum(math.log((n_docs + 1) / (df[t] + 0.5)) for t in (q & toks)))
+        return [
+            (label, agent, sum(sorted(per_agent[i], reverse=True)[:_TOP_TOOLS]))
+            for i, (label, agent) in enumerate(self.agents)
+        ]
+
+    def _llm_route(self, query: str) -> str | None:
+        """One small call: which agent should take this question? None on any failure —
+        the lexical order stands (routing never blocks on the LLM)."""
+        blocks = []
+        for label, agent in self.agents:
+            ranked = sorted(
+                agent.catalog.tools.values(),
+                key=lambda t: -len(set(_query_tokens(query))
+                                   & set(_tokens(f"{t.name} {t.description}"))))
+            sample = "\n".join(
+                f"  - {t.name}: {(t.description.splitlines() or [''])[0][:120]}"
+                for t in ranked[:_SUMMARY_TOOLS])
+            blocks.append(f"AGENT \"{label}\":\n{sample}")
+        try:
+            resp = self.llm.complete(
+                [{"role": "system", "content": SUPERVISOR_SYSTEM},
+                 {"role": "user",
+                  "content": f"QUESTION:\n{query}\n\n" + "\n\n".join(blocks)}],
+                model=self.config.router_model or self.config.generator_model,
+                json_mode=True,
+            )
+            pick = extract_json(resp.choices[0].message.content).get("agent")
+            return pick if pick in {label for label, _ in self.agents} else None
+        except Exception:
+            return None
+
+    # ------------------------------------------------------------------ #
+    @staticmethod
+    def _source_down(res: AgentResult) -> bool:
+        """True when the result is an abstain caused by the agent's data source being
+        unreachable (kind taxonomy on the trace, with the abstain note as belt) — the one
+        failure class where another backend over the same data can still answer."""
+        if res.status != "abstained":
+            return False
+        if any(not t.ok and getattr(t, "kind", None) == "unreachable" for t in res.trace):
+            return True
+        notes = (res.faithfulness.notes if res.faithfulness else "") or ""
+        return "unreachable" in notes.lower()
+
+    @staticmethod
+    def _emit(progress, message: str) -> None:
+        _LOG.info("supervisor: %s", message)
+        if progress is None:
+            return
+        try:
+            progress({"stage": "route", "message": message})
+        except Exception:
+            pass