feat(agents): research loop with budget guards, compaction, depth nudge, finalize

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016DvdRZkyCy9NJpHYDdPqmS

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

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\": \"<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,
+    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

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 <model>")
+        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())