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
This commit is contained in:
Yijia-Xiao
2026-09-23 19:28:56 +00:00
parent a9cc3be731
commit 96daaf1152
3 changed files with 176 additions and 11 deletions
+139
View File
@@ -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"]
+32 -8
View File
@@ -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)
+4 -2
View File
@@ -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,6 +565,7 @@ class TradingAgentsGraph:
that is done analyzing a ticker (a backtest sweep, a scheduled job) calls
this to settle it now.
"""
with run_config(self.config):
self._resolve_pending_entries(company_name)
def record_decision(self, company_name, trade_date, final_state):