From b2ae3d6ceda968d0af50ec503f9689a245f2ba64 Mon Sep 17 00:00:00 2001 From: Aryan Singh Date: Wed, 1 Jul 2026 16:29:29 -0500 Subject: [PATCH] feat(agents): research loop with budget guards, compaction, depth nudge, finalize Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_016DvdRZkyCy9NJpHYDdPqmS --- agents/compaction.py | 65 ++++++++ agents/graph.py | 284 +++++++++++++++++++++++++++++++++++ scripts/smoke_toolcalling.py | 90 +++++++++++ 3 files changed, 439 insertions(+) create mode 100644 agents/compaction.py create mode 100644 agents/graph.py create mode 100644 scripts/smoke_toolcalling.py diff --git a/agents/compaction.py b/agents/compaction.py new file mode 100644 index 0000000..e2ffdc1 --- /dev/null +++ b/agents/compaction.py @@ -0,0 +1,65 @@ +"""Rolling message-history compaction — the context safety net. + +We let the research history grow into the model's full context window (131k+ on gpt-oss, +absorbed by oMLX's paged SSD KV cache) and only compact when the prompt actually approaches +that limit. When it does, we fold the older turns into a running text summary and keep the +most recent turns verbatim. +""" + +from __future__ import annotations + +from langchain_core.messages import HumanMessage, SystemMessage + +COMPACT_RATIO = 0.9 # compact once the prompt reaches this fraction of the context window +KEEP_LAST = 20 # verbatim recent turns to retain after compaction +_CHARS_PER_TOKEN = 4 # rough fallback when the server doesn't report token counts + + +def estimate_tokens(messages: list) -> int: + chars = sum( + len(m.content if isinstance(m.content, str) else str(m.content)) for m in messages + ) + return chars // _CHARS_PER_TOKEN + + +def should_compact(prompt_tokens: int, context_length: int, ratio: float = COMPACT_RATIO) -> bool: + return prompt_tokens >= ratio * context_length + + +def _render(messages: list) -> str: + out = [] + for m in messages: + role = m.__class__.__name__.replace("Message", "") + content = m.content if isinstance(m.content, str) else str(m.content) + out.append(f"{role}: {content[:800]}") + return "\n".join(out) + + +def compact(summarizer_llm, messages: list, running_summary: str, keep_last: int = KEEP_LAST): + """Return (new_messages, new_summary). Keeps system[0] + a rolling summary + recent turns.""" + if len(messages) <= keep_last + 1: + return messages, running_summary + system = messages[0] + older = messages[1:-keep_last] + recent = messages[-keep_last:] + if not older: + return messages, running_summary + + prompt = ( + "Summarize the research progress so far into concise bullet points. Preserve tickers, " + "prices, key numbers, findings, and any tentative picks.\n\n" + f"Previous summary:\n{running_summary or '(none)'}\n\nNew activity:\n{_render(older)}" + ) + try: + resp = summarizer_llm.invoke( + [ + SystemMessage(content="You compress research notes faithfully and briefly."), + HumanMessage(content=prompt), + ] + ) + summary = resp.content if isinstance(resp.content, str) else str(resp.content) + except Exception: + summary = running_summary # on failure, keep the old summary and just trim + + new_messages = [system, SystemMessage(content=f"Research summary so far:\n{summary}")] + recent + return new_messages, summary diff --git a/agents/graph.py b/agents/graph.py new file mode 100644 index 0000000..ca2203e --- /dev/null +++ b/agents/graph.py @@ -0,0 +1,284 @@ +"""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." +) + +DEEPEN_NUDGE = ( + "You're not finished. Keep researching before finalizing: evaluate additional candidate " + "tickers you haven't examined yet, pull bars/news/fundamentals for them, verify with " + "web_search, and stress-test your leading ideas against the alternatives. Continue using tools." +) + +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\": \"\", \"expiry\": \"YYYY-MM-DD\", \"strike\": . 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, + 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 + 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 wants to stop early, nudge it to keep researching + # until it has done at least min_steps (still bounded by the wall-clock deadline). + if steps_used < min_steps and time.time() < deadline_ts: + messages.append(HumanMessage(content=DEEPEN_NUDGE)) + reasoning_trace.append(f"[nudge] deepen research ({steps_used}/{min_steps} steps)") + continue + if ai.content: + reasoning_trace.append(f"[note] {str(ai.content)[:200]}") + ended_reason = "natural" + break + + for tc in tool_calls: + name, args = tc["name"], 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 diff --git a/scripts/smoke_toolcalling.py b/scripts/smoke_toolcalling.py new file mode 100644 index 0000000..74cdf2e --- /dev/null +++ b/scripts/smoke_toolcalling.py @@ -0,0 +1,90 @@ +"""Phase-1 GATE: does the local model reliably emit tool calls over MULTIPLE rounds? + +This is the single load-bearing risk of the whole design. A single tool call often works; +the failure mode is multi-turn (the model stops calling tools, or emits malformed calls). +Run this with LM Studio up and a model loaded: + + lms server start && lms load gpt-oss-20b + uv run python scripts/smoke_toolcalling.py + +PASS = the model chained >=3 real tool calls and produced a final answer. +""" + +from __future__ import annotations + +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage + +from llm.client import build_llm, list_loaded_models, resolve_model_id +from tools.market_data import get_quote +from tools.python_exec import run_python +from tools.web import web_search + +TOOLS = [run_python, web_search, get_quote] +TOOLS_BY_NAME = {t.name: t for t in TOOLS} + +SYSTEM = ( + "You are a test agent. Use a tool for EVERY factual or numeric step — never compute " + "in your head. Call tools one step at a time, then give a final summary." +) +TASK = ( + "Do these steps, each via a tool call, then summarize the results:\n" + "1) Get the last price of AAPL.\n" + "2) With run_python, compute a 7% stop-loss price (price * 0.93).\n" + "3) With run_python, compute a 15% take-profit price (price * 1.15).\n" + "4) web_search for one recent Apple (AAPL) news headline.\n" + "Report the stop price, take price, and the headline." +) +MAX_ROUNDS = 10 + + +def main() -> int: + loaded = list_loaded_models() + if not loaded: + print("FAIL: LM Studio not reachable at /v1/models.") + print("Start it: lms server start then lms load ") + return 2 + print("Loaded models :", loaded) + print("Resolved model:", resolve_model_id()) + print("-" * 60) + + llm = build_llm(TOOLS) + messages: list = [SystemMessage(content=SYSTEM), HumanMessage(content=TASK)] + tool_calls_made = 0 + rounds = 0 + + while rounds < MAX_ROUNDS: + rounds += 1 + ai: AIMessage = llm.invoke(messages) + messages.append(ai) + if not ai.tool_calls: + print(f"\n[round {rounds}] FINAL:\n{ai.content}\n") + break + for tc in ai.tool_calls: + name, args = tc["name"], tc["args"] + print(f"[round {rounds}] call {name}({args})") + tool = TOOLS_BY_NAME.get(name) + if tool is None: + result = f"ERROR: unknown tool {name}" + else: + try: + result = tool.invoke(args) + except Exception as e: + result = f"ERROR: {type(e).__name__}: {e}" + tool_calls_made += 1 + print(f" -> {str(result)[:180]}") + messages.append(ToolMessage(content=str(result), tool_call_id=tc["id"])) + + print("=" * 60) + print(f"rounds={rounds} tool_calls_executed={tool_calls_made}") + ok = tool_calls_made >= 3 + print("RESULT:", "PASS — multi-turn tool calling works" if ok else "FAIL — tool calling unreliable") + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main())