fix(graph): key checkpoints on graph shape and expose the LLM retry budget

- checkpoint resume keyed only by ticker+date silently continued the old graph
  under a different analyst selection / depth / asset mode; fold a run signature
  into the thread id #1089
- add llm_max_retries + TRADINGAGENTS_LLM_MAX_RETRIES, forwarded to every provider
  when set (int-coerced, rejects negatives/booleans), so a 429 burst can't abort
  a run #1091
This commit is contained in:
Yijia-Xiao
2026-07-05 14:29:07 +00:00
parent b47a828a4f
commit daf1da9c35
5 changed files with 244 additions and 13 deletions

View File

@@ -140,5 +140,79 @@ class TestCheckpointResume(unittest.TestCase):
self.assertTrue(has_checkpoint(self.tmpdir, self.ticker, self.date))
class TestCheckpointSignature(unittest.TestCase):
"""A different graph shape (analyst selection / depth / asset mode) must not
resume the previous run's checkpoint (#1089)."""
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
self.ticker = "TEST"
self.date = "2026-04-20"
def test_empty_signature_is_legacy_id(self):
self.assertEqual(
thread_id(self.ticker, self.date),
thread_id(self.ticker, self.date, ""),
)
def test_signature_changes_thread_id(self):
legacy = thread_id(self.ticker, self.date)
sig_a = thread_id(self.ticker, self.date, "analysts=market,news|asset=stock")
sig_b = thread_id(self.ticker, self.date, "analysts=market|asset=stock")
self.assertNotEqual(sig_a, sig_b) # different graph shapes differ
self.assertNotEqual(legacy, sig_a) # signature-keyed differs from legacy
self.assertEqual( # same inputs are stable
sig_a, thread_id(self.ticker, self.date, "analysts=market,news|asset=stock")
)
def test_different_signature_starts_fresh(self):
global _should_crash
builder = _build_graph()
sig1 = "analysts=market,news,fundamentals|asset=stock"
sig2 = "analysts=market|asset=stock" # dropped analysts -> different graph
_should_crash = True
tid1 = thread_id(self.ticker, self.date, sig1)
with get_checkpointer(self.tmpdir, self.ticker) as saver:
graph = builder.compile(checkpointer=saver)
with self.assertRaises(RuntimeError):
graph.invoke({"count": 0}, config={"configurable": {"thread_id": tid1}})
self.assertTrue(has_checkpoint(self.tmpdir, self.ticker, self.date, sig1))
# A different graph shape has no checkpoint to resume from.
self.assertFalse(has_checkpoint(self.tmpdir, self.ticker, self.date, sig2))
_should_crash = False
tid2 = thread_id(self.ticker, self.date, sig2)
self.assertNotEqual(tid1, tid2)
with get_checkpointer(self.tmpdir, self.ticker) as saver:
graph = builder.compile(checkpointer=saver)
result = graph.invoke({"count": 0}, config={"configurable": {"thread_id": tid2}})
self.assertEqual(result["count"], 11)
# sig1's checkpoint remains untouched.
self.assertTrue(has_checkpoint(self.tmpdir, self.ticker, self.date, sig1))
def test_run_signature_captures_graph_shape(self):
from tradingagents.graph.trading_graph import TradingAgentsGraph
# Build a bare instance to exercise the pure helper without heavy __init__.
g = object.__new__(TradingAgentsGraph)
g.selected_analysts = ("market", "news")
g.config = {"max_debate_rounds": 1, "max_risk_discuss_rounds": 1}
base = g._run_signature("stock")
self.assertNotEqual(base, g._run_signature("crypto")) # asset mode
g.selected_analysts = ("market",)
self.assertNotEqual(base, g._run_signature("stock")) # analyst selection
g.selected_analysts = ("market", "news")
g.config = {"max_debate_rounds": 3, "max_risk_discuss_rounds": 1}
self.assertNotEqual(base, g._run_signature("stock")) # debate depth
g.config = {"max_debate_rounds": 1, "max_risk_discuss_rounds": 5}
self.assertNotEqual(base, g._run_signature("stock")) # risk depth
# Stable for identical inputs.
g.config = {"max_debate_rounds": 1, "max_risk_discuss_rounds": 1}
self.assertEqual(base, g._run_signature("stock"))
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,100 @@
"""Configurable LLM SDK retry budget (#1090/#1091).
A single transient 429 burst used to kill an otherwise-healthy multi-agent run
because each provider SDK's max_retries (default 2) was not exposed. This adds an
opt-in llm_max_retries knob forwarded to every provider chat client.
"""
from __future__ import annotations
import importlib
import pytest
import tradingagents.default_config as default_config_module
from tradingagents.graph.trading_graph import TradingAgentsGraph, _coerce_max_retries
# --- coercion / validation -------------------------------------------------
@pytest.mark.unit
@pytest.mark.parametrize("value,expected", [(0, 0), (2, 2), (10, 10), ("6", 6)])
def test_coerce_accepts_non_negative_ints_and_numeric_strings(value, expected):
assert _coerce_max_retries(value) == expected
@pytest.mark.unit
@pytest.mark.parametrize("bad", [-1, "-3"])
def test_coerce_rejects_negative(bad):
with pytest.raises(ValueError, match=">= 0"):
_coerce_max_retries(bad)
@pytest.mark.unit
@pytest.mark.parametrize("bad", [True, False])
def test_coerce_rejects_booleans(bad):
with pytest.raises(ValueError, match="boolean"):
_coerce_max_retries(bad)
@pytest.mark.unit
@pytest.mark.parametrize("bad", ["abc", "1.5", None])
def test_coerce_rejects_non_integers(bad):
with pytest.raises(ValueError, match="integer"):
_coerce_max_retries(bad)
# --- forwarding into provider kwargs --------------------------------------
def _bare_graph(config):
g = object.__new__(TradingAgentsGraph)
g.config = config
return g
@pytest.mark.unit
def test_not_forwarded_when_unset():
kwargs = _bare_graph({"llm_provider": "openai", "llm_max_retries": None})._get_provider_kwargs()
assert "max_retries" not in kwargs
@pytest.mark.unit
@pytest.mark.parametrize("provider", ["openai", "anthropic", "google"])
def test_forwarded_across_providers(provider):
kwargs = _bare_graph({"llm_provider": provider, "llm_max_retries": 6})._get_provider_kwargs()
assert kwargs["max_retries"] == 6
@pytest.mark.unit
def test_forwarded_env_string_is_coerced():
# env vars arrive as strings; the consumer coerces (like temperature)
kwargs = _bare_graph({"llm_provider": "openai", "llm_max_retries": "4"})._get_provider_kwargs()
assert kwargs["max_retries"] == 4
@pytest.mark.unit
def test_invalid_config_value_fails_loudly():
with pytest.raises(ValueError):
_bare_graph({"llm_provider": "openai", "llm_max_retries": -1})._get_provider_kwargs()
# --- env overlay -----------------------------------------------------------
def _reload_with_env(monkeypatch, **overrides):
for key in list(default_config_module._ENV_OVERRIDES):
monkeypatch.delenv(key, raising=False)
for key, val in overrides.items():
monkeypatch.setenv(key, val)
return importlib.reload(default_config_module)
@pytest.mark.unit
def test_default_is_none(monkeypatch):
dc = _reload_with_env(monkeypatch)
assert dc.DEFAULT_CONFIG["llm_max_retries"] is None
@pytest.mark.unit
def test_env_override_sets_config(monkeypatch):
dc = _reload_with_env(monkeypatch, TRADINGAGENTS_LLM_MAX_RETRIES="8")
# None-default key: env value arrives as a string and is coerced downstream.
assert dc.DEFAULT_CONFIG["llm_max_retries"] == "8"
assert _coerce_max_retries(dc.DEFAULT_CONFIG["llm_max_retries"]) == 8

