From 36140c674605e0f5aa4db17a874be5f5d921ddaa Mon Sep 17 00:00:00 2001 From: Garrick Date: Tue, 24 Mar 2026 16:29:55 -0700 Subject: [PATCH] refactor: add structured stock underwriting state --- tests/test_stock_role_wiring.py | 111 +++++++++++++++ .../agents/managers/portfolio_manager.py | 24 ++++ .../agents/managers/research_manager.py | 27 +++- tradingagents/agents/utils/agent_states.py | 126 +++++++++++++++++- tradingagents/graph/propagation.py | 2 + 5 files changed, 286 insertions(+), 4 deletions(-) create mode 100644 tests/test_stock_role_wiring.py diff --git a/tests/test_stock_role_wiring.py b/tests/test_stock_role_wiring.py new file mode 100644 index 00000000..777e26df --- /dev/null +++ b/tests/test_stock_role_wiring.py @@ -0,0 +1,111 @@ +from copy import deepcopy + +from tradingagents.agents.managers.portfolio_manager import create_portfolio_manager +from tradingagents.agents.managers.research_manager import create_research_manager +from tradingagents.graph.propagation import Propagator + + +EXPECTED_VALUATION_DATA = { + "fair_value_range": {"low": None, "high": None}, + "expected_return_pct": None, + "primary_method": "", + "thesis": "", +} + +EXPECTED_SEGMENT_DATA = { + "segments": [], + "dominant_segment": "", + "thesis": "", +} + +EXPECTED_SCENARIO_CATALYST_DATA = { + "bull_case": {"probability": None, "price_target": None, "thesis": ""}, + "base_case": {"probability": None, "price_target": None, "thesis": ""}, + "bear_case": {"probability": None, "price_target": None, "thesis": ""}, + "catalysts": [], + "invalidation_triggers": [], +} + +EXPECTED_POSITION_SIZING_DATA = { + "conviction": "", + "target_weight_pct": None, + "initial_weight_pct": None, + "max_loss_pct": None, +} + +EXPECTED_CHIEF_ANALYST_DATA = { + "action": "", + "summary": "", + "thesis": "", + "confidence": "", +} + + +class DummyMemory: + def get_memories(self, _situation, n_matches=2): + return [] + + +class DummyResponse: + def __init__(self, content): + self.content = content + + +class DummyLLM: + def __init__(self, content): + self.content = content + + def invoke(self, _prompt): + return DummyResponse(self.content) + + +def assert_structured_stock_fields(payload): + assert payload["valuation_data"] == EXPECTED_VALUATION_DATA + assert payload["segment_data"] == EXPECTED_SEGMENT_DATA + assert payload["scenario_catalyst_data"] == EXPECTED_SCENARIO_CATALYST_DATA + assert payload["position_sizing_data"] == EXPECTED_POSITION_SIZING_DATA + assert payload["chief_analyst_data"] == EXPECTED_CHIEF_ANALYST_DATA + + +def test_propagator_initializes_structured_stock_underwriting_fields(): + initial_state = Propagator().create_initial_state("NVDA", "2026-03-24") + + assert_structured_stock_fields(initial_state) + + +def test_manager_nodes_preserve_structured_stock_underwriting_fields(monkeypatch): + monkeypatch.setattr( + "tradingagents.agents.managers.research_manager.build_instrument_context", + lambda _ticker: "instrument context", + ) + monkeypatch.setattr( + "tradingagents.agents.managers.portfolio_manager.build_instrument_context", + lambda _ticker: "instrument context", + ) + + state = Propagator().create_initial_state("NVDA", "2026-03-24") + state["investment_plan"] = "Existing investment plan" + + research_manager = create_research_manager( + DummyLLM("Research manager output"), + DummyMemory(), + ) + research_result = research_manager(deepcopy(state)) + + assert research_result["investment_plan"] == "Research manager output" + assert research_result["investment_debate_state"]["judge_decision"] == ( + "Research manager output" + ) + assert_structured_stock_fields(research_result) + + portfolio_manager = create_portfolio_manager( + DummyLLM("Portfolio manager output"), + DummyMemory(), + ) + portfolio_result = portfolio_manager(deepcopy(state)) + + assert portfolio_result["final_trade_decision"] == "Portfolio manager output" + assert portfolio_result["risk_debate_state"]["judge_decision"] == ( + "Portfolio manager output" + ) + assert_structured_stock_fields(portfolio_result) diff --git a/tradingagents/agents/managers/portfolio_manager.py b/tradingagents/agents/managers/portfolio_manager.py index acdf940b..69da4d56 100644 --- a/tradingagents/agents/managers/portfolio_manager.py +++ b/tradingagents/agents/managers/portfolio_manager.py @@ -1,8 +1,12 @@ from tradingagents.agents.utils.agent_utils import build_instrument_context +from tradingagents.agents.utils.agent_states import ( + make_default_structured_stock_underwriting_state, +) def create_portfolio_manager(llm, memory): def portfolio_manager_node(state) -> dict: + structured_stock_defaults = make_default_structured_stock_underwriting_state() instrument_context = build_instrument_context(state["company_of_interest"]) @@ -70,6 +74,26 @@ Be decisive and ground every conclusion in specific evidence from the analysts." return { "risk_debate_state": new_risk_debate_state, "final_trade_decision": response.content, + "valuation_data": state.get( + "valuation_data", + structured_stock_defaults["valuation_data"], + ), + "segment_data": state.get( + "segment_data", + structured_stock_defaults["segment_data"], + ), + "scenario_catalyst_data": state.get( + "scenario_catalyst_data", + structured_stock_defaults["scenario_catalyst_data"], + ), + "position_sizing_data": state.get( + "position_sizing_data", + structured_stock_defaults["position_sizing_data"], + ), + "chief_analyst_data": state.get( + "chief_analyst_data", + structured_stock_defaults["chief_analyst_data"], + ), } return portfolio_manager_node diff --git a/tradingagents/agents/managers/research_manager.py b/tradingagents/agents/managers/research_manager.py index 3ac4b150..884ce196 100644 --- a/tradingagents/agents/managers/research_manager.py +++ b/tradingagents/agents/managers/research_manager.py @@ -1,11 +1,12 @@ -import time -import json - from tradingagents.agents.utils.agent_utils import build_instrument_context +from tradingagents.agents.utils.agent_states import ( + make_default_structured_stock_underwriting_state, +) def create_research_manager(llm, memory): def research_manager_node(state) -> dict: + structured_stock_defaults = make_default_structured_stock_underwriting_state() instrument_context = build_instrument_context(state["company_of_interest"]) history = state["investment_debate_state"].get("history", "") market_research_report = state["market_report"] @@ -55,6 +56,26 @@ Debate History: return { "investment_debate_state": new_investment_debate_state, "investment_plan": response.content, + "valuation_data": state.get( + "valuation_data", + structured_stock_defaults["valuation_data"], + ), + "segment_data": state.get( + "segment_data", + structured_stock_defaults["segment_data"], + ), + "scenario_catalyst_data": state.get( + "scenario_catalyst_data", + structured_stock_defaults["scenario_catalyst_data"], + ), + "position_sizing_data": state.get( + "position_sizing_data", + structured_stock_defaults["position_sizing_data"], + ), + "chief_analyst_data": state.get( + "chief_analyst_data", + structured_stock_defaults["chief_analyst_data"], + ), } return research_manager_node diff --git a/tradingagents/agents/utils/agent_states.py b/tradingagents/agents/utils/agent_states.py index 813b00ee..fbc9b27d 100644 --- a/tradingagents/agents/utils/agent_states.py +++ b/tradingagents/agents/utils/agent_states.py @@ -1,4 +1,4 @@ -from typing import Annotated, Sequence +from typing import Annotated, Any, Sequence from datetime import date, timedelta, datetime from typing_extensions import TypedDict, Optional from langchain_openai import ChatOpenAI @@ -47,6 +47,115 @@ class RiskDebateState(TypedDict): count: Annotated[int, "Length of the current conversation"] # Conversation length +class FairValueRange(TypedDict): + low: Optional[float] + high: Optional[float] + + +class ValuationData(TypedDict): + fair_value_range: FairValueRange + expected_return_pct: Optional[float] + primary_method: str + thesis: str + + +class SegmentData(TypedDict): + segments: list[dict[str, Any]] + dominant_segment: str + thesis: str + + +class ScenarioCaseData(TypedDict): + probability: Optional[float] + price_target: Optional[float] + thesis: str + + +class ScenarioCatalystData(TypedDict): + bull_case: ScenarioCaseData + base_case: ScenarioCaseData + bear_case: ScenarioCaseData + catalysts: list[dict[str, Any]] + invalidation_triggers: list[str] + + +class PositionSizingData(TypedDict): + conviction: str + target_weight_pct: Optional[float] + initial_weight_pct: Optional[float] + max_loss_pct: Optional[float] + + +class ChiefAnalystData(TypedDict): + action: str + summary: str + thesis: str + confidence: str + + +def make_default_valuation_data() -> ValuationData: + return { + "fair_value_range": {"low": None, "high": None}, + "expected_return_pct": None, + "primary_method": "", + "thesis": "", + } + + +def make_default_segment_data() -> SegmentData: + return { + "segments": [], + "dominant_segment": "", + "thesis": "", + } + + +def make_default_scenario_case_data() -> ScenarioCaseData: + return { + "probability": None, + "price_target": None, + "thesis": "", + } + + +def make_default_scenario_catalyst_data() -> ScenarioCatalystData: + return { + "bull_case": make_default_scenario_case_data(), + "base_case": make_default_scenario_case_data(), + "bear_case": make_default_scenario_case_data(), + "catalysts": [], + "invalidation_triggers": [], + } + + +def make_default_position_sizing_data() -> PositionSizingData: + return { + "conviction": "", + "target_weight_pct": None, + "initial_weight_pct": None, + "max_loss_pct": None, + } + + +def make_default_chief_analyst_data() -> ChiefAnalystData: + return { + "action": "", + "summary": "", + "thesis": "", + "confidence": "", + } + + +def make_default_structured_stock_underwriting_state() -> dict[str, Any]: + return { + "valuation_data": make_default_valuation_data(), + "segment_data": make_default_segment_data(), + "scenario_catalyst_data": make_default_scenario_catalyst_data(), + "position_sizing_data": make_default_position_sizing_data(), + "chief_analyst_data": make_default_chief_analyst_data(), + } + + class AgentState(MessagesState): company_of_interest: Annotated[str, "Company that we are interested in trading"] trade_date: Annotated[str, "What date we are trading at"] @@ -60,6 +169,21 @@ class AgentState(MessagesState): str, "Report from the News Researcher of current world affairs" ] fundamentals_report: Annotated[str, "Report from the Fundamentals Researcher"] + valuation_data: Annotated[ + ValuationData, "Structured valuation underwriting output" + ] + segment_data: Annotated[ + SegmentData, "Structured segment underwriting output" + ] + scenario_catalyst_data: Annotated[ + ScenarioCatalystData, "Structured scenario and catalyst underwriting output" + ] + position_sizing_data: Annotated[ + PositionSizingData, "Structured position sizing underwriting output" + ] + chief_analyst_data: Annotated[ + ChiefAnalystData, "Structured chief analyst summary output" + ] # researcher team discussion step investment_debate_state: Annotated[ diff --git a/tradingagents/graph/propagation.py b/tradingagents/graph/propagation.py index 0fd10c0c..75b3be3b 100644 --- a/tradingagents/graph/propagation.py +++ b/tradingagents/graph/propagation.py @@ -5,6 +5,7 @@ from tradingagents.agents.utils.agent_states import ( AgentState, InvestDebateState, RiskDebateState, + make_default_structured_stock_underwriting_state, ) @@ -51,6 +52,7 @@ class Propagator: "fundamentals_report": "", "sentiment_report": "", "news_report": "", + **make_default_structured_stock_underwriting_state(), } def get_graph_args(self, callbacks: Optional[List] = None) -> Dict[str, Any]: