From daf1da9c351fb333c4af6b9efb041ab591187dd1 Mon Sep 17 00:00:00 2001 From: Yijia-Xiao Date: Sun, 5 Jul 2026 14:29:07 +0000 Subject: [PATCH] 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 --- tests/test_checkpoint_resume.py | 74 ++++++++++++++++++++ tests/test_llm_max_retries.py | 100 +++++++++++++++++++++++++++ tradingagents/default_config.py | 5 ++ tradingagents/graph/checkpointer.py | 26 ++++--- tradingagents/graph/trading_graph.py | 52 ++++++++++++-- 5 files changed, 244 insertions(+), 13 deletions(-) create mode 100644 tests/test_llm_max_retries.py diff --git a/tests/test_checkpoint_resume.py b/tests/test_checkpoint_resume.py index bf801b073..6c5be500c 100644 --- a/tests/test_checkpoint_resume.py +++ b/tests/test_checkpoint_resume.py @@ -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() diff --git a/tests/test_llm_max_retries.py b/tests/test_llm_max_retries.py new file mode 100644 index 000000000..632060a5e --- /dev/null +++ b/tests/test_llm_max_retries.py @@ -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 diff --git a/tradingagents/default_config.py b/tradingagents/default_config.py index a25d0f76b..ab75f7508 100644 --- a/tradingagents/default_config.py +++ b/tradingagents/default_config.py @@ -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, diff --git a/tradingagents/graph/checkpointer.py b/tradingagents/graph/checkpointer.py index a32cb4cb2..750d8d12a 100644 --- a/tradingagents/graph/checkpointer.py +++ b/tradingagents/graph/checkpointer.py @@ -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"): diff --git a/tradingagents/graph/trading_graph.py b/tradingagents/graph/trading_graph.py index ece0513e1..a1902177d 100644 --- a/tradingagents/graph/trading_graph.py +++ b/tradingagents/graph/trading_graph.py @@ -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"])