From 96daaf1152334f3465e01ed45f9f99e8085da895 Mon Sep 17 00:00:00 2001 From: Yijia-Xiao Date: Wed, 23 Sep 2026 19:28:56 +0000 Subject: [PATCH] fix(dataflows): read the vendors of the run in progress (#1369) - propagate and settle_pending bind the graph's config for the length of the run - a graph built later, or running concurrently, no longer changes another graph's vendors --- tests/test_dataflows_config.py | 139 +++++++++++++++++++++++++++ tradingagents/dataflows/config.py | 40 ++++++-- tradingagents/graph/trading_graph.py | 8 +- 3 files changed, 176 insertions(+), 11 deletions(-) diff --git a/tests/test_dataflows_config.py b/tests/test_dataflows_config.py index ab0800eee..480793782 100644 --- a/tests/test_dataflows_config.py +++ b/tests/test_dataflows_config.py @@ -59,3 +59,142 @@ class DataflowsConfigIsolationTests(unittest.TestCase): fresh = get_config() self.assertEqual(fresh["tool_vendors"]["get_stock_data"], "alpha_vantage") self.assertEqual(fresh["tool_vendors"]["get_news"], "alpha_vantage") + + +# --- the config of the run in progress (#1369) -------------------------------- + +def _graph(config): + from tradingagents.graph.trading_graph import TradingAgentsGraph + + g = object.__new__(TradingAgentsGraph) + g.config = config + g._checkpointer_ctx = None + return g + + +def _vendors_seen_by_a_run(graph, ticker="AAPL"): + from tradingagents.dataflows.interface import get_vendor + + seen = [] + + def _run(*a, **k): + seen.append(get_vendor("fundamental_data", "get_balance_sheet")) + return {}, "Hold" + + graph._run_graph = _run + graph.propagate(ticker, "2026-09-01") + return seen + + +@pytest.mark.unit +def test_a_run_reads_its_own_graphs_vendors_not_the_last_graph_built(): + """Building a graph sets the process-wide config, and set_config merges, so + a second graph built with the defaults was served the first one's vendors.""" + first = copy.deepcopy(default_config.DEFAULT_CONFIG) + first["tool_vendors"] = {"get_balance_sheet": "sec_edgar,yfinance"} + set_config(first) # graph A is built + second = _graph(copy.deepcopy(default_config.DEFAULT_CONFIG)) + + assert _vendors_seen_by_a_run(second) == ["yfinance"] + + +@pytest.mark.unit +def test_a_graph_built_earlier_still_runs_with_its_own_config(): + """Scoping at construction would hand graph A graph B's config if B was built + after A; the config must be bound when the run starts.""" + a_config = copy.deepcopy(default_config.DEFAULT_CONFIG) + a_config["tool_vendors"] = {"get_balance_sheet": "sec_edgar,yfinance"} + a = _graph(a_config) + set_config(copy.deepcopy(default_config.DEFAULT_CONFIG)) # graph B is built + + assert _vendors_seen_by_a_run(a) == ["sec_edgar,yfinance"] + + +@pytest.mark.unit +def test_concurrent_runs_each_read_their_own_config(): + import threading + + barrier = threading.Barrier(2) + results = {} + + def run(name, vendor): + config = copy.deepcopy(default_config.DEFAULT_CONFIG) + config["tool_vendors"] = {"get_balance_sheet": vendor} + graph = _graph(config) + from tradingagents.dataflows.interface import get_vendor + + def _run(*a, **k): + barrier.wait(timeout=5) # both runs are in flight + results[name] = get_vendor("fundamental_data", "get_balance_sheet") + return {}, "Hold" + + graph._run_graph = _run + graph.propagate("AAPL", "2026-09-01") + + threads = [threading.Thread(target=run, args=("a", "alpha_vantage")), + threading.Thread(target=run, args=("b", "sec_edgar,yfinance"))] + [t.start() for t in threads] + [t.join() for t in threads] + + assert results == {"a": "alpha_vantage", "b": "sec_edgar,yfinance"} + + +@pytest.mark.unit +def test_settling_reads_the_graphs_own_config(): + from tradingagents.dataflows.interface import get_vendor + + config = copy.deepcopy(default_config.DEFAULT_CONFIG) + config["tool_vendors"] = {"get_stock_data": "alpha_vantage"} + graph = _graph(config) + seen = [] + graph._resolve_pending_entries = lambda ticker: seen.append( + get_vendor("core_stock_apis", "get_stock_data")) + + graph.settle_pending("AAPL") + + assert seen == ["alpha_vantage"] + + +@pytest.mark.unit +def test_tools_inside_a_langgraph_run_see_the_run_config(): + """The fix rests on LangGraph carrying the caller's context into tool calls.""" + from langchain_core.messages import AIMessage + from langchain_core.tools import tool + from langgraph.graph import END, START, MessagesState, StateGraph + from langgraph.prebuilt import ToolNode + + from tradingagents.dataflows.config import run_config + from tradingagents.dataflows.interface import get_vendor + + @tool + def probe() -> str: + """Report the vendor the run would use.""" + return get_vendor("fundamental_data", "get_balance_sheet") + + def call(state): + return {"messages": [AIMessage("", tool_calls=[{"name": "probe", "args": {}, "id": "1"}])]} + + g = StateGraph(MessagesState) + g.add_node("call", call) + g.add_node("tools", ToolNode([probe])) + g.add_edge(START, "call") + g.add_edge("call", "tools") + g.add_edge("tools", END) + + config = copy.deepcopy(default_config.DEFAULT_CONFIG) + config["tool_vendors"] = {"get_balance_sheet": "sec_edgar,yfinance"} + with run_config(config): + out = g.compile().invoke({"messages": [("user", "go")]}) + + assert out["messages"][-1].content == "sec_edgar,yfinance" + + +@pytest.mark.unit +def test_a_run_config_missing_a_newer_key_still_reads_the_default(): + """A config saved before a key existed must not fail inside a run.""" + from tradingagents.dataflows.config import get_config, run_config + + config = copy.deepcopy(default_config.DEFAULT_CONFIG) + del config["news_article_limit"] + with run_config(config): + assert get_config()["news_article_limit"] == default_config.DEFAULT_CONFIG["news_article_limit"] diff --git a/tradingagents/dataflows/config.py b/tradingagents/dataflows/config.py index c4790e847..27d6785c6 100644 --- a/tradingagents/dataflows/config.py +++ b/tradingagents/dataflows/config.py @@ -1,3 +1,5 @@ +from contextlib import contextmanager +from contextvars import ContextVar from copy import deepcopy import tradingagents.default_config as default_config @@ -5,6 +7,11 @@ import tradingagents.default_config as default_config # Use default config but allow it to be overridden _config: dict | None = None +# The config of the run in progress. A graph binds its own for the length of a +# run, so the data tools it calls read that graph's vendors even when several +# graphs share a process. LangGraph carries the context into tool calls. +_run_config: ContextVar[dict | None] = ContextVar("tradingagents_run_config", default=None) + def initialize_config(): """Initialize the configuration with default values.""" @@ -13,6 +20,16 @@ def initialize_config(): _config = deepcopy(default_config.DEFAULT_CONFIG) +def _merge(base: dict, config: dict) -> dict: + """Merge ``config`` into ``base``: dict-valued keys one level deep, scalars replaced.""" + for key, value in deepcopy(config).items(): + if isinstance(value, dict) and isinstance(base.get(key), dict): + base[key].update(value) + else: + base[key] = value + return base + + def set_config(config: dict): """Update the configuration with custom values. @@ -20,18 +37,25 @@ def set_config(config: dict): partial update like ``{"data_vendors": {"core_stock_apis": "alpha_vantage"}}`` keeps the other nested keys from the default; scalar keys are replaced. """ - global _config initialize_config() - incoming = deepcopy(config) - for key, value in incoming.items(): - if isinstance(value, dict) and isinstance(_config.get(key), dict): - _config[key].update(value) - else: - _config[key] = value + _merge(_config, config) + + +@contextmanager +def run_config(config: dict): + """Serve ``config``, over the defaults, to every read made inside the block.""" + token = _run_config.set(_merge(deepcopy(default_config.DEFAULT_CONFIG), config)) + try: + yield + finally: + _run_config.reset(token) def get_config() -> dict: - """Get the current configuration.""" + """Get the configuration of the run in progress, else the process-wide one.""" + scoped = _run_config.get() + if scoped is not None: + return deepcopy(scoped) if _config is None: initialize_config() return deepcopy(_config) diff --git a/tradingagents/graph/trading_graph.py b/tradingagents/graph/trading_graph.py index a3510ddc2..8495c32d7 100644 --- a/tradingagents/graph/trading_graph.py +++ b/tradingagents/graph/trading_graph.py @@ -29,7 +29,7 @@ from tradingagents.agents.utils.agent_utils import ( resolve_instrument_identity, ) from tradingagents.agents.utils.memory import TradingMemoryLog -from tradingagents.dataflows.config import set_config +from tradingagents.dataflows.config import run_config, set_config from tradingagents.dataflows.utils import get_current_date, safe_ticker_component from tradingagents.default_config import DEFAULT_CONFIG from tradingagents.llm_clients import create_llm_client @@ -452,7 +452,8 @@ class TradingAgentsGraph: trade_date = _validate_trade_date(trade_date) self.ticker = company_name - with self.checkpoint_scope(company_name, trade_date, asset_type, portfolio) as thread_id_value: + with run_config(self.config), \ + self.checkpoint_scope(company_name, trade_date, asset_type, portfolio) as thread_id_value: return self._run_graph( company_name, trade_date, asset_type=asset_type, checkpoint_thread_id=thread_id_value, portfolio=portfolio, @@ -564,7 +565,8 @@ class TradingAgentsGraph: that is done analyzing a ticker (a backtest sweep, a scheduled job) calls this to settle it now. """ - self._resolve_pending_entries(company_name) + with run_config(self.config): + self._resolve_pending_entries(company_name) def record_decision(self, company_name, trade_date, final_state): """Log a finished run's decision for reflection on the next same-ticker run."""