305 lines
13 KiB
Python
305 lines
13 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, 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 stocks 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."
|
|
)
|
|
|
|
# Rotating directive nudges — cycled so a long (~2hr) research session stays PRODUCTIVE
|
|
# (does new, varied work) instead of looping "am I done?". Keep using tools.
|
|
DEEPEN_NUDGES = [
|
|
"Not done yet. Screen 5-10 MORE candidate tickers in your style you haven't examined; pull "
|
|
"get_bars / get_news / get_fundamentals on them.",
|
|
"Build the BEAR case for your current top pick: what would make it wrong? Find data that "
|
|
"challenges your thesis (web_search + web_fetch recent coverage).",
|
|
"Use run_analysis to write a real script: fetch prices with yfinance, compute momentum / "
|
|
"moving averages / volatility / correlations, and save a chart to back your best idea.",
|
|
"Zoom out: check the macro + sector context (rates, the sector ETF, peers) via web_search and "
|
|
"get_bars. Does it support or undercut your picks?",
|
|
"Compare your leading picks head-to-head against 2-3 sector peers on valuation and momentum "
|
|
"(get_fundamentals + run_analysis). Are you picking the best expression of the thesis?",
|
|
"Re-examine catalysts and timing for each candidate (earnings dates, events via get_news / "
|
|
"web_search) and refine entry, stop, and take-profit levels with run_python.",
|
|
]
|
|
|
|
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 _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 _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."""
|
|
findings = "\n".join(reasoning_trace[-50:])
|
|
if running_summary:
|
|
findings = f"Summary:\n{running_summary}\n\nDetails:\n{findings}"
|
|
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)}"
|
|
)
|
|
llm = build_llm(max_tokens=4096) # no tools, no json_mode (json_mode returns empty on oMLX)
|
|
errors: list[str] = []
|
|
for _ in range(2):
|
|
try:
|
|
resp = llm.invoke(
|
|
[SystemMessage(content=FINALIZE_SYSTEM), HumanMessage(content=human)]
|
|
)
|
|
content = resp.content if isinstance(resp.content, str) else str(resp.content)
|
|
picks, errs = parse_picks(content, persona.default_stop_pct, persona.default_take_pct)
|
|
if picks:
|
|
return picks, errs
|
|
errors = errs
|
|
except Exception as e:
|
|
errors = [f"finalize llm error: {e}"]
|
|
return [], errors
|
|
|
|
|
|
def run_agent(
|
|
persona: Persona,
|
|
*,
|
|
deadline_ts: float,
|
|
step_budget: int,
|
|
min_steps: int = 0,
|
|
min_seconds: float = 0.0,
|
|
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
|
|
ctx_len = get_settings().llm_context_length
|
|
|
|
date_str = datetime.now().astimezone().strftime("%A %Y-%m-%d")
|
|
messages: list = [
|
|
SystemMessage(content=persona.system_prompt),
|
|
HumanMessage(content=RESEARCH_KICKOFF.format(date=date_str, n=persona.target_pick_count)),
|
|
]
|
|
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
|
|
try:
|
|
ai = tool_llm.invoke(messages)
|
|
except Exception as e:
|
|
reasoning_trace.append(f"[error] llm.invoke: {e}")
|
|
ended_reason = "error"
|
|
logger.warning("{}: llm.invoke failed: {}", persona.id, e)
|
|
break
|
|
messages.append(ai)
|
|
|
|
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 rotating directive nudge so it keeps doing NEW research. Bounded by deadline.
|
|
elapsed = time.time() - start
|
|
if (elapsed < min_seconds or steps_used < min_steps) and time.time() < deadline_ts:
|
|
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:
|
|
# 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), ctx_len):
|
|
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 on timeout, to salvage picks) ---
|
|
picks: list[Pick] = []
|
|
errors: list[str] = []
|
|
if ended_reason != "error":
|
|
picks, errors = _finalize(persona, running_summary, reasoning_trace)
|
|
|
|
if ended_reason == "error":
|
|
status = 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
|