Files
trading-agent/agents/graph.py
T
Aryan Singh e1ac514f30 Add dynamic agent registry and options mandate
Introduces a dashboard-managed agent registry (`state/agents.json`) used by nightly research and market-open execution, including support for custom/disabled agents and per-agent time budgets. Expands personas to four agents (adds Maverick), shifts pick targets to include mandatory options, and updates prompts/finalization nudges to enforce option research without hard-failing runs.

Strengthens execution safety by adding option-specific guardrails (conviction floor, premium cap, reserved option slot through trade/day/buying-power caps), validates OCC option symbols in schemas, and normalizes invalid option direction cases. The dashboard was refactored into modular views, adds an Agents admin tab, and improves live portfolio visibility with Alpaca-first data and snapshot fallback.
2026-07-07 22:00:20 -05:00

491 lines
23 KiB
Python

"""The agent runner: one persona researches with tools under a hard step + wall-clock
budget, then finalizes structured picks. Hand-rolled (not create_react_agent) for precise
budget control and guaranteed partial-output salvage on timeouts — the two things that
matter for unattended overnight runs.
Flow: research loop (LLM <-> tools, compaction as needed) -> finalize (structured JSON
picks). Finalize is ALWAYS attempted, even after a timeout, so a stuck agent still yields
any picks it formed rather than nothing.
"""
from __future__ import annotations
import json
import time
from datetime import datetime, timezone
from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
from loguru import logger
from agents import compaction
from agents.schemas import AgentOutput, AgentStatus, Instrument, Pick
from config.personas import Persona
from config.settings import get_settings
from llm.client import build_llm
from tools import build_toolset
RESEARCH_KICKOFF = (
"Today is {date}. This is DEEP overnight research — be exhaustive, not a quick scan. "
"Choose {n} strong picks that fit your style.\n\n"
"1. Call get_portfolio first (your current holdings).\n"
"2. Build a candidate list of 8-12 tickers in your style, then investigate EACH one: "
"get_quote, get_bars (trend/momentum), get_news (catalysts), get_fundamentals "
"(quality/valuation), and web_search / web_fetch for context and second opinions.\n"
"3. Use run_python for ALL math (position sizing, %s, ratios) — never do arithmetic yourself.\n"
"4. Compare candidates, then narrow to your best {n} and pressure-test each thesis "
"(what would make it wrong?) before finalizing.\n"
"Work one tool call at a time. Do NOT rush to finalize — keep researching until you have "
"genuinely vetted a broad set of names.\n"
"You may go LONG (bet a name rises) or SHORT (bet it falls). You may trade the STOCK or a "
"single-leg OPTION: buy a CALL for a bullish view or a PUT for a bearish view. For options, "
"call get_option_chain(underlying, 'call' or 'put') and copy an exact contract_symbol."
)
# Appended to the kickoff when the persona has an options quota (all built-ins do).
OPTIONS_MANDATE = (
"\n\nOPTIONS MANDATE — {opt} of your {n} picks MUST be a single-leg option trade, and it "
"must be your HIGHEST-probability idea, not an afterthought:\n"
"- Take your single strongest thesis and express it as a CALL (bullish) or PUT (bearish).\n"
"- Pull get_option_chain for it and evaluate 2-3 contracts: premium, bid/ask spread, "
"~30-45 days to expiration.\n"
"- Use run_python to compute the breakeven price and the premium at risk. The underlying "
"must have a specific catalyst INSIDE the contract window.\n"
"- BUDGET CHECK: one contract costs premium x 100, and the risk system caps an option "
"trade near 10% of your account equity (check equity via get_portfolio). If your favorite "
"contract doesn't fit, choose a different expression that does — further strike, nearer "
"expiry, or a cheaper underlying — your call based on the market.\n"
"- Only submit an option pick you'd give conviction >= 0.7 after pressure-testing; "
"low-conviction option picks are dropped by the risk system and waste a trade slot."
)
# Appended to every kickoff: how the two Python tools differ and how the sandbox behaves.
# Most historical agent errors were sandbox misuse (pandas in run_python, expecting files to
# persist, forgetting print) — this onboarding block prevents them up front.
SANDBOX_GUIDE = (
"\n\nPYTHON TOOLS GUIDE:\n"
"- run_python: quick exact math ONLY (stdlib: math, statistics). No pandas/yfinance here.\n"
"- run_analysis: full scripts with pandas, numpy, yfinance, matplotlib (Agg), scipy. Each "
"call runs in a FRESH empty directory — files do NOT persist between calls, so every "
"script must be self-contained (re-download data each time).\n"
"- Both: you MUST print() anything you want to see; save charts with plt.savefig('name.png')."
)
# Rotating directive nudges — cycled so a long (~2hr) research session stays PRODUCTIVE
# (does new, varied work) instead of looping "am I done?". Phrased as coverage guidance,
# not commands — the agent keeps ownership of the plan and only fills genuine gaps.
DEEPEN_NUDGES = [
"There's time left — where is your coverage thinnest? If you've examined fewer than 8-10 "
"candidates, screen a few MORE tickers in your style (get_bars / get_news / get_fundamentals).",
"If you haven't yet stress-tested your favorite idea: build the BEAR case for it. What would "
"make it wrong? Look for data that challenges the thesis (web_search + web_fetch).",
"Consider hard numbers for your leading idea: a run_analysis script (yfinance download, "
"momentum / moving averages / volatility) is far more convincing than a headline.",
"Have you checked the bigger picture? Macro + sector context (rates, the sector ETF, peers) "
"via web_search and get_bars can confirm or undercut your picks.",
"If your top pick hasn't been compared to alternatives: line it up against 2-3 sector peers "
"on valuation and momentum. Are you holding the best expression of the thesis?",
"Check timing if you haven't: catalysts and dates (earnings, events via get_news / "
"web_search), then refine entry, stop, and take-profit levels with run_python.",
]
# Takes priority over the rotation while the options mandate has had NO research yet.
OPTIONS_NUDGE = (
"Before finalizing, remember the options mandate: your strongest thesis must also be "
"expressed as a single-leg option. Call get_option_chain on your best idea, compare 2-3 "
"contracts (premium, spread, ~30-45 DTE), and compute breakeven with run_python."
)
# Selected in the final stretch of the time floor: converge instead of adding breadth.
CONSOLIDATE_NUDGE = (
"You're in the final stretch. Stop adding new names — consolidate: rank your candidates, "
"confirm entry/stop/take numbers with run_python, and double-check the catalyst dates for "
"the picks you intend to submit."
)
FINALIZE_PROMPT = (
"Research complete. Return a JSON object of the form {{\"picks\": [ ... ]}} containing about "
"{n} picks. Each pick object has these fields:\n"
'{{"ticker": "SYM", "direction": "long", "conviction": 0.72, "thesis": "...", '
'"catalysts": ["..."], "time_horizon": "swing", "requested_weight_pct": 0.08, '
'"entry_type": "market", "stop_loss_pct": 0.07, "take_profit_pct": 0.15, '
'"risk_notes": "...", "data_refs": ["..."]}}\n'
"For an OPTION pick (long call/put) also add: \"instrument\": \"option\", "
"\"option_right\": \"call\" or \"put\", \"contract_symbol\": \"<exact OCC symbol from "
"get_option_chain>\", \"expiry\": \"YYYY-MM-DD\", \"strike\": <number>. Omit these for stocks.\n"
"Rules: go long OR short as your thesis dictates (direction 'long' = bet it rises, "
"'short' = bet it falls). requested_weight_pct is a fraction of the portfolio (0.08 = 8%). "
"Base every number on your research. Output ONLY the JSON object."
)
# Deliberately NOT the persona prompt: the persona says "always use tools / never estimate",
# which makes a reasoning model defer to tools and emit an empty final channel at finalize
# time. A clean "research is done" system makes gpt-oss commit to the JSON answer.
FINALIZE_SYSTEM = (
"You are finalizing stock picks from COMPLETED research. Do not ask for tools or more "
"data — research is done. Output ONLY the requested JSON object, nothing else."
)
def _short(args: dict) -> str:
try:
return ", ".join(f"{k}={str(v)[:30]}" for k, v in args.items())
except Exception:
return str(args)[:60]
def _sanitize_tool_calls(messages: list) -> None:
"""Re-encode historical tool-call arguments to valid JSON, in place.
gpt-oss sometimes emits a tool call (especially run_analysis with a big multi-line code
string) whose raw `arguments` aren't valid JSON. It executes fine once, but when that
assistant turn is echoed back in the next request, oMLX's FunctionCall validator rejects
it with a 500 ('arguments must be valid JSON') and the whole agent dies. We rewrite the
stored raw arguments from the cleanly-parsed args dict so every echo is valid JSON.
"""
for m in messages:
ak = getattr(m, "additional_kwargs", None)
if not isinstance(ak, dict):
continue
raws = ak.get("tool_calls")
if not raws:
continue
parsed = getattr(m, "tool_calls", None) or []
for i, raw in enumerate(raws):
fn = raw.get("function")
if not isinstance(fn, dict):
continue
args = parsed[i]["args"] if i < len(parsed) and isinstance(parsed[i].get("args"), dict) else {}
fn["arguments"] = json.dumps(args)
def _repair_invalid_tool_calls(ai) -> list[tuple[str, str]]:
"""Convert `invalid_tool_calls` (malformed-JSON args) into valid empty-args calls, in place.
LangChain serializes invalid_tool_calls with their RAW malformed argument string — which
oMLX's validator 500s on at the next request (this bypasses additional_kwargs, so the
sanitizer above can't catch it). We convert each into a valid call with args={} and return
(id, name) pairs so the loop can attach an error ToolMessage telling the model to retry
with valid JSON.
"""
import uuid
repaired: list[tuple[str, str]] = []
invalid = getattr(ai, "invalid_tool_calls", None) or []
if not invalid:
return repaired
for itc in invalid:
name = ((itc.get("name") or "unknown").split("<|")[0].strip()) or "unknown"
tc_id = itc.get("id") or f"call_{uuid.uuid4().hex[:8]}"
ai.tool_calls.append({"name": name, "args": {}, "id": tc_id, "type": "tool_call"})
repaired.append((tc_id, name))
ai.invalid_tool_calls = []
return repaired
def _prompt_tokens(ai, messages: list) -> int:
"""Best available count of tokens the last prompt consumed (for compaction timing)."""
um = getattr(ai, "usage_metadata", None) or {}
if um.get("input_tokens"):
return um["input_tokens"]
tu = (getattr(ai, "response_metadata", {}) or {}).get("token_usage") or {}
if tu.get("prompt_tokens"):
return tu["prompt_tokens"]
return compaction.estimate_tokens(messages)
def _extract_between(text: str, open_ch: str, close_ch: str) -> str | None:
start, end = text.find(open_ch), text.rfind(close_ch)
if start != -1 and end != -1 and end > start:
return text[start : end + 1]
return None
def parse_picks(text: str, default_stop: float, default_take: float) -> tuple[list[Pick], list[str]]:
"""Parse the finalize output into validated Picks, collecting per-item errors.
Accepts an object {"picks": [...]}, a bare array, or those embedded in surrounding text
/ code fences. Robust to reasoning-model chatter around the JSON.
"""
t = text.strip()
if t.startswith("```"):
t = t.strip("`")
if t.lower().startswith("json"):
t = t[4:]
t = t.strip()
data = None
for candidate in (t, _extract_between(t, "{", "}"), _extract_between(t, "[", "]")):
if not candidate:
continue
try:
parsed = json.loads(candidate)
except Exception:
continue
if isinstance(parsed, dict) and isinstance(parsed.get("picks"), list):
data = parsed["picks"]
break
if isinstance(parsed, list):
data = parsed
break
if data is None:
return [], ["finalize: no JSON picks found in model output"]
picks, errors = [], []
for item in data:
if not isinstance(item, dict):
errors.append(f"finalize: skipped non-object item {item!r}")
continue
item.setdefault("stop_loss_pct", default_stop)
item.setdefault("take_profit_pct", default_take)
try:
picks.append(Pick(**item))
except Exception as e:
errors.append(f"finalize: invalid pick {item.get('ticker', '?')}: {e}")
return picks, errors
def _count_option_picks(picks: list[Pick]) -> int:
return sum(1 for p in picks if p.instrument == Instrument.OPTION)
def _finalize(persona: Persona, running_summary: str, reasoning_trace: list[str]):
"""Ask the model for structured picks using a clean finalize context (not the persona's
tool-pushing prompt), fed the research findings we gathered. Retries once on a parse
failure, and once more if the options quota is unmet and the trace actually contains
option-chain data to draw from. Fail-open: an unmet quota never discards the night."""
findings = "\n".join(reasoning_trace[-50:])
if running_summary:
findings = f"Summary:\n{running_summary}\n\nDetails:\n{findings}"
quota = persona.options_quota
option_line = (
f"\nAt least {quota} pick(s) must be OPTION picks (instrument \"option\") — your "
"highest-conviction thesis expressed as a call or put, using an exact contract_symbol "
"copied from your get_option_chain research. If your findings contain NO real contract "
"symbol from get_option_chain, do NOT fabricate one — submit stock picks instead; an "
"invented symbol is rejected and wastes the pick." if quota > 0 else ""
)
human = (
f"You are {persona.display_name}. Choose picks that fit your style "
f"(typical stop ~{persona.default_stop_pct * 100:.0f}%, "
f"take ~{persona.default_take_pct * 100:.0f}%, horizon {persona.preferred_horizon}).\n\n"
f"RESEARCH FINDINGS:\n{findings or '(sparse — use your best judgment)'}\n\n"
f"{FINALIZE_PROMPT.format(n=persona.target_pick_count)}{option_line}"
)
llm = build_llm(max_tokens=4096) # no tools, no json_mode (json_mode returns empty on oMLX)
def _ask(extra: str = "") -> tuple[list[Pick], list[str]]:
resp = llm.invoke(
[SystemMessage(content=FINALIZE_SYSTEM), HumanMessage(content=human + extra)]
)
content = resp.content if isinstance(resp.content, str) else str(resp.content)
return parse_picks(content, persona.default_stop_pct, persona.default_take_pct)
errors: list[str] = []
picks: list[Pick] = []
for _ in range(2):
try:
picks, errs = _ask()
if picks:
errors = errs
break
errors = errs
except Exception as e:
errors = [f"finalize llm error: {e}"]
if not picks:
return [], errors
# options quota: one corrective retry, only if the trace has real contracts to use
if quota > 0 and _count_option_picks(picks) < quota:
trace_has_chain = any("get_option_chain" in line for line in reasoning_trace)
if trace_has_chain:
try:
retry_picks, retry_errs = _ask(
f"\n\nYour previous answer had too few option picks ({_count_option_picks(picks)}"
f"/{quota} required). Re-issue ALL picks, converting your strongest thesis into "
"a single-leg option pick using an exact contract_symbol that appears in the "
"research findings. Do not invent contract symbols."
)
if retry_picks and _count_option_picks(retry_picks) > _count_option_picks(picks):
return retry_picks, retry_errs
except Exception as e:
errors.append(f"options-quota retry failed: {e}")
errors.append(
f"options mandate unmet: {_count_option_picks(picks)}/{quota} option picks (fail-open)"
)
return picks, errors
def run_agent(
persona: Persona,
*,
deadline_ts: float,
step_budget: int,
min_steps: int = 0,
min_seconds: float = 0.0,
context_blurb: str = "",
checkpointer=None,
model_name: str = "",
) -> AgentOutput:
start = time.time()
started_at = datetime.now(timezone.utc)
tools = build_toolset(persona)
tools_by_name = {t.name: t for t in tools}
tool_llm = build_llm(tools=tools)
plain_llm = build_llm() # no tools — used for summarization
# Compaction ceiling is the MEMORY-safe KV budget, not the model's 131k window: gpt-oss
# hits oMLX's prefill memory guard around ~47k tokens on a 24GB Mac, so we compact well
# below that to keep long (2hr) runs alive.
compact_budget = get_settings().compaction_token_budget
date_str = datetime.now().astimezone().strftime("%A %Y-%m-%d")
kickoff = RESEARCH_KICKOFF.format(date=date_str, n=persona.target_pick_count)
if persona.options_quota > 0:
kickoff += OPTIONS_MANDATE.format(opt=persona.options_quota, n=persona.target_pick_count)
kickoff += SANDBOX_GUIDE
if context_blurb: # market regime + the agent's own track record (self-learning)
kickoff = f"{context_blurb}\n\n{kickoff}"
messages: list = [
SystemMessage(content=persona.system_prompt),
HumanMessage(content=kickoff),
]
running_summary = ""
reasoning_trace: list[str] = []
steps_used = 0
nudge_count = 0
ended_reason = "natural"
def snapshot(status: AgentStatus, picks: list[Pick] | None = None, errors=None) -> AgentOutput:
return AgentOutput(
agent_name=persona.display_name,
persona_id=persona.id,
started_at=started_at,
ended_at=datetime.now(timezone.utc),
status=status,
steps_used=steps_used,
wall_seconds=round(time.time() - start, 1),
market_summary=running_summary,
picks=picks or [],
reasoning_trace=list(reasoning_trace),
errors=errors or [],
)
# --- research loop ---
while steps_used < step_budget:
if time.time() >= deadline_ts:
ended_reason = "timeout"
break
steps_used += 1
_sanitize_tool_calls(messages) # keep echoed tool-call args valid JSON (oMLX 500 guard)
ai = None
for attempt in range(4):
try:
ai = tool_llm.invoke(messages)
break
except Exception as e:
err = str(e)[:100]
reasoning_trace.append(f"[retry {attempt + 1}/4] invoke failed: {err}")
logger.warning("{}: invoke failed (attempt {}/4): {}", persona.id, attempt + 1, err)
time.sleep(4)
if attempt >= 1: # watchdog: restart oMLX if it's wedged, then keep going
from llm import server as llm_server
if not llm_server.server_up():
logger.warning("{}: oMLX down — watchdog restarting it", persona.id)
llm_server.ensure_ready(warm=False)
if ai is None:
reasoning_trace.append("[error] llm.invoke failed after 4 retries")
ended_reason = "error"
logger.warning("{}: giving up research after 4 failed invokes (will still finalize)", persona.id)
break
messages.append(ai)
# malformed-JSON tool calls: make them serializable + tell the model to retry
repaired = _repair_invalid_tool_calls(ai)
repaired_ids = {tc_id for tc_id, _ in repaired}
for tc_id, name in repaired:
messages.append(
ToolMessage(
content=f"ERROR: your {name} call had malformed JSON arguments. "
"Call it again — arguments must be a valid JSON object.",
tool_call_id=tc_id,
)
)
reasoning_trace.append(f"[repair] malformed {name} args -> asked model to retry")
tool_calls = getattr(ai, "tool_calls", None)
if not tool_calls:
# Force depth: if the model tries to finalize before the TIME floor (or step floor),
# inject a directive nudge so it keeps doing NEW research. Bounded by deadline.
# Priority: unmet options mandate first, converge in the final stretch, else rotate.
elapsed = time.time() - start
if (elapsed < min_seconds or steps_used < min_steps) and time.time() < deadline_ts:
options_pending = persona.options_quota > 0 and not any(
"get_option_chain" in line for line in reasoning_trace
)
if options_pending:
nudge = OPTIONS_NUDGE
elif min_seconds and elapsed > 0.8 * min_seconds:
nudge = CONSOLIDATE_NUDGE
else:
nudge = DEEPEN_NUDGES[nudge_count % len(DEEPEN_NUDGES)]
nudge_count += 1
messages.append(HumanMessage(content=nudge))
reasoning_trace.append(
f"[nudge #{nudge_count}] {elapsed/60:.0f}/{min_seconds/60:.0f}min elapsed, step {steps_used}"
)
continue
if ai.content:
reasoning_trace.append(f"[note] {str(ai.content)[:200]}")
ended_reason = "natural"
break
for tc in tool_calls:
if tc["id"] in repaired_ids: # already answered with a retry-error ToolMessage
continue
# gpt-oss can leak harmony channel tokens into the tool name (e.g. "get_bars<|channel|>")
name = tc["name"].split("<|")[0].strip()
args = tc.get("args", {}) or {}
tool = tools_by_name.get(name)
try:
result = tool.invoke(args) if tool else f"ERROR: unknown tool {name}"
except Exception as e:
result = f"ERROR: {type(e).__name__}: {e}"
messages.append(ToolMessage(content=str(result), tool_call_id=tc["id"]))
reasoning_trace.append(f"{name}({_short(args)}) -> {str(result)[:160]}")
if compaction.should_compact(_prompt_tokens(ai, messages), compact_budget):
messages, running_summary = compaction.compact(plain_llm, messages, running_summary)
if checkpointer: # incremental save after each round
checkpointer.write(snapshot(AgentStatus.PARTIAL))
if steps_used >= step_budget and ended_reason == "natural":
ended_reason = "budget"
# --- finalize (ALWAYS attempted — even after invoke failures, salvage the findings) ---
picks: list[Pick] = []
errors: list[str] = []
try:
picks, errors = _finalize(persona, running_summary, reasoning_trace)
except Exception as e:
errors = [f"finalize crashed: {e}"]
if ended_reason == "error":
# salvaged picks still count — mark partial so execution can use them
status = AgentStatus.PARTIAL if picks else AgentStatus.ERROR
elif ended_reason == "timeout":
status = AgentStatus.TIMEOUT
elif picks:
status = AgentStatus.COMPLETE
else:
status = AgentStatus.PARTIAL
out = snapshot(status, picks=picks, errors=errors)
if checkpointer:
checkpointer.write(out)
logger.info(
"{}: status={} steps={} picks={} wall={}s",
persona.id, status.value, steps_used, len(picks), out.wall_seconds,
)
return out