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 committed 2026-07-05 14:29:07 +00:00
1 parent b47a828a4f
commit daf1da9c35
5 files changed
+244 -13

No files matched your search

+48 -4
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"])