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:
Yijia-Xiao
2026-09-25 06:35:54 +00:00
parent fc1ab1db07
commit 9968bd8dd1
15 changed files with 163 additions and 193 deletions
+2
View File
@@ -81,6 +81,8 @@ Our framework decomposes complex trading tasks into specialized roles.
- News Analyst: Monitors global news and macroeconomic indicators, interpreting the impact of events on market conditions. - News Analyst: Monitors global news and macroeconomic indicators, interpreting the impact of events on market conditions.
- Technical Analyst: Utilizes technical indicators (like MACD and RSI) to detect trading patterns and forecast price movements. - Technical Analyst: Utilizes technical indicators (like MACD and RSI) to detect trading patterns and forecast price movements.
The selected analysts work at the same time, each on its own tools, and the research debate starts once all of their reports are in.
<p align="center"> <p align="center">
<img src="assets/analyst.png" width="100%" style="display: inline-block; margin: 0 2%;"> <img src="assets/analyst.png" width="100%" style="display: inline-block; margin: 0 2%;">
</p> </p>
+12 -30
View File
@@ -463,22 +463,17 @@ ANALYST_REPORT_MAP = {
def update_analyst_statuses(message_buffer, chunk, wall_time_tracker=None): def update_analyst_statuses(message_buffer, chunk, wall_time_tracker=None):
"""Update analyst statuses based on accumulated report state. """Update analyst statuses from the reports filed so far.
Logic: The analysts run together: each is in progress until its own report lands.
- Store new report content from the current chunk if present When every selected analyst has filed, the research debate is in progress.
- Check accumulated report_sections (not just current chunk) for status
- Analysts with reports = completed
- First analyst without report = in_progress
- Remaining analysts without reports = pending
- When all analysts done, set Bull Researcher to in_progress
""" """
selected = message_buffer.selected_analysts selected = message_buffer.selected_analysts
found_active = False
if wall_time_tracker is not None: if wall_time_tracker is not None:
sync_analyst_tracker_from_chunk(wall_time_tracker, chunk) sync_analyst_tracker_from_chunk(wall_time_tracker, chunk)
all_filed = True
for analyst_key in ANALYST_ORDER: for analyst_key in ANALYST_ORDER:
if analyst_key not in selected: if analyst_key not in selected:
continue continue
@@ -490,20 +485,15 @@ def update_analyst_statuses(message_buffer, chunk, wall_time_tracker=None):
if chunk.get(report_key): if chunk.get(report_key):
message_buffer.update_report_section(report_key, chunk[report_key]) message_buffer.update_report_section(report_key, chunk[report_key])
# Determine status from accumulated sections, not just current chunk # Status comes from accumulated sections, not just the current chunk.
has_report = bool(message_buffer.report_sections.get(report_key)) if message_buffer.report_sections.get(report_key):
if has_report:
message_buffer.update_agent_status(agent_name, "completed") message_buffer.update_agent_status(agent_name, "completed")
elif not found_active:
message_buffer.update_agent_status(agent_name, "in_progress")
found_active = True
else: else:
message_buffer.update_agent_status(agent_name, "pending") message_buffer.update_agent_status(agent_name, "in_progress")
all_filed = False
# When all analysts complete, transition research team to in_progress
if ( if (
not found_active all_filed
and selected and selected
and message_buffer.agent_status.get("Bull Researcher") == "pending" and message_buffer.agent_status.get("Bull Researcher") == "pending"
): ):
@@ -624,17 +614,9 @@ def sync_analyst_tracker_from_chunk(
chunk: dict[str, str], chunk: dict[str, str],
now: float | None = None, now: float | None = None,
) -> None: ) -> None:
"""The analysts start together; each stops its clock when its report lands."""
current_time = monotonic() if now is None else now current_time = monotonic() if now is None else now
active_found = False
for spec in tracker.plan.specs: for spec in tracker.plan.specs:
has_report = bool(chunk.get(spec.report_key)) tracker.mark_started(spec.key, started_at=current_time)
if chunk.get(spec.report_key):
if has_report:
tracker.mark_started(spec.key, started_at=current_time)
tracker.mark_completed(spec.key, completed_at=current_time) tracker.mark_completed(spec.key, completed_at=current_time)
continue
if not active_found:
tracker.mark_started(spec.key, started_at=current_time)
active_found = True
+10 -5
View File
@@ -200,9 +200,10 @@ def run_analysis(checkpoint: bool | None = None, portfolio=None, flags=None):
) )
update_display(layout, stats_handler=stats_handler, start_time=start_time) update_display(layout, stats_handler=stats_handler, start_time=start_time)
first_analyst = analyst_execution_plan.specs[0].agent_node # The analysts start together.
message_buffer.update_agent_status(first_analyst, "in_progress") for spec in analyst_execution_plan.specs:
analyst_wall_time_tracker.mark_started(selected_analyst_keys[0]) message_buffer.update_agent_status(spec.agent_node, "in_progress")
analyst_wall_time_tracker.mark_started(spec.key)
update_display(layout, stats_handler=stats_handler, start_time=start_time) update_display(layout, stats_handler=stats_handler, start_time=start_time)
spinner_text = ( spinner_text = (
@@ -234,8 +235,8 @@ def run_analysis(checkpoint: bool | None = None, portfolio=None, flags=None):
# try/finally tears the checkpointer down even if the stream raises. # try/finally tears the checkpointer down even if the stream raises.
trace = [] trace = []
try: try:
for chunk in graph.graph.stream(graph.checkpoint_input(init_agent_state), **args): for messages, chunk in graph.stream_run(graph.checkpoint_input(init_agent_state), **args):
for message in chunk.get("messages", []): for message in messages:
msg_id = getattr(message, "id", None) msg_id = getattr(message, "id", None)
if msg_id is not None: if msg_id is not None:
if msg_id in message_buffer._processed_message_ids: if msg_id in message_buffer._processed_message_ids:
@@ -253,6 +254,10 @@ def run_analysis(checkpoint: bool | None = None, portfolio=None, flags=None):
else: else:
message_buffer.add_tool_call(tool_call.name, tool_call.args) message_buffer.add_tool_call(tool_call.name, tool_call.args)
if chunk is None: # a step inside an analyst's graph: messages only
update_display(layout, stats_handler=stats_handler, start_time=start_time)
continue
update_analyst_statuses( update_analyst_statuses(
message_buffer, message_buffer,
chunk, chunk,
+2 -2
View File
@@ -11,8 +11,8 @@ class AnalystExecutionPlanTests(unittest.TestCase):
self.assertEqual([spec.key for spec in plan.specs], ["news", "market"]) 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].agent_node, "News Analyst")
self.assertEqual(plan.specs[0].tool_node, "tools_news") self.assertEqual(plan.specs[0].report_key, "news_report")
self.assertEqual(plan.specs[0].clear_node, "Msg Clear News") self.assertFalse(hasattr(plan.specs[0], "clear_node"))
def test_rejects_unknown_analyst_keys(self): def test_rejects_unknown_analyst_keys(self):
with self.assertRaises(ValueError): with self.assertRaises(ValueError):
+4
View File
@@ -210,6 +210,10 @@ class TestCheckpointSignature(unittest.TestCase):
# Stable for identical inputs. # Stable for identical inputs.
g.config = {"max_debate_rounds": 1, "max_risk_discuss_rounds": 1} g.config = {"max_debate_rounds": 1, "max_risk_discuss_rounds": 1}
self.assertEqual(base, g._run_signature("stock")) 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__": if __name__ == "__main__":
+18 -13
View File
@@ -129,23 +129,28 @@ class AnalystWallTimeTrackerTests(unittest.TestCase):
"Analyst wall time: News 4.00s | Market 2.25s", "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"]) plan = build_analyst_execution_plan(["market", "news"])
tracker = AnalystWallTimeTracker(plan) tracker = AnalystWallTimeTracker(plan)
sync_analyst_tracker_from_chunk(tracker, {}, now=10.0) sync_analyst_tracker_from_chunk(tracker, {}, now=10.0)
self.assertEqual(tracker.format_summary(), "Analyst wall time: pending") self.assertEqual(tracker.format_summary(), "Analyst wall time: pending")
sync_analyst_tracker_from_chunk( sync_analyst_tracker_from_chunk(tracker, {"news_report": "done"}, now=13.0)
tracker, self.assertEqual(tracker.format_summary(), "Analyst wall time: News 3.00s")
{"market_report": "done"},
now=13.0,
)
self.assertEqual(tracker.format_summary(), "Analyst wall time: Market 3.00s")
sync_analyst_tracker_from_chunk( sync_analyst_tracker_from_chunk(tracker, {"news_report": "done", "market_report": "done"}, now=18.0)
tracker, self.assertEqual(tracker.format_summary(), "Analyst wall time: Market 8.00s | News 3.00s")
{"market_report": "done", "news_report": "done"},
now=18.0,
) @pytest.mark.unit
self.assertEqual(tracker.format_summary(), "Analyst wall time: Market 3.00s | News 5.00s") 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"
+3 -3
View File
@@ -95,9 +95,9 @@ class _FakeGraph:
def end_checkpoint(self): def end_checkpoint(self):
pass pass
def stream(self, graph_input, **kwargs): def stream_run(self, graph_input, **kwargs):
yield {"messages": [], "market_report": "M"} yield [], {"messages": [], "market_report": "M"}
yield {"messages": [], "final_trade_decision": "Rating: Buy\n\nBuy NVDA.", "final_rating": "Buy"} yield [], {"messages": [], "final_trade_decision": "Rating: Buy\n\nBuy NVDA.", "final_rating": "Buy"}
class _NullLive: class _NullLive:
+46 -2
View File
@@ -52,6 +52,7 @@ class ScriptedModel(BaseChatModel):
structured: bool = False structured: bool = False
tools: tuple = () tools: tuple = ()
calls: list = Field(default_factory=list) # shared across bound copies 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 fail_at: int | None = None # raise on this call, once
@property @property
@@ -73,6 +74,11 @@ class ScriptedModel(BaseChatModel):
def _generate(self, messages, stop=None, run_manager=None, **kwargs) -> ChatResult: def _generate(self, messages, stop=None, run_manager=None, **kwargs) -> ChatResult:
self._count() 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): if self.tools and not isinstance(messages[-1], ToolMessage):
calls = [{"name": t.name, "id": f"call_{i}", calls = [{"name": t.name, "id": f"call_{i}",
"args": {k: v for k, v in ARGS.items() "args": {k: v for k, v in ARGS.items()
@@ -113,12 +119,12 @@ def offline(monkeypatch, tmp_path):
return called return called
def _graph(tmp_path, monkeypatch, model, **config): def _graph(tmp_path, monkeypatch, model, debug=False, **config):
cfg = copy.deepcopy(DEFAULT_CONFIG) cfg = copy.deepcopy(DEFAULT_CONFIG)
cfg.update(results_dir=str(tmp_path / "results"), data_cache_dir=str(tmp_path / "cache"), cfg.update(results_dir=str(tmp_path / "results"), data_cache_dir=str(tmp_path / "cache"),
memory_log_path=str(tmp_path / "log.md"), **config) memory_log_path=str(tmp_path / "log.md"), **config)
monkeypatch.setattr(trading_graph, "create_llm_client", lambda **k: _Client(model)) 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 @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] held = [v for v in vars(graph).values() if isinstance(v, dict) and TRADE_DATE in v]
assert held == [] assert held == []
assert len(list(tmp_path.glob("results/NVDA/TradingAgentsStrategy_logs/*.json"))) == 2 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
-50
View File
@@ -5,11 +5,9 @@ import unittest
from unittest.mock import patch from unittest.mock import patch
import pytest import pytest
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
from tradingagents.agents.context import ( from tradingagents.agents.context import (
build_instrument_context, build_instrument_context,
create_msg_delete,
get_instrument_context_from_state, get_instrument_context_from_state,
resolve_instrument_identity, resolve_instrument_identity,
) )
@@ -118,53 +116,5 @@ class GetInstrumentContextFromStateTests(unittest.TestCase):
self.assertIn("crypto asset", context) 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__": if __name__ == "__main__":
unittest.main() unittest.main()
+2 -2
View File
@@ -136,8 +136,8 @@ def test_the_cli_says_when_a_run_produced_no_usable_rating(monkeypatch, tmp_path
def end_checkpoint(self): def end_checkpoint(self):
pass pass
def stream(self, *a, **k): def stream_run(self, *a, **k):
yield {"messages": [], "final_trade_decision": REFUSAL, "final_rating": RATING_REVIEW} yield [], {"messages": [], "final_trade_decision": REFUSAL, "final_rating": RATING_REVIEW}
fake = _Graph() fake = _Graph()
fake.graph = fake fake.graph = fake
-2
View File
@@ -2,7 +2,6 @@ from .analysts.fundamentals_analyst import create_fundamentals_analyst
from .analysts.market_analyst import create_market_analyst from .analysts.market_analyst import create_market_analyst
from .analysts.news_analyst import create_news_analyst from .analysts.news_analyst import create_news_analyst
from .analysts.sentiment_analyst import create_sentiment_analyst from .analysts.sentiment_analyst import create_sentiment_analyst
from .context import create_msg_delete
from .managers.portfolio_manager import create_portfolio_manager from .managers.portfolio_manager import create_portfolio_manager
from .managers.research_manager import create_research_manager from .managers.research_manager import create_research_manager
from .researchers.bear_researcher import create_bear_researcher from .researchers.bear_researcher import create_bear_researcher
@@ -15,7 +14,6 @@ from .trader.trader import create_trader
__all__ = [ __all__ = [
"AgentState", "AgentState",
"create_msg_delete",
"InvestDebateState", "InvestDebateState",
"RiskDebateState", "RiskDebateState",
"create_bear_researcher", "create_bear_researcher",
-32
View File
@@ -6,8 +6,6 @@ import logging
from collections.abc import Mapping from collections.abc import Mapping
from typing import Any from typing import Any
from langchain_core.messages import HumanMessage, RemoveMessage
from tradingagents.dataflows.date_window import get_current_date from tradingagents.dataflows.date_window import get_current_date
from tradingagents.dataflows.vendors.yahoo.fundamentals import get_company_profile from tradingagents.dataflows.vendors.yahoo.fundamentals import get_company_profile
@@ -206,33 +204,3 @@ def get_portfolio_context_from_state(state: Mapping[str, Any]) -> str:
"holdings or cash, so do not assume a flat book; give direction and " "holdings or cash, so do not assume a flat book; give direction and "
"sizing guidance in terms the caller can apply to their own position." "sizing guidance in terms the caller can apply to their own position."
) )
def create_msg_delete():
def delete_messages(state):
"""Clear messages and add a context-anchored placeholder.
The placeholder must not be a bare ``"Continue"``: some
OpenAI-compatible providers interpret that literally as the user task
and produce output about the word "continue" instead of analysing the
instrument (#888). Anchoring it to the resolved instrument context and
date keeps the next analyst on-task even if the provider treats the
placeholder as a standalone request.
"""
messages = state["messages"]
removal_operations = [RemoveMessage(id=m.id) for m in messages]
instrument_context = get_instrument_context_from_state(state)
trade_date = state.get("trade_date", "the requested date")
placeholder = HumanMessage(
content=(
f"Proceed with your assigned analysis for this workflow. "
f"{instrument_context} The analysis date is {trade_date}."
)
)
return {"messages": removal_operations + [placeholder]}
return delete_messages
-10
View File
@@ -8,15 +8,9 @@ from tradingagents.agents.analysts import fundamentals_analyst, market_analyst,
class AnalystNodeSpec: class AnalystNodeSpec:
key: str key: str
agent_node: str agent_node: str
clear_node: str
report_key: str report_key: str
tools: tuple = () tools: tuple = ()
@property
def tool_node(self) -> str | None:
"""The node that runs this analyst's tool calls; None when it has no tools."""
return f"tools_{self.key}" if self.tools else None
@dataclass(frozen=True) @dataclass(frozen=True)
class AnalystExecutionPlan: class AnalystExecutionPlan:
@@ -27,7 +21,6 @@ ANALYST_NODE_SPECS: dict[str, AnalystNodeSpec] = {
"market": AnalystNodeSpec( "market": AnalystNodeSpec(
key="market", key="market",
agent_node="Market Analyst", agent_node="Market Analyst",
clear_node="Msg Clear Market",
report_key="market_report", report_key="market_report",
tools=market_analyst.TOOLS, tools=market_analyst.TOOLS,
), ),
@@ -36,20 +29,17 @@ ANALYST_NODE_SPECS: dict[str, AnalystNodeSpec] = {
# sources before calling the model, so it has no tools. # sources before calling the model, so it has no tools.
key="social", key="social",
agent_node="Sentiment Analyst", agent_node="Sentiment Analyst",
clear_node="Msg Clear Sentiment",
report_key="sentiment_report", report_key="sentiment_report",
), ),
"news": AnalystNodeSpec( "news": AnalystNodeSpec(
key="news", key="news",
agent_node="News Analyst", agent_node="News Analyst",
clear_node="Msg Clear News",
report_key="news_report", report_key="news_report",
tools=news_analyst.TOOLS, tools=news_analyst.TOOLS,
), ),
"fundamentals": AnalystNodeSpec( "fundamentals": AnalystNodeSpec(
key="fundamentals", key="fundamentals",
agent_node="Fundamentals Analyst", agent_node="Fundamentals Analyst",
clear_node="Msg Clear Fundamentals",
report_key="fundamentals_report", report_key="fundamentals_report",
tools=fundamentals_analyst.TOOLS, tools=fundamentals_analyst.TOOLS,
), ),
+30 -25
View File
@@ -1,4 +1,4 @@
from typing import Any from typing import Any, TypedDict
from langgraph.graph import END, START, StateGraph from langgraph.graph import END, START, StateGraph
from langgraph.prebuilt import ToolNode from langgraph.prebuilt import ToolNode
@@ -10,7 +10,6 @@ from tradingagents.agents import (
create_conservative_debator, create_conservative_debator,
create_fundamentals_analyst, create_fundamentals_analyst,
create_market_analyst, create_market_analyst,
create_msg_delete,
create_neutral_debator, create_neutral_debator,
create_news_analyst, create_news_analyst,
create_portfolio_manager, create_portfolio_manager,
@@ -40,11 +39,28 @@ RISK_ANALYSIS_PATH_MAP = {
} }
def _tools_or_clear(spec): def _tools_or_done(state) -> str:
"""Route an analyst's turn: run its tool calls, or finish its report.""" """Route an analyst's turn: run its tool calls, or finish with its report."""
def route(state) -> str: return "tools" if state["messages"][-1].tool_calls else END
return spec.tool_node if state["messages"][-1].tool_calls else spec.clear_node
return route
def _analyst_graph(spec, agent):
"""One analyst as a graph of its own: the model and its tools, on a private message history.
It returns only its report, so analysts running side by side never write the
same key, and its tool calls never reach the other analysts' messages.
"""
output = TypedDict(f"{spec.key.capitalize()}Report", {spec.report_key: str})
graph = StateGraph(AgentState, output_schema=output)
graph.add_node("agent", agent)
graph.add_edge(START, "agent")
if spec.tools:
graph.add_node("tools", ToolNode(list(spec.tools)))
graph.add_conditional_edges("agent", _tools_or_done, ["tools", END])
graph.add_edge("tools", "agent")
else:
graph.add_edge("agent", END)
return graph.compile()
class GraphSetup: class GraphSetup:
@@ -95,10 +111,7 @@ class GraphSetup:
workflow = StateGraph(AgentState) workflow = StateGraph(AgentState)
for spec in plan.specs: for spec in plan.specs:
workflow.add_node(spec.agent_node, analyst_factories[spec.key]()) workflow.add_node(spec.agent_node, _analyst_graph(spec, analyst_factories[spec.key]()))
workflow.add_node(spec.clear_node, create_msg_delete())
if spec.tools:
workflow.add_node(spec.tool_node, ToolNode(list(spec.tools)))
workflow.add_node("Bull Researcher", bull_researcher_node) workflow.add_node("Bull Researcher", bull_researcher_node)
workflow.add_node("Bear Researcher", bear_researcher_node) workflow.add_node("Bear Researcher", bear_researcher_node)
@@ -109,20 +122,12 @@ class GraphSetup:
workflow.add_node("Conservative Analyst", conservative_analyst) workflow.add_node("Conservative Analyst", conservative_analyst)
workflow.add_node("Portfolio Manager", portfolio_manager_node) workflow.add_node("Portfolio Manager", portfolio_manager_node)
workflow.add_edge(START, plan.specs[0].agent_node) # The analysts work at the same time; the research debate starts once
# every one of them has filed its report.
for i, spec in enumerate(plan.specs): analysts = [spec.agent_node for spec in plan.specs]
if spec.tools: for node in analysts:
workflow.add_conditional_edges( workflow.add_edge(START, node)
spec.agent_node, _tools_or_clear(spec), [spec.tool_node, spec.clear_node] workflow.add_edge(analysts, "Bull Researcher")
)
workflow.add_edge(spec.tool_node, spec.agent_node)
else:
workflow.add_edge(spec.agent_node, spec.clear_node)
# The last analyst hands over to the research debate.
following = plan.specs[i + 1].agent_node if i < len(plan.specs) - 1 else "Bull Researcher"
workflow.add_edge(spec.clear_node, following)
# Both research-debate edges share the complete DEBATE_PATH_MAP (#1088). # Both research-debate edges share the complete DEBATE_PATH_MAP (#1088).
for debate_node in ("Bull Researcher", "Bear Researcher"): for debate_node in ("Bull Researcher", "Bear Researcher"):
+34 -17
View File
@@ -152,6 +152,9 @@ class TradingAgentsGraph:
f"asset={asset_type}", f"asset={asset_type}",
# None, an empty book and a changed book are three different runs. # None, an empty book and a changed book are three different runs.
f"portfolio={portfolio.fingerprint() if portfolio is not None else 'none'}", f"portfolio={portfolio.fingerprint() if portfolio is not None else 'none'}",
# The layout itself: a checkpoint saved when analysts ran one after
# another has pending nodes this graph no longer has.
"analysts=parallel",
]) ])
def propagate(self, company_name, trade_date, asset_type: str = "stock", portfolio=None): def propagate(self, company_name, trade_date, asset_type: str = "stock", portfolio=None):
@@ -333,24 +336,16 @@ class TradingAgentsGraph:
# None resumes an existing checkpoint; init_agent_state starts fresh (#1249). # None resumes an existing checkpoint; init_agent_state starts fresh (#1249).
graph_input = self.checkpoint_input(init_agent_state) graph_input = self.checkpoint_input(init_agent_state)
if self.debug: if self.debug:
trace = [] # A state repeats the messages before it, so each prints once (#1027).
last_printed = None final_state, printed = {}, set()
for chunk in self.graph.stream(graph_input, **args): for messages, state in self.stream_run(graph_input, **args):
if chunk["messages"]: for msg in messages:
msg = chunk["messages"][-1] key = getattr(msg, "id", None) or (type(msg).__name__, getattr(msg, "content", None))
# Nodes after the trader don't append to messages, so the if key not in printed:
# same trailing message repeats across chunks. Print it only printed.add(key)
# when it changes (#1027); the trace/state merge is unchanged.
signature = (type(msg).__name__, getattr(msg, "content", None))
if signature != last_printed:
msg.pretty_print() msg.pretty_print()
last_printed = signature if state is not None:
trace.append(chunk) final_state.update(state)
# Streamed chunks are per-node deltas. Merge them so the returned
# state matches what graph.invoke() yields in the non-debug path.
final_state = {}
for chunk in trace:
final_state.update(chunk)
else: else:
final_state = self.graph.invoke(graph_input, **args) final_state = self.graph.invoke(graph_input, **args)
@@ -364,6 +359,28 @@ class TradingAgentsGraph:
return final_state, run_rating(final_state) return final_state, run_rating(final_state)
def stream_run(self, graph_input, **args):
"""Stream a run as ``(messages, state)`` pairs.
``messages`` are the agents' messages, the analysts' included. ``state``
is the run's state after a top-level step; for a step inside an analyst's
graph it is that analyst's report once filed, else None.
Each analyst works in a graph of its own, and the run's state takes the
analysts' reports only when the slowest has finished, so their messages
and reports come from their own finished steps ("tasks") as they happen.
"""
args = {**args, "stream_mode": ["values", "tasks"]}
for namespace, mode, chunk in self.graph.stream(graph_input, subgraphs=True, **args):
if namespace:
result = chunk.get("result") if mode == "tasks" else None
if isinstance(result, dict):
report = {k: v for k, v in result.items() if k != "messages" and v}
if result.get("messages") or report:
yield result.get("messages", []), report or None
elif mode == "values":
yield chunk.get("messages", []), chunk
def _log_state(self, trade_date, final_state): def _log_state(self, trade_date, final_state):
"""Write a run's final state to JSON under the run's own ticker.""" """Write a run's final state to JSON under the run's own ticker."""
entry = { entry = {