refactor: add structured stock underwriting state

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)

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

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

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[

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]: