mirror of
https://github.com/TauricResearch/TradingAgents.git
synced 2026-05-01 14:33:10 +03:00
feat: structured-output Trader and Research Manager (#434, finishes the trio)
Extends the canonical structured-output pattern from the Portfolio Manager to the other two decision-making agents. Each of the three agents now returns a typed Pydantic instance via llm.with_structured_output() in a single primary call, and a render helper turns the result into the same markdown shape downstream agents and saved reports already consume. - ResearchPlan: 5-tier recommendation, conversational rationale, concrete strategic actions for the trader. - TraderProposal: 3-tier action (transaction direction is naturally Buy / Hold / Sell — position sizing happens later at the Portfolio Manager), reasoning, and optional entry_price / stop_loss / position_sizing. Rendered output preserves the trailing "FINAL TRANSACTION PROPOSAL: **BUY/HOLD/SELL**" line for backward compatibility with the analyst stop-signal text. - PortfolioDecision: 5-tier rating, executive summary, investment thesis, optional price_target / time_horizon (unchanged). The shared try-structured-then-fallback pattern is extracted into tradingagents/agents/utils/structured.py (bind_structured + invoke_structured_or_freetext) so all three agents go through the same code path and log the same warning when a provider lacks structured output and the agent falls back to free-text generation. Net effect for users: every saved markdown report (research/manager.md, trading/trader.md, portfolio/decision.md) now has consistent section headers across runs and providers, easier to scan. Net effect for the runtime: the rating extraction round-trip is gone — the rating comes from the structured response itself, not a second LLM call. SignalProcessor was already simplified to a heuristic adapter in the previous commit. 11 new tests in tests/test_structured_agents.py cover the Trader and Research Manager render functions, structured-output happy paths, and free-text fallback. Full suite: 88 tests pass in ~2s without API keys.
This commit is contained in:
232
tests/test_structured_agents.py
Normal file
232
tests/test_structured_agents.py
Normal file
@@ -0,0 +1,232 @@
|
||||
"""Tests for structured-output agents (Trader and Research Manager).
|
||||
|
||||
The Portfolio Manager has its own coverage in tests/test_memory_log.py
|
||||
(which exercises the full memory-log → PM injection cycle). This file
|
||||
covers the parallel schemas, render functions, and graceful-fallback
|
||||
behavior we added for the Trader and Research Manager so all three
|
||||
decision-making agents share the same shape.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from tradingagents.agents.managers.research_manager import create_research_manager
|
||||
from tradingagents.agents.schemas import (
|
||||
PortfolioRating,
|
||||
ResearchPlan,
|
||||
TraderAction,
|
||||
TraderProposal,
|
||||
render_research_plan,
|
||||
render_trader_proposal,
|
||||
)
|
||||
from tradingagents.agents.trader.trader import create_trader
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Render functions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestRenderTraderProposal:
|
||||
def test_minimal_required_fields(self):
|
||||
p = TraderProposal(action=TraderAction.HOLD, reasoning="Balanced setup; no edge.")
|
||||
md = render_trader_proposal(p)
|
||||
assert "**Action**: Hold" in md
|
||||
assert "**Reasoning**: Balanced setup; no edge." in md
|
||||
# The trailing FINAL TRANSACTION PROPOSAL line is preserved for the
|
||||
# analyst stop-signal text and any external code that greps for it.
|
||||
assert "FINAL TRANSACTION PROPOSAL: **HOLD**" in md
|
||||
|
||||
def test_optional_fields_included_when_present(self):
|
||||
p = TraderProposal(
|
||||
action=TraderAction.BUY,
|
||||
reasoning="Strong technicals + fundamentals.",
|
||||
entry_price=189.5,
|
||||
stop_loss=178.0,
|
||||
position_sizing="6% of portfolio",
|
||||
)
|
||||
md = render_trader_proposal(p)
|
||||
assert "**Action**: Buy" in md
|
||||
assert "**Entry Price**: 189.5" in md
|
||||
assert "**Stop Loss**: 178.0" in md
|
||||
assert "**Position Sizing**: 6% of portfolio" in md
|
||||
assert "FINAL TRANSACTION PROPOSAL: **BUY**" in md
|
||||
|
||||
def test_optional_fields_omitted_when_absent(self):
|
||||
p = TraderProposal(action=TraderAction.SELL, reasoning="Guidance cut.")
|
||||
md = render_trader_proposal(p)
|
||||
assert "Entry Price" not in md
|
||||
assert "Stop Loss" not in md
|
||||
assert "Position Sizing" not in md
|
||||
assert "FINAL TRANSACTION PROPOSAL: **SELL**" in md
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestRenderResearchPlan:
|
||||
def test_required_fields(self):
|
||||
p = ResearchPlan(
|
||||
recommendation=PortfolioRating.OVERWEIGHT,
|
||||
rationale="Bull case carried; tailwinds intact.",
|
||||
strategic_actions="Build position over two weeks; cap at 5%.",
|
||||
)
|
||||
md = render_research_plan(p)
|
||||
assert "**Recommendation**: Overweight" in md
|
||||
assert "**Rationale**: Bull case carried" in md
|
||||
assert "**Strategic Actions**: Build position" in md
|
||||
|
||||
def test_all_5_tier_ratings_render(self):
|
||||
for rating in PortfolioRating:
|
||||
p = ResearchPlan(
|
||||
recommendation=rating,
|
||||
rationale="r",
|
||||
strategic_actions="s",
|
||||
)
|
||||
md = render_research_plan(p)
|
||||
assert f"**Recommendation**: {rating.value}" in md
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Trader agent: structured happy path + fallback
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_trader_state():
|
||||
return {
|
||||
"company_of_interest": "NVDA",
|
||||
"investment_plan": "**Recommendation**: Buy\n**Rationale**: ...\n**Strategic Actions**: ...",
|
||||
}
|
||||
|
||||
|
||||
def _structured_trader_llm(captured: dict, proposal: TraderProposal | None = None):
|
||||
"""Build a MagicMock LLM whose with_structured_output binding captures the
|
||||
prompt and returns a real TraderProposal so render_trader_proposal works.
|
||||
"""
|
||||
if proposal is None:
|
||||
proposal = TraderProposal(
|
||||
action=TraderAction.BUY,
|
||||
reasoning="Strong setup.",
|
||||
)
|
||||
structured = MagicMock()
|
||||
structured.invoke.side_effect = lambda prompt: (
|
||||
captured.__setitem__("prompt", prompt) or proposal
|
||||
)
|
||||
llm = MagicMock()
|
||||
llm.with_structured_output.return_value = structured
|
||||
return llm
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestTraderAgent:
|
||||
def test_structured_path_produces_rendered_markdown(self):
|
||||
captured = {}
|
||||
proposal = TraderProposal(
|
||||
action=TraderAction.BUY,
|
||||
reasoning="AI capex cycle intact; institutional flows constructive.",
|
||||
entry_price=189.5,
|
||||
stop_loss=178.0,
|
||||
position_sizing="6% of portfolio",
|
||||
)
|
||||
llm = _structured_trader_llm(captured, proposal)
|
||||
trader = create_trader(llm)
|
||||
result = trader(_make_trader_state())
|
||||
plan = result["trader_investment_plan"]
|
||||
assert "**Action**: Buy" in plan
|
||||
assert "**Entry Price**: 189.5" in plan
|
||||
assert "FINAL TRANSACTION PROPOSAL: **BUY**" in plan
|
||||
# The same rendered markdown is also added to messages for downstream agents.
|
||||
assert plan in result["messages"][0].content
|
||||
|
||||
def test_prompt_includes_investment_plan(self):
|
||||
captured = {}
|
||||
llm = _structured_trader_llm(captured)
|
||||
trader = create_trader(llm)
|
||||
trader(_make_trader_state())
|
||||
# The investment plan is in the user message of the captured prompt.
|
||||
prompt = captured["prompt"]
|
||||
assert any("Proposed Investment Plan" in m["content"] for m in prompt)
|
||||
|
||||
def test_falls_back_to_freetext_when_structured_unavailable(self):
|
||||
plain_response = (
|
||||
"**Action**: Sell\n\nGuidance cut hits margins.\n\n"
|
||||
"FINAL TRANSACTION PROPOSAL: **SELL**"
|
||||
)
|
||||
llm = MagicMock()
|
||||
llm.with_structured_output.side_effect = NotImplementedError("provider unsupported")
|
||||
llm.invoke.return_value = MagicMock(content=plain_response)
|
||||
trader = create_trader(llm)
|
||||
result = trader(_make_trader_state())
|
||||
assert result["trader_investment_plan"] == plain_response
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Research Manager agent: structured happy path + fallback
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_rm_state():
|
||||
return {
|
||||
"company_of_interest": "NVDA",
|
||||
"investment_debate_state": {
|
||||
"history": "Bull and bear arguments here.",
|
||||
"bull_history": "Bull says...",
|
||||
"bear_history": "Bear says...",
|
||||
"current_response": "",
|
||||
"judge_decision": "",
|
||||
"count": 1,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _structured_rm_llm(captured: dict, plan: ResearchPlan | None = None):
|
||||
if plan is None:
|
||||
plan = ResearchPlan(
|
||||
recommendation=PortfolioRating.HOLD,
|
||||
rationale="Balanced view across both sides.",
|
||||
strategic_actions="Hold current position; reassess after earnings.",
|
||||
)
|
||||
structured = MagicMock()
|
||||
structured.invoke.side_effect = lambda prompt: (
|
||||
captured.__setitem__("prompt", prompt) or plan
|
||||
)
|
||||
llm = MagicMock()
|
||||
llm.with_structured_output.return_value = structured
|
||||
return llm
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestResearchManagerAgent:
|
||||
def test_structured_path_produces_rendered_markdown(self):
|
||||
captured = {}
|
||||
plan = ResearchPlan(
|
||||
recommendation=PortfolioRating.OVERWEIGHT,
|
||||
rationale="Bull case is stronger; AI tailwind intact.",
|
||||
strategic_actions="Build position gradually over two weeks.",
|
||||
)
|
||||
llm = _structured_rm_llm(captured, plan)
|
||||
rm = create_research_manager(llm)
|
||||
result = rm(_make_rm_state())
|
||||
ip = result["investment_plan"]
|
||||
assert "**Recommendation**: Overweight" in ip
|
||||
assert "**Rationale**: Bull case" in ip
|
||||
assert "**Strategic Actions**: Build position" in ip
|
||||
|
||||
def test_prompt_uses_5_tier_rating_scale(self):
|
||||
"""The RM prompt must list all five tiers so the schema enum matches user expectations."""
|
||||
captured = {}
|
||||
llm = _structured_rm_llm(captured)
|
||||
rm = create_research_manager(llm)
|
||||
rm(_make_rm_state())
|
||||
prompt = captured["prompt"]
|
||||
for tier in ("Buy", "Overweight", "Hold", "Underweight", "Sell"):
|
||||
assert f"**{tier}**" in prompt, f"missing {tier} in prompt"
|
||||
|
||||
def test_falls_back_to_freetext_when_structured_unavailable(self):
|
||||
plain_response = "**Recommendation**: Sell\n\n**Rationale**: ...\n\n**Strategic Actions**: ..."
|
||||
llm = MagicMock()
|
||||
llm.with_structured_output.side_effect = NotImplementedError("provider unsupported")
|
||||
llm.invoke.return_value = MagicMock(content=plain_response)
|
||||
rm = create_research_manager(llm)
|
||||
result = rm(_make_rm_state())
|
||||
assert result["investment_plan"] == plain_response
|
||||
Reference in New Issue
Block a user