Spaces:
Running
Running
File size: 6,039 Bytes
69e310f | 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 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | """Supervisor — plans, fans out in parallel, retries once, and knows when to stop.
Two independent brakes guarantee termination:
* a hard iteration cap (default 15), and
* the per-run budget guard, which trips ``BUDGET_ABORT`` from inside any node.
A ticker gets exactly one retry. Once its attempt count reaches two it is dropped
from the pending set, so a permanently broken symbol can never loop the graph.
"""
from __future__ import annotations
import logging
from typing import Any, Literal
from langgraph.graph import END
from app.core.budget import BudgetExceededError
from app.core.claude import AgentRole, LLMRequest, PromptHint
from app.core.events import EventKind
from app.graph.context import current_context
from app.graph.llm import call_model
from app.graph.prompts import SUPERVISOR_SYSTEM
from app.graph.state import (
RunState,
completeness,
pending_market_tickers,
pending_news_tickers,
)
from app.models.run import RunError, RunStatus
logger = logging.getLogger(__name__)
MAX_ATTEMPTS_PER_TICKER = 2
DispatchTarget = Literal["data_agent", "news_agent", "writer", "__end__"]
async def supervisor_node(state: RunState) -> dict[str, Any]:
"""Graph node: decide what still needs doing and dispatch the workers."""
ctx = current_context()
settings = ctx.settings
current_iteration = int(state.get("iterations", 0))
next_iteration = current_iteration + 1
# ---- brake 1: hard iteration cap -------------------------------------
if next_iteration > settings.max_iterations:
reason = (
f"supervisor hit the hard iteration cap of {settings.max_iterations} "
f"without completing the watchlist"
)
await ctx.emit(EventKind.RUN_FAILED, reason, {"iterations": current_iteration})
return {
"status": RunStatus.ITERATION_ABORT,
"abort_reason": reason,
"errors": [RunError(stage="supervisor", message=reason)],
"plan": {"dispatch": [], "market_tickers": [], "news_tickers": []},
}
# ---- brake 2: a worker already tripped the budget guard ---------------
if state.get("status") in (RunStatus.BUDGET_ABORT, RunStatus.ITERATION_ABORT):
return {"plan": {"dispatch": [], "market_tickers": [], "news_tickers": []}}
market_pending = pending_market_tickers(state, MAX_ATTEMPTS_PER_TICKER)
news_pending = pending_news_tickers(state, MAX_ATTEMPTS_PER_TICKER)
dispatch: list[str] = []
if market_pending:
dispatch.append("data_agent")
if news_pending:
dispatch.append("news_agent")
plan: dict[str, Any] = {
"dispatch": dispatch,
"market_tickers": market_pending,
"news_tickers": news_pending,
"iteration": next_iteration,
}
update: dict[str, Any] = {"iterations": 1, "plan": plan}
# The supervisor is routed by Haiku for cost discipline. Its reasoning is
# advisory: the dispatch set itself is computed from state, so a model hiccup
# can slow the run down but can never send it somewhere unsafe.
if dispatch:
try:
outcome = await call_model(
ctx,
LLMRequest(
role=AgentRole.SUPERVISOR,
hint=PromptHint.SUPERVISOR_PLAN,
system=SUPERVISOR_SYSTEM,
messages=[
{
"role": "user",
"content": (
f"Iteration {next_iteration} of {settings.max_iterations}.\n"
f"Watchlist: {', '.join(state.get('tickers', []))}\n"
f"Missing market data: {', '.join(market_pending) or 'none'}\n"
f"Missing news/sentiment: {', '.join(news_pending) or 'none'}\n"
"State one sentence on what you are dispatching and why."
),
}
],
max_tokens=300,
context={"pending_tickers": sorted(set(market_pending) | set(news_pending))},
),
)
plan["reason"] = outcome.result.text[:400]
update["token_spend"] = outcome.spend
except BudgetExceededError as exc:
await ctx.emit(EventKind.RUN_FAILED, "Budget exhausted during planning", {})
return {
"status": RunStatus.BUDGET_ABORT,
"abort_reason": str(exc),
"errors": [RunError(stage="supervisor", message=str(exc))],
"plan": {"dispatch": [], "market_tickers": [], "news_tickers": []},
}
await ctx.emit(
EventKind.SUPERVISOR_PLAN,
(
f"Dispatching {' + '.join(dispatch)} (parallel)"
if len(dispatch) > 1
else (
f"Dispatching {dispatch[0]}"
if dispatch
else "Watchlist complete — routing to writer"
)
),
{
"iteration": next_iteration,
"dispatch": dispatch,
"market_tickers": market_pending,
"news_tickers": news_pending,
"reason": plan.get("reason"),
},
)
await ctx.emit(
EventKind.STATE_PROGRESS,
"State completeness",
{"completeness": completeness(state), "iteration": next_iteration},
)
return update
def route_from_supervisor(state: RunState) -> list[str] | str:
"""Conditional edge: parallel fan-out, straight to the writer, or abort.
Returning a *list* of node names is what makes the two workers run in the
same LangGraph superstep — that is the parallelism the reducers exist for.
"""
status = state.get("status")
if status in (RunStatus.BUDGET_ABORT, RunStatus.ITERATION_ABORT):
return END
dispatch = list((state.get("plan") or {}).get("dispatch", []))
if not dispatch:
return "writer"
return dispatch
|