mirror of
https://github.com/TauricResearch/TradingAgents.git
synced 2026-09-27 15:02:39 +03:00
feat(graph): run the analysts at the same time (#1255)
- each analyst is a graph of its own (model and tools on a private message history) that returns only its report; all start together and the research debate waits for every report - the message-clearing nodes are gone; a checkpoint saved by the sequential layout starts fresh - TradingAgentsGraph.stream_run streams the analysts' messages for debug mode and the CLI, whose status and timing now track the analysts side by side
This commit is contained in:
@@ -11,8 +11,8 @@ class AnalystExecutionPlanTests(unittest.TestCase):
|
||||
|
||||
self.assertEqual([spec.key for spec in plan.specs], ["news", "market"])
|
||||
self.assertEqual(plan.specs[0].agent_node, "News Analyst")
|
||||
self.assertEqual(plan.specs[0].tool_node, "tools_news")
|
||||
self.assertEqual(plan.specs[0].clear_node, "Msg Clear News")
|
||||
self.assertEqual(plan.specs[0].report_key, "news_report")
|
||||
self.assertFalse(hasattr(plan.specs[0], "clear_node"))
|
||||
|
||||
def test_rejects_unknown_analyst_keys(self):
|
||||
with self.assertRaises(ValueError):
|
||||
|
||||
@@ -210,6 +210,10 @@ class TestCheckpointSignature(unittest.TestCase):
|
||||
# Stable for identical inputs.
|
||||
g.config = {"max_debate_rounds": 1, "max_risk_discuss_rounds": 1}
|
||||
self.assertEqual(base, g._run_signature("stock"))
|
||||
# A checkpoint saved by the sequential layout is not resumed on the
|
||||
# parallel one: its pending node no longer exists, and the join would
|
||||
# never fire.
|
||||
self.assertIn("analysts=parallel", base)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
+18
-13
@@ -129,23 +129,28 @@ class AnalystWallTimeTrackerTests(unittest.TestCase):
|
||||
"Analyst wall time: News 4.00s | Market 2.25s",
|
||||
)
|
||||
|
||||
def test_syncs_wall_time_from_sequential_chunks(self):
|
||||
def test_analysts_run_together_and_finish_on_their_own_reports(self):
|
||||
plan = build_analyst_execution_plan(["market", "news"])
|
||||
tracker = AnalystWallTimeTracker(plan)
|
||||
|
||||
sync_analyst_tracker_from_chunk(tracker, {}, now=10.0)
|
||||
self.assertEqual(tracker.format_summary(), "Analyst wall time: pending")
|
||||
|
||||
sync_analyst_tracker_from_chunk(
|
||||
tracker,
|
||||
{"market_report": "done"},
|
||||
now=13.0,
|
||||
)
|
||||
self.assertEqual(tracker.format_summary(), "Analyst wall time: Market 3.00s")
|
||||
sync_analyst_tracker_from_chunk(tracker, {"news_report": "done"}, now=13.0)
|
||||
self.assertEqual(tracker.format_summary(), "Analyst wall time: News 3.00s")
|
||||
|
||||
sync_analyst_tracker_from_chunk(
|
||||
tracker,
|
||||
{"market_report": "done", "news_report": "done"},
|
||||
now=18.0,
|
||||
)
|
||||
self.assertEqual(tracker.format_summary(), "Analyst wall time: Market 3.00s | News 5.00s")
|
||||
sync_analyst_tracker_from_chunk(tracker, {"news_report": "done", "market_report": "done"}, now=18.0)
|
||||
self.assertEqual(tracker.format_summary(), "Analyst wall time: Market 8.00s | News 3.00s")
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_every_selected_analyst_is_in_progress_until_its_report_lands():
|
||||
from cli.display import MessageBuffer, update_analyst_statuses
|
||||
|
||||
buffer = MessageBuffer()
|
||||
buffer.init_for_analysis(["market", "news", "fundamentals"])
|
||||
update_analyst_statuses(buffer, {"news_report": "done"})
|
||||
|
||||
assert buffer.agent_status["Market Analyst"] == "in_progress"
|
||||
assert buffer.agent_status["Fundamentals Analyst"] == "in_progress"
|
||||
assert buffer.agent_status["News Analyst"] == "completed"
|
||||
|
||||
@@ -95,9 +95,9 @@ class _FakeGraph:
|
||||
def end_checkpoint(self):
|
||||
pass
|
||||
|
||||
def stream(self, graph_input, **kwargs):
|
||||
yield {"messages": [], "market_report": "M"}
|
||||
yield {"messages": [], "final_trade_decision": "Rating: Buy\n\nBuy NVDA.", "final_rating": "Buy"}
|
||||
def stream_run(self, graph_input, **kwargs):
|
||||
yield [], {"messages": [], "market_report": "M"}
|
||||
yield [], {"messages": [], "final_trade_decision": "Rating: Buy\n\nBuy NVDA.", "final_rating": "Buy"}
|
||||
|
||||
|
||||
class _NullLive:
|
||||
|
||||
@@ -52,6 +52,7 @@ class ScriptedModel(BaseChatModel):
|
||||
structured: bool = False
|
||||
tools: tuple = ()
|
||||
calls: list = Field(default_factory=list) # shared across bound copies
|
||||
threads: set = Field(default_factory=set) # threads that served a tool-bound call
|
||||
fail_at: int | None = None # raise on this call, once
|
||||
|
||||
@property
|
||||
@@ -73,6 +74,11 @@ class ScriptedModel(BaseChatModel):
|
||||
|
||||
def _generate(self, messages, stop=None, run_manager=None, **kwargs) -> ChatResult:
|
||||
self._count()
|
||||
if self.tools:
|
||||
import threading
|
||||
import time
|
||||
self.threads.add(threading.current_thread().name)
|
||||
time.sleep(0.05) # long enough for concurrent analysts to overlap
|
||||
if self.tools and not isinstance(messages[-1], ToolMessage):
|
||||
calls = [{"name": t.name, "id": f"call_{i}",
|
||||
"args": {k: v for k, v in ARGS.items()
|
||||
@@ -113,12 +119,12 @@ def offline(monkeypatch, tmp_path):
|
||||
return called
|
||||
|
||||
|
||||
def _graph(tmp_path, monkeypatch, model, **config):
|
||||
def _graph(tmp_path, monkeypatch, model, debug=False, **config):
|
||||
cfg = copy.deepcopy(DEFAULT_CONFIG)
|
||||
cfg.update(results_dir=str(tmp_path / "results"), data_cache_dir=str(tmp_path / "cache"),
|
||||
memory_log_path=str(tmp_path / "log.md"), **config)
|
||||
monkeypatch.setattr(trading_graph, "create_llm_client", lambda **k: _Client(model))
|
||||
return trading_graph.TradingAgentsGraph(config=cfg)
|
||||
return trading_graph.TradingAgentsGraph(config=cfg, debug=debug)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
@@ -169,3 +175,41 @@ def test_a_graph_reused_across_runs_keeps_no_run_state(tmp_path, monkeypatch, of
|
||||
held = [v for v in vars(graph).values() if isinstance(v, dict) and TRADE_DATE in v]
|
||||
assert held == []
|
||||
assert len(list(tmp_path.glob("results/NVDA/TradingAgentsStrategy_logs/*.json"))) == 2
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_the_analysts_run_at_the_same_time(tmp_path, monkeypatch, offline):
|
||||
model = ScriptedModel()
|
||||
graph = _graph(tmp_path, monkeypatch, model)
|
||||
|
||||
assert not [n for n in graph.graph.get_graph().nodes if n.startswith("Msg Clear")]
|
||||
graph.propagate("NVDA", TRADE_DATE)
|
||||
|
||||
assert len(model.threads) > 1
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_a_debug_run_prints_the_analysts_work_and_reaches_the_same_decision(tmp_path, monkeypatch, offline, capsys):
|
||||
"""Debug mode streams the analysts' own graphs, so their tool calls still print."""
|
||||
graph = _graph(tmp_path, monkeypatch, ScriptedModel(), debug=True)
|
||||
|
||||
state, signal = graph.propagate("NVDA", TRADE_DATE)
|
||||
|
||||
assert signal == "Overweight"
|
||||
assert state["market_report"].strip() and state["fundamentals_report"].strip()
|
||||
printed = capsys.readouterr().out
|
||||
assert "get_stock_data" in printed and "get_balance_sheet" in printed
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
def test_each_report_streams_as_soon_as_its_analyst_files_it(tmp_path, monkeypatch, offline):
|
||||
"""The main state takes the analysts' reports only when the slowest one is
|
||||
done; the CLI shows each report, and stops each clock, as it lands."""
|
||||
graph = _graph(tmp_path, monkeypatch, ScriptedModel())
|
||||
reports = ("market_report", "sentiment_report", "news_report", "fundamentals_report")
|
||||
|
||||
first = next(state for _, state in graph.stream_run(graph.create_run_state("NVDA", TRADE_DATE),
|
||||
**graph.propagator.get_graph_args())
|
||||
if state and any(state.get(k) for k in reports))
|
||||
|
||||
assert sum(bool(first.get(k)) for k in reports) == 1
|
||||
|
||||
@@ -5,11 +5,9 @@ import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
||||
|
||||
from tradingagents.agents.context import (
|
||||
build_instrument_context,
|
||||
create_msg_delete,
|
||||
get_instrument_context_from_state,
|
||||
resolve_instrument_identity,
|
||||
)
|
||||
@@ -118,53 +116,5 @@ class GetInstrumentContextFromStateTests(unittest.TestCase):
|
||||
self.assertIn("crypto asset", context)
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class ContextAnchoredPlaceholderTests(unittest.TestCase):
|
||||
"""#888 — the message-clear placeholder must not be a bare 'Continue'."""
|
||||
|
||||
def _run(self, state_extra):
|
||||
state = {
|
||||
"messages": [
|
||||
HumanMessage(content="old", id="h1"),
|
||||
AIMessage(content="reply", id="a1"),
|
||||
],
|
||||
**state_extra,
|
||||
}
|
||||
return create_msg_delete()(state)
|
||||
|
||||
def test_placeholder_is_not_bare_continue(self):
|
||||
result = self._run(
|
||||
{"company_of_interest": "EC", "asset_type": "stock", "trade_date": "2026-05-28"}
|
||||
)
|
||||
placeholder = result["messages"][-1]
|
||||
self.assertIsInstance(placeholder, HumanMessage)
|
||||
self.assertNotEqual(placeholder.content.strip(), "Continue")
|
||||
|
||||
def test_placeholder_carries_resolved_identity(self):
|
||||
result = self._run(
|
||||
{
|
||||
"company_of_interest": "EC",
|
||||
"instrument_context": "The instrument to analyze is `EC`. Resolved identity: Company: Ecopetrol.",
|
||||
"trade_date": "2026-05-28",
|
||||
}
|
||||
)
|
||||
content = result["messages"][-1].content
|
||||
self.assertIn("Ecopetrol", content)
|
||||
self.assertIn("2026-05-28", content)
|
||||
|
||||
def test_old_messages_are_removed(self):
|
||||
result = self._run({"company_of_interest": "EC", "trade_date": "2026-05-28"})
|
||||
removals = [m for m in result["messages"] if isinstance(m, RemoveMessage)]
|
||||
humans = [m for m in result["messages"] if isinstance(m, HumanMessage)]
|
||||
self.assertEqual(len(removals), 2)
|
||||
self.assertEqual(len(humans), 1)
|
||||
|
||||
def test_safe_defaults_when_state_minimal(self):
|
||||
result = create_msg_delete()({"messages": [], "company_of_interest": "EC"})
|
||||
placeholder = result["messages"][-1]
|
||||
self.assertNotEqual(placeholder.content.strip(), "Continue")
|
||||
self.assertIn("EC", placeholder.content)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -136,8 +136,8 @@ def test_the_cli_says_when_a_run_produced_no_usable_rating(monkeypatch, tmp_path
|
||||
def end_checkpoint(self):
|
||||
pass
|
||||
|
||||
def stream(self, *a, **k):
|
||||
yield {"messages": [], "final_trade_decision": REFUSAL, "final_rating": RATING_REVIEW}
|
||||
def stream_run(self, *a, **k):
|
||||
yield [], {"messages": [], "final_trade_decision": REFUSAL, "final_rating": RATING_REVIEW}
|
||||
|
||||
fake = _Graph()
|
||||
fake.graph = fake
|
||||
|
||||
Reference in New Issue
Block a user