mirror of
https://github.com/simonlin1212/TradingAgents-astock.git
synced 2026-08-31 01:23:38 +00:00
d55820c08d
BREAKING CHANGE: 上一提交引入的 enable_execution_levels 开关一并移除。 改为直接删除而非默认关闭：荐股软件的认定看软件是否「具备」该功能， 留一个开关在代码里、README 还写着怎么打开，那软件依然具备该功能。 删掉才是真的不具备，同时少一个开关、少两个 schema 变体、少两处分支。 - 删除 TraderProposalWithLevels / PortfolioDecisionWithTarget 两个变体 与 trader_proposal_model() / portfolio_decision_model() 选择器 - 删除 entry_price / stop_loss / position_sizing / price_target 字段 - 渲染函数不再输出对应四节，getattr 兼容层一并移除 - create_trader / create_portfolio_manager / GraphSetup 去掉开关参数 - default_config 去掉 enable_execution_levels - 提示词保持收紧（仅删字段挡不住模型写进散文字段） - 测试 TestExecutionLevelsFlag → TestNoExecutionLevels： 锁定「schema 无价位字段 / 提示词禁止 / 渲染永不输出」 需要该能力的使用者可自行 fork 添加（Apache-2.0 允许）。 测试：161 passed + 48 subtests passed。
240 lines
9.1 KiB
Python
240 lines
9.1 KiB
Python
"""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_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 TestNoExecutionLevels:
|
|
"""The framework ships no executable price levels at all — not a default,
|
|
not a flag. See TraderProposal's docstring for why."""
|
|
|
|
def test_schema_has_no_level_fields(self):
|
|
fields = TraderProposal.model_fields
|
|
assert "action" in fields and "reasoning" in fields
|
|
for f in ("entry_price", "stop_loss", "position_sizing"):
|
|
assert f not in fields
|
|
|
|
def test_prompt_forbids_levels(self):
|
|
captured = {}
|
|
trader = create_trader(_structured_trader_llm(captured))
|
|
trader(_make_trader_state())
|
|
system = next(m["content"] for m in captured["prompt"] if m["role"] == "system")
|
|
assert "Do NOT state entry prices" in system
|
|
|
|
def test_render_never_emits_levels(self):
|
|
p = TraderProposal(action=TraderAction.BUY, reasoning="Trend intact.")
|
|
md = render_trader_proposal(p)
|
|
for label in ("Entry Price", "Stop Loss", "Position Sizing"):
|
|
assert label not in md
|
|
|
|
|
|
@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.",
|
|
)
|
|
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 "**Reasoning**: AI capex cycle intact" 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
|