diff --git a/README.md b/README.md index c92b5dce6..d3cbb663b 100644 --- a/README.md +++ b/README.md @@ -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. - 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. +

diff --git a/cli/display.py b/cli/display.py index 1510e9581..0f9bd0943 100644 --- a/cli/display.py +++ b/cli/display.py @@ -463,22 +463,17 @@ ANALYST_REPORT_MAP = { 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: - - Store new report content from the current chunk if present - - 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 + The analysts run together: each is in progress until its own report lands. + When every selected analyst has filed, the research debate is in progress. """ selected = message_buffer.selected_analysts - found_active = False if wall_time_tracker is not None: sync_analyst_tracker_from_chunk(wall_time_tracker, chunk) + all_filed = True for analyst_key in ANALYST_ORDER: if analyst_key not in selected: continue @@ -490,20 +485,15 @@ def update_analyst_statuses(message_buffer, chunk, wall_time_tracker=None): if chunk.get(report_key): message_buffer.update_report_section(report_key, chunk[report_key]) - # Determine status from accumulated sections, not just current chunk - has_report = bool(message_buffer.report_sections.get(report_key)) - - if has_report: + # Status comes from accumulated sections, not just the current chunk. + if message_buffer.report_sections.get(report_key): 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: - 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 ( - not found_active + all_filed and selected and message_buffer.agent_status.get("Bull Researcher") == "pending" ): @@ -624,17 +614,9 @@ def sync_analyst_tracker_from_chunk( chunk: dict[str, str], now: float | None = None, ) -> None: + """The analysts start together; each stops its clock when its report lands.""" current_time = monotonic() if now is None else now - active_found = False - for spec in tracker.plan.specs: - has_report = bool(chunk.get(spec.report_key)) - - if has_report: - tracker.mark_started(spec.key, started_at=current_time) + tracker.mark_started(spec.key, started_at=current_time) + if chunk.get(spec.report_key): 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 diff --git a/cli/run.py b/cli/run.py index dc6408cb2..4070d3e65 100644 --- a/cli/run.py +++ b/cli/run.py @@ -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) - first_analyst = analyst_execution_plan.specs[0].agent_node - message_buffer.update_agent_status(first_analyst, "in_progress") - analyst_wall_time_tracker.mark_started(selected_analyst_keys[0]) + # The analysts start together. + for spec in analyst_execution_plan.specs: + 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) 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. trace = [] try: - for chunk in graph.graph.stream(graph.checkpoint_input(init_agent_state), **args): - for message in chunk.get("messages", []): + for messages, chunk in graph.stream_run(graph.checkpoint_input(init_agent_state), **args): + for message in messages: msg_id = getattr(message, "id", None) if msg_id is not None: 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: 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( message_buffer, chunk, diff --git a/tests/test_analyst_execution.py b/tests/test_analyst_execution.py index f668b3b16..213189c4b 100644 --- a/tests/test_analyst_execution.py +++ b/tests/test_analyst_execution.py @@ -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): diff --git a/tests/test_checkpoint_resume.py b/tests/test_checkpoint_resume.py index 6d134f6bb..c85f3d3c5 100644 --- a/tests/test_checkpoint_resume.py +++ b/tests/test_checkpoint_resume.py @@ -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__": diff --git a/tests/test_cli_display.py b/tests/test_cli_display.py index 12a8f1084..702a31c8f 100644 --- a/tests/test_cli_display.py +++ b/tests/test_cli_display.py @@ -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" diff --git a/tests/test_cli_memory_log.py b/tests/test_cli_memory_log.py index fa55dfd63..4e3170660 100644 --- a/tests/test_cli_memory_log.py +++ b/tests/test_cli_memory_log.py @@ -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: diff --git a/tests/test_graph_end_to_end.py b/tests/test_graph_end_to_end.py index 7b791220c..e51df46ad 100644 --- a/tests/test_graph_end_to_end.py +++ b/tests/test_graph_end_to_end.py @@ -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 diff --git a/tests/test_instrument_identity.py b/tests/test_instrument_identity.py index 4c12a4769..11c79e67c 100644 --- a/tests/test_instrument_identity.py +++ b/tests/test_instrument_identity.py @@ -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() diff --git a/tests/test_rating_integrity.py b/tests/test_rating_integrity.py index fabe76082..61def3de2 100644 --- a/tests/test_rating_integrity.py +++ b/tests/test_rating_integrity.py @@ -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 diff --git a/tradingagents/agents/__init__.py b/tradingagents/agents/__init__.py index 5675c7d2e..169b70910 100644 --- a/tradingagents/agents/__init__.py +++ b/tradingagents/agents/__init__.py @@ -2,7 +2,6 @@ from .analysts.fundamentals_analyst import create_fundamentals_analyst from .analysts.market_analyst import create_market_analyst from .analysts.news_analyst import create_news_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.research_manager import create_research_manager from .researchers.bear_researcher import create_bear_researcher @@ -15,7 +14,6 @@ from .trader.trader import create_trader __all__ = [ "AgentState", - "create_msg_delete", "InvestDebateState", "RiskDebateState", "create_bear_researcher", diff --git a/tradingagents/agents/context.py b/tradingagents/agents/context.py index 9842f2d5b..354531e27 100644 --- a/tradingagents/agents/context.py +++ b/tradingagents/agents/context.py @@ -6,8 +6,6 @@ import logging from collections.abc import Mapping from typing import Any -from langchain_core.messages import HumanMessage, RemoveMessage - from tradingagents.dataflows.date_window import get_current_date 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 " "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 - - - diff --git a/tradingagents/graph/analyst_execution.py b/tradingagents/graph/analyst_execution.py index 741685c50..3a4d0f70b 100644 --- a/tradingagents/graph/analyst_execution.py +++ b/tradingagents/graph/analyst_execution.py @@ -8,15 +8,9 @@ from tradingagents.agents.analysts import fundamentals_analyst, market_analyst, class AnalystNodeSpec: key: str agent_node: str - clear_node: str report_key: str 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) class AnalystExecutionPlan: @@ -27,7 +21,6 @@ ANALYST_NODE_SPECS: dict[str, AnalystNodeSpec] = { "market": AnalystNodeSpec( key="market", agent_node="Market Analyst", - clear_node="Msg Clear Market", report_key="market_report", tools=market_analyst.TOOLS, ), @@ -36,20 +29,17 @@ ANALYST_NODE_SPECS: dict[str, AnalystNodeSpec] = { # sources before calling the model, so it has no tools. key="social", agent_node="Sentiment Analyst", - clear_node="Msg Clear Sentiment", report_key="sentiment_report", ), "news": AnalystNodeSpec( key="news", agent_node="News Analyst", - clear_node="Msg Clear News", report_key="news_report", tools=news_analyst.TOOLS, ), "fundamentals": AnalystNodeSpec( key="fundamentals", agent_node="Fundamentals Analyst", - clear_node="Msg Clear Fundamentals", report_key="fundamentals_report", tools=fundamentals_analyst.TOOLS, ), diff --git a/tradingagents/graph/setup.py b/tradingagents/graph/setup.py index df4577fa0..d17215d96 100644 --- a/tradingagents/graph/setup.py +++ b/tradingagents/graph/setup.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, TypedDict from langgraph.graph import END, START, StateGraph from langgraph.prebuilt import ToolNode @@ -10,7 +10,6 @@ from tradingagents.agents import ( create_conservative_debator, create_fundamentals_analyst, create_market_analyst, - create_msg_delete, create_neutral_debator, create_news_analyst, create_portfolio_manager, @@ -40,11 +39,28 @@ RISK_ANALYSIS_PATH_MAP = { } -def _tools_or_clear(spec): - """Route an analyst's turn: run its tool calls, or finish its report.""" - def route(state) -> str: - return spec.tool_node if state["messages"][-1].tool_calls else spec.clear_node - return route +def _tools_or_done(state) -> str: + """Route an analyst's turn: run its tool calls, or finish with its report.""" + return "tools" if state["messages"][-1].tool_calls else END + + +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: @@ -95,10 +111,7 @@ class GraphSetup: workflow = StateGraph(AgentState) for spec in plan.specs: - workflow.add_node(spec.agent_node, 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(spec.agent_node, _analyst_graph(spec, analyst_factories[spec.key]())) workflow.add_node("Bull Researcher", bull_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("Portfolio Manager", portfolio_manager_node) - workflow.add_edge(START, plan.specs[0].agent_node) - - for i, spec in enumerate(plan.specs): - if spec.tools: - workflow.add_conditional_edges( - spec.agent_node, _tools_or_clear(spec), [spec.tool_node, spec.clear_node] - ) - 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) + # The analysts work at the same time; the research debate starts once + # every one of them has filed its report. + analysts = [spec.agent_node for spec in plan.specs] + for node in analysts: + workflow.add_edge(START, node) + workflow.add_edge(analysts, "Bull Researcher") # Both research-debate edges share the complete DEBATE_PATH_MAP (#1088). for debate_node in ("Bull Researcher", "Bear Researcher"): diff --git a/tradingagents/graph/trading_graph.py b/tradingagents/graph/trading_graph.py index 0c3e59b95..05953eb82 100644 --- a/tradingagents/graph/trading_graph.py +++ b/tradingagents/graph/trading_graph.py @@ -152,6 +152,9 @@ class TradingAgentsGraph: f"asset={asset_type}", # None, an empty book and a changed book are three different runs. 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): @@ -333,24 +336,16 @@ class TradingAgentsGraph: # None resumes an existing checkpoint; init_agent_state starts fresh (#1249). graph_input = self.checkpoint_input(init_agent_state) if self.debug: - trace = [] - last_printed = None - for chunk in self.graph.stream(graph_input, **args): - if chunk["messages"]: - msg = chunk["messages"][-1] - # Nodes after the trader don't append to messages, so the - # same trailing message repeats across chunks. Print it only - # when it changes (#1027); the trace/state merge is unchanged. - signature = (type(msg).__name__, getattr(msg, "content", None)) - if signature != last_printed: + # A state repeats the messages before it, so each prints once (#1027). + final_state, printed = {}, set() + for messages, state in self.stream_run(graph_input, **args): + for msg in messages: + key = getattr(msg, "id", None) or (type(msg).__name__, getattr(msg, "content", None)) + if key not in printed: + printed.add(key) msg.pretty_print() - last_printed = signature - trace.append(chunk) - # 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) + if state is not None: + final_state.update(state) else: final_state = self.graph.invoke(graph_input, **args) @@ -364,6 +359,28 @@ class TradingAgentsGraph: 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): """Write a run's final state to JSON under the run's own ticker.""" entry = {