View File

@@ -18,6 +18,7 @@ _ENV_OVERRIDES = {
"TRADINGAGENTS_CHECKPOINT_ENABLED": "checkpoint_enabled",
"TRADINGAGENTS_BENCHMARK_TICKER": "benchmark_ticker",
"TRADINGAGENTS_TEMPERATURE": "temperature",
"TRADINGAGENTS_LLM_MAX_RETRIES": "llm_max_retries",
# Provider-specific reasoning/thinking knobs (None = each provider's own
# default). Settable here for non-interactive runs; the CLI also offers an
# interactive choice, which is skipped when the matching var is set.
@@ -95,6 +96,10 @@ DEFAULT_CONFIG = _apply_env_overrides({
# variation on models that honor it; reasoning models largely ignore it
# and no setting makes LLM output bit-identical across runs (see README).
"temperature": None,
# SDK retry budget forwarded to every provider chat client. None leaves each
# provider/SDK at its own default (usually 2). Raise it to ride out bursty
# 429 throttling on rate-limited deployments instead of aborting a run (#1091).
"llm_max_retries": None,
# Checkpoint/resume: when True, LangGraph saves state after each node
# so a crashed run can resume from the last successful step.
"checkpoint_enabled": False,

View File

@@ -25,9 +25,17 @@ def _db_path(data_dir: str | Path, ticker: str) -> Path:
return p / f"{safe}.db"
def thread_id(ticker: str, date: str) -> str:
"""Deterministic thread ID for a ticker+date pair."""
return hashlib.sha256(f"{ticker.upper()}:{date}".encode()).hexdigest()[:16]
def thread_id(ticker: str, date: str, signature: str = "") -> str:
"""Deterministic thread ID for a ticker+date pair.
``signature`` folds in graph-shape-affecting run choices so a resume under a
different graph can't reuse this checkpoint (#1089); omitting it keeps the
legacy ID.
"""
base = f"{ticker.upper()}:{date}"
if signature:
base = f"{base}:{signature}"
return hashlib.sha256(base.encode()).hexdigest()[:16]
@contextmanager
@@ -43,17 +51,17 @@ def get_checkpointer(data_dir: str | Path, ticker: str) -> Generator[SqliteSaver
conn.close()
def has_checkpoint(data_dir: str | Path, ticker: str, date: str) -> bool:
def has_checkpoint(data_dir: str | Path, ticker: str, date: str, signature: str = "") -> bool:
"""Check whether a resumable checkpoint exists for ticker+date."""
return checkpoint_step(data_dir, ticker, date) is not None
return checkpoint_step(data_dir, ticker, date, signature) is not None
def checkpoint_step(data_dir: str | Path, ticker: str, date: str) -> int | None:
def checkpoint_step(data_dir: str | Path, ticker: str, date: str, signature: str = "") -> int | None:
"""Return the step number of the latest checkpoint, or None if none exists."""
db = _db_path(data_dir, ticker)
if not db.exists():
return None
tid = thread_id(ticker, date)
tid = thread_id(ticker, date, signature)
with get_checkpointer(data_dir, ticker) as saver:
config = {"configurable": {"thread_id": tid}}
cp = saver.get_tuple(config)
@@ -73,12 +81,12 @@ def clear_all_checkpoints(data_dir: str | Path) -> int:
return len(dbs)
def clear_checkpoint(data_dir: str | Path, ticker: str, date: str) -> None:
def clear_checkpoint(data_dir: str | Path, ticker: str, date: str, signature: str = "") -> None:
"""Remove checkpoint for a specific ticker+date by deleting the thread's rows."""
db = _db_path(data_dir, ticker)
if not db.exists():
return
tid = thread_id(ticker, date)
tid = thread_id(ticker, date, signature)
conn = sqlite3.connect(str(db))
try:
for table in ("writes", "checkpoints"):

View File

@@ -44,6 +44,24 @@ from .signal_processing import SignalProcessor
logger = logging.getLogger(__name__)
def _coerce_max_retries(value):
"""Validate an ``llm_max_retries`` value to a non-negative int.
Accepts an int or a numeric string (env vars arrive as strings). Rejects
booleans and negatives loudly so a misconfiguration fails at startup rather
than silently disabling retries.
"""
if isinstance(value, bool):
raise ValueError(f"llm_max_retries must be an integer, not a boolean: {value!r}")
try:
n = int(value)
except (TypeError, ValueError) as exc:
raise ValueError(f"llm_max_retries must be an integer, got {value!r}") from exc
if n < 0:
raise ValueError(f"llm_max_retries must be >= 0, got {n}")
return n
class TradingAgentsGraph:
"""Main class that orchestrates the trading agents framework."""
@@ -124,6 +142,9 @@ class TradingAgentsGraph:
self.ticker = None
self.log_states_dict = {} # date to full state dict
# Graph-shape-affecting run choices, kept for the checkpoint signature.
self.selected_analysts = tuple(selected_analysts)
# Set up the graph: keep the workflow for recompilation with a checkpointer.
self.workflow = self.graph_setup.setup_graph(selected_analysts)
self.graph = self.workflow.compile()
@@ -156,6 +177,12 @@ class TradingAgentsGraph:
if temperature is not None and temperature != "":
kwargs["temperature"] = float(temperature)
# SDK retry budget is cross-provider. Forward it only when explicitly set
# so each provider keeps its own default (usually 2) otherwise (#1091).
max_retries = self.config.get("llm_max_retries")
if max_retries is not None and max_retries != "":
kwargs["max_retries"] = _coerce_max_retries(max_retries)
return kwargs
def _create_tool_nodes(self) -> dict[str, ToolNode]:
@@ -318,6 +345,20 @@ class TradingAgentsGraph:
identity = resolve_instrument_identity(ticker)
return build_instrument_context(ticker, asset_type, identity)
def _run_signature(self, asset_type: str) -> str:
"""Graph-shape inputs that must invalidate a checkpoint if changed.
Keyed into the checkpoint thread ID so a resume under a different analyst
selection, debate/risk depth, or asset mode starts fresh instead of
silently continuing the previous graph (#1089).
"""
return "|".join([
"analysts=" + ",".join(self.selected_analysts),
f"debate={self.config['max_debate_rounds']}",
f"risk={self.config['max_risk_discuss_rounds']}",
f"asset={asset_type}",
])
def propagate(self, company_name, trade_date, asset_type: str = "stock"):
"""Run the trading agents graph for a company on a specific date.
@@ -342,7 +383,8 @@ class TradingAgentsGraph:
self.graph = self.workflow.compile(checkpointer=saver)
step = checkpoint_step(
self.config["data_cache_dir"], company_name, str(trade_date)
self.config["data_cache_dir"], company_name, str(trade_date),
self._run_signature(asset_type),
)
if step is not None:
logger.info(
@@ -389,9 +431,10 @@ class TradingAgentsGraph:
)
args = self.propagator.get_graph_args()
# Inject thread_id so same ticker+date resumes, different date starts fresh.
# Inject thread_id so same ticker+date+graph-shape resumes; a different
# date or graph shape starts fresh (#1089).
if self.config.get("checkpoint_enabled"):
tid = thread_id(company_name, str(trade_date))
tid = thread_id(company_name, str(trade_date), self._run_signature(asset_type))
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = tid
if self.debug:
@@ -432,7 +475,8 @@ class TradingAgentsGraph:
# Clear checkpoint on successful completion to avoid stale state.
if self.config.get("checkpoint_enabled"):
clear_checkpoint(
self.config["data_cache_dir"], company_name, str(trade_date)
self.config["data_cache_dir"], company_name, str(trade_date),
self._run_signature(asset_type),
)
return final_state, self.process_signal(final_state["final_trade_decision"])