466 lines
15 KiB
Python
466 lines
15 KiB
Python
import json
|
|
|
|
from tradingagents.graph.setup import GraphSetup
|
|
from tradingagents.graph.trading_graph import TradingAgentsGraph
|
|
|
|
|
|
class DummyStateGraph:
|
|
def __init__(self, _state_type):
|
|
self.nodes = {}
|
|
self.conditional_edges = {}
|
|
|
|
def add_node(self, name, node):
|
|
self.nodes[name] = node
|
|
|
|
def add_edge(self, *_args, **_kwargs):
|
|
return None
|
|
|
|
def add_conditional_edges(self, source, condition, destinations):
|
|
self.conditional_edges[source] = {
|
|
"condition": condition,
|
|
"destinations": destinations,
|
|
}
|
|
|
|
def compile(self):
|
|
return {
|
|
"nodes": self.nodes,
|
|
"conditional_edges": self.conditional_edges,
|
|
}
|
|
|
|
|
|
class DummyToolNode:
|
|
def __init__(self, tools):
|
|
self.tools = tools
|
|
|
|
|
|
def test_scenario_tools_route_to_vendor(monkeypatch):
|
|
import tradingagents.dataflows.interface as interface
|
|
from tradingagents.agents.utils.scenario_tools import (
|
|
get_catalyst_calendar,
|
|
get_scenario_fundamentals,
|
|
get_scenario_news,
|
|
)
|
|
|
|
calls = []
|
|
|
|
def fake_route_to_vendor(method, *args, **kwargs):
|
|
calls.append((method, args, kwargs))
|
|
return f"{method}-result"
|
|
|
|
monkeypatch.setattr(interface, "route_to_vendor", fake_route_to_vendor)
|
|
|
|
assert (
|
|
get_scenario_fundamentals.invoke({"ticker": "AAPL", "curr_date": "2026-03-24"})
|
|
== "get_fundamentals-result"
|
|
)
|
|
assert (
|
|
get_scenario_news.invoke(
|
|
{"query": "AAPL product launch catalyst", "start_date": "2026-03-01", "end_date": "2026-03-24"}
|
|
)
|
|
== "get_news-result"
|
|
)
|
|
assert (
|
|
get_catalyst_calendar.invoke({"curr_date": "2026-03-24"})
|
|
== "get_fed_calendar-result"
|
|
)
|
|
assert calls == [
|
|
(
|
|
"get_fundamentals",
|
|
(),
|
|
{"ticker": "AAPL", "curr_date": "2026-03-24"},
|
|
),
|
|
(
|
|
"get_news",
|
|
(),
|
|
{
|
|
"query": "AAPL product launch catalyst",
|
|
"start_date": "2026-03-01",
|
|
"end_date": "2026-03-24",
|
|
},
|
|
),
|
|
("get_fed_calendar", (), {"curr_date": "2026-03-24"}),
|
|
]
|
|
|
|
|
|
def test_graph_setup_wires_scenario_catalyst_analyst_and_tools(monkeypatch):
|
|
recorded_llms = {}
|
|
|
|
monkeypatch.setattr("tradingagents.graph.setup.StateGraph", DummyStateGraph)
|
|
monkeypatch.setattr("tradingagents.graph.setup.create_msg_delete", lambda: "delete")
|
|
|
|
def make_factory(node_name):
|
|
def factory(llm, *_args):
|
|
recorded_llms[node_name] = llm
|
|
return node_name
|
|
|
|
return factory
|
|
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_market_analyst",
|
|
make_factory("Market Analyst"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_scenario_catalyst_analyst",
|
|
make_factory("Scenario Analyst"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_social_media_analyst",
|
|
make_factory("Social Analyst"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_news_analyst",
|
|
make_factory("News Analyst"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_fundamentals_analyst",
|
|
make_factory("Fundamentals Analyst"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_bull_researcher",
|
|
make_factory("Bull Researcher"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_bear_researcher",
|
|
make_factory("Bear Researcher"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_research_manager",
|
|
make_factory("Research Manager"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_trader",
|
|
make_factory("Trader"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_aggressive_debator",
|
|
make_factory("Aggressive Analyst"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_neutral_debator",
|
|
make_factory("Neutral Analyst"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_conservative_debator",
|
|
make_factory("Conservative Analyst"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"tradingagents.graph.setup.create_portfolio_manager",
|
|
make_factory("Portfolio Manager"),
|
|
)
|
|
|
|
class PartialConditionalLogic:
|
|
def should_continue_market(self, _state):
|
|
return "Msg Clear Market"
|
|
|
|
def should_continue_scenario(self, _state):
|
|
return "Msg Clear Scenario"
|
|
|
|
def should_continue_debate(self, _state):
|
|
return "Research Manager"
|
|
|
|
def should_continue_risk_analysis(self, _state):
|
|
return "Portfolio Manager"
|
|
|
|
setup = GraphSetup(
|
|
quick_thinking_llm="quick-llm",
|
|
deep_thinking_llm="deep-llm",
|
|
tool_nodes={"market": "market-tools", "scenario": "scenario-tools"},
|
|
bull_memory=object(),
|
|
bear_memory=object(),
|
|
trader_memory=object(),
|
|
invest_judge_memory=object(),
|
|
portfolio_manager_memory=object(),
|
|
conditional_logic=PartialConditionalLogic(),
|
|
role_llms={"scenario": "scenario-llm"},
|
|
)
|
|
|
|
graph = setup.setup_graph(selected_analysts=["market", "scenario"])
|
|
|
|
assert recorded_llms["Scenario Analyst"] == "scenario-llm"
|
|
assert graph["nodes"]["Scenario Analyst"] == "Scenario Analyst"
|
|
assert graph["nodes"]["tools_scenario"] == "scenario-tools"
|
|
assert "Scenario Analyst" in graph["conditional_edges"]
|
|
|
|
|
|
def test_trading_graph_creates_scenario_tool_node(monkeypatch):
|
|
monkeypatch.setattr("tradingagents.graph.trading_graph.ToolNode", DummyToolNode)
|
|
|
|
graph = TradingAgentsGraph.__new__(TradingAgentsGraph)
|
|
tool_nodes = TradingAgentsGraph._create_tool_nodes(graph)
|
|
|
|
assert [tool.name for tool in tool_nodes["scenario"].tools] == [
|
|
"get_scenario_fundamentals",
|
|
"get_scenario_news",
|
|
"get_catalyst_calendar",
|
|
]
|
|
|
|
|
|
class DummyPrompt:
|
|
def __init__(self, result):
|
|
self.result = result
|
|
|
|
def partial(self, **_kwargs):
|
|
return self
|
|
|
|
def __or__(self, _other):
|
|
return DummyChain(self.result)
|
|
|
|
|
|
class DummyChain:
|
|
def __init__(self, result):
|
|
self.result = result
|
|
|
|
def invoke(self, _messages):
|
|
return self.result
|
|
|
|
|
|
class DummyResult:
|
|
def __init__(self, content, tool_calls):
|
|
self.content = content
|
|
self.tool_calls = tool_calls
|
|
|
|
|
|
class DummyLLM:
|
|
def __init__(self, result):
|
|
self.result = result
|
|
self.bound_tool_names = []
|
|
|
|
def bind_tools(self, tools):
|
|
self.bound_tool_names = [tool.name for tool in tools]
|
|
return object()
|
|
|
|
|
|
def test_scenario_catalyst_analyst_returns_structured_data(monkeypatch):
|
|
from tradingagents.agents.analysts.scenario_catalyst_analyst import (
|
|
create_scenario_catalyst_analyst,
|
|
)
|
|
|
|
result = DummyResult(
|
|
content="""## Scenario and Catalyst Summary
|
|
|
|
Bull case is an AI-led re-rating with operating leverage, base case is steady
|
|
execution with modest multiple expansion, and bear case is demand rollover.
|
|
|
|
```json
|
|
{
|
|
"scenario_map": [
|
|
{
|
|
"name": "bull",
|
|
"probability_pct": 30,
|
|
"thesis": "AI-driven demand acceleration",
|
|
"valuation_implication": "multiple expansion toward upper historical range",
|
|
"signposts": ["order lead-times extend", "gross margin beats guidance"]
|
|
},
|
|
{
|
|
"name": "base",
|
|
"probability_pct": 50,
|
|
"thesis": "stable demand and disciplined opex",
|
|
"valuation_implication": "range-bound multiple with EPS carry",
|
|
"signposts": ["in-line guidance", "services growth steady"]
|
|
},
|
|
{
|
|
"name": "bear",
|
|
"probability_pct": 20,
|
|
"thesis": "weaker upgrade cycle and pricing pressure",
|
|
"valuation_implication": "derating to cycle-low valuation band",
|
|
"signposts": ["inventory builds", "discounting rises"]
|
|
}
|
|
],
|
|
"dated_catalyst_map": [
|
|
{
|
|
"catalyst": "FOMC rate decision",
|
|
"date_or_window": "2026-05-06",
|
|
"related_scenarios": ["bull", "base", "bear"],
|
|
"expected_impact": "changes discount-rate pressure on valuation",
|
|
"confidence": "medium"
|
|
}
|
|
],
|
|
"invalidation_triggers": [
|
|
{
|
|
"trigger": "two consecutive quarters of revenue miss versus guidance midpoint",
|
|
"affected_scenarios": ["bull", "base"],
|
|
"severity": "high",
|
|
"evidence_to_watch": "quarterly filings and management commentary"
|
|
}
|
|
]
|
|
}
|
|
```""",
|
|
tool_calls=[],
|
|
)
|
|
llm = DummyLLM(result)
|
|
monkeypatch.setattr(
|
|
"tradingagents.agents.analysts.scenario_catalyst_analyst.ChatPromptTemplate.from_messages",
|
|
lambda _messages: DummyPrompt(result),
|
|
)
|
|
|
|
node = create_scenario_catalyst_analyst(llm)
|
|
output = node(
|
|
{
|
|
"trade_date": "2026-03-24",
|
|
"company_of_interest": "AAPL",
|
|
"messages": ["analyze scenario tree"],
|
|
}
|
|
)
|
|
|
|
assert llm.bound_tool_names == [
|
|
"get_scenario_fundamentals",
|
|
"get_scenario_news",
|
|
"get_catalyst_calendar",
|
|
]
|
|
assert output["scenario_catalyst_report"] == result.content
|
|
assert output["scenario_catalyst_data"] == {
|
|
"ticker": "AAPL",
|
|
"analysis_date": "2026-03-24",
|
|
"scenario_map": [
|
|
{
|
|
"name": "bull",
|
|
"probability_pct": 30,
|
|
"thesis": "AI-driven demand acceleration",
|
|
"valuation_implication": "multiple expansion toward upper historical range",
|
|
"signposts": [
|
|
"order lead-times extend",
|
|
"gross margin beats guidance",
|
|
],
|
|
},
|
|
{
|
|
"name": "base",
|
|
"probability_pct": 50,
|
|
"thesis": "stable demand and disciplined opex",
|
|
"valuation_implication": "range-bound multiple with EPS carry",
|
|
"signposts": ["in-line guidance", "services growth steady"],
|
|
},
|
|
{
|
|
"name": "bear",
|
|
"probability_pct": 20,
|
|
"thesis": "weaker upgrade cycle and pricing pressure",
|
|
"valuation_implication": "derating to cycle-low valuation band",
|
|
"signposts": ["inventory builds", "discounting rises"],
|
|
},
|
|
],
|
|
"dated_catalyst_map": [
|
|
{
|
|
"catalyst": "FOMC rate decision",
|
|
"date_or_window": "2026-05-06",
|
|
"related_scenarios": ["bull", "base", "bear"],
|
|
"expected_impact": "changes discount-rate pressure on valuation",
|
|
"confidence": "medium",
|
|
}
|
|
],
|
|
"invalidation_triggers": [
|
|
{
|
|
"trigger": "two consecutive quarters of revenue miss versus guidance midpoint",
|
|
"affected_scenarios": ["bull", "base"],
|
|
"severity": "high",
|
|
"evidence_to_watch": "quarterly filings and management commentary",
|
|
}
|
|
],
|
|
}
|
|
assert output["messages"] == [result]
|
|
|
|
|
|
def test_extract_scenario_payload_tolerates_common_model_json_variants():
|
|
from tradingagents.agents.analysts.scenario_catalyst_analyst import (
|
|
_extract_scenario_catalyst_payload,
|
|
)
|
|
|
|
expected = {
|
|
"scenario_map": [{"name": "base", "probability_pct": 60}],
|
|
"dated_catalyst_map": [{"catalyst": "earnings", "date_or_window": "Q2"}],
|
|
"invalidation_triggers": [{"trigger": "demand miss"}],
|
|
}
|
|
|
|
uppercase_fence = """
|
|
```JSON
|
|
{"scenario_map":[{"name":"base","probability_pct":60}],"dated_catalyst_map":[{"catalyst":"earnings","date_or_window":"Q2"}],"invalidation_triggers":[{"trigger":"demand miss"}]}
|
|
```
|
|
"""
|
|
plain_fence = """
|
|
```
|
|
{"scenario_map":[{"name":"base","probability_pct":60}],"dated_catalyst_map":[{"catalyst":"earnings","date_or_window":"Q2"}],"invalidation_triggers":[{"trigger":"demand miss"}]}
|
|
```
|
|
"""
|
|
raw_json = """
|
|
Narrative intro before payload.
|
|
{"scenario_map":[{"name":"base","probability_pct":60}],"dated_catalyst_map":[{"catalyst":"earnings","date_or_window":"Q2"}],"invalidation_triggers":[{"trigger":"demand miss"}]}
|
|
Tail note after payload.
|
|
"""
|
|
|
|
assert _extract_scenario_catalyst_payload(uppercase_fence) == expected
|
|
assert _extract_scenario_catalyst_payload(plain_fence) == expected
|
|
assert _extract_scenario_catalyst_payload(raw_json) == expected
|
|
|
|
|
|
def test_propagator_initial_state_seeds_scenario_defaults():
|
|
from tradingagents.graph.propagation import Propagator
|
|
|
|
state = Propagator().create_initial_state("AAPL", "2026-03-24")
|
|
|
|
assert state["scenario_catalyst_report"] == ""
|
|
assert state["scenario_catalyst_data"] == {
|
|
"ticker": "",
|
|
"analysis_date": "",
|
|
"scenario_map": [],
|
|
"dated_catalyst_map": [],
|
|
"invalidation_triggers": [],
|
|
}
|
|
|
|
|
|
def test_log_state_persists_scenario_catalyst_report_and_data(tmp_path, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
graph = TradingAgentsGraph.__new__(TradingAgentsGraph)
|
|
graph.ticker = "AAPL"
|
|
graph.log_states_dict = {}
|
|
|
|
final_state = {
|
|
"company_of_interest": "Apple",
|
|
"trade_date": "2026-03-24",
|
|
"market_report": "market",
|
|
"sentiment_report": "sentiment",
|
|
"news_report": "news",
|
|
"fundamentals_report": "fundamentals",
|
|
"segment_report": "segment report",
|
|
"segment_data": {"ticker": "AAPL"},
|
|
"macro_report": "macro report",
|
|
"scenario_catalyst_report": "scenario report",
|
|
"scenario_catalyst_data": {
|
|
"ticker": "AAPL",
|
|
"analysis_date": "2026-03-24",
|
|
"scenario_map": [{"name": "base", "probability_pct": 60}],
|
|
"dated_catalyst_map": [{"catalyst": "earnings", "date_or_window": "Q2"}],
|
|
"invalidation_triggers": [{"trigger": "demand miss"}],
|
|
},
|
|
"investment_debate_state": {
|
|
"bull_history": "bull",
|
|
"bear_history": "bear",
|
|
"history": "debate history",
|
|
"current_response": "current",
|
|
"judge_decision": "judge",
|
|
},
|
|
"trader_investment_plan": "trader plan",
|
|
"risk_debate_state": {
|
|
"aggressive_history": "agg",
|
|
"conservative_history": "cons",
|
|
"neutral_history": "neutral",
|
|
"history": "risk history",
|
|
"judge_decision": "risk judge",
|
|
},
|
|
"investment_plan": "investment plan",
|
|
"final_trade_decision": "buy",
|
|
}
|
|
|
|
graph._log_state("2026-03-24", final_state)
|
|
|
|
output_path = (
|
|
tmp_path
|
|
/ "eval_results"
|
|
/ "AAPL"
|
|
/ "TradingAgentsStrategy_logs"
|
|
/ "full_states_log_2026-03-24.json"
|
|
)
|
|
payload = json.loads(output_path.read_text())
|
|
logged = payload["2026-03-24"]
|
|
|
|
assert logged["scenario_catalyst_report"] == "scenario report"
|
|
assert logged["scenario_catalyst_data"] == final_state["scenario_catalyst_data"]
|