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:
Yijia-Xiao
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
|
||||
@@ -5,35 +5,24 @@ Uses LangChain's ``with_structured_output`` so the LLM produces a typed
|
||||
back to markdown for storage in ``final_trade_decision`` so memory log,
|
||||
CLI display, and saved reports continue to consume the same shape they do
|
||||
today. When a provider does not expose structured output, the agent falls
|
||||
back to a free-text invocation and the existing heuristic rating parser.
|
||||
back gracefully to free-text generation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from tradingagents.agents.schemas import PortfolioDecision, render_pm_decision
|
||||
from tradingagents.agents.utils.agent_utils import (
|
||||
build_instrument_context,
|
||||
get_language_instruction,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
from tradingagents.agents.utils.structured import (
|
||||
bind_structured,
|
||||
invoke_structured_or_freetext,
|
||||
)
|
||||
|
||||
|
||||
def create_portfolio_manager(llm):
|
||||
# Wrap once at agent construction; if the provider does not support
|
||||
# structured output we keep ``structured_llm`` as None and use the
|
||||
# free-text fallback for every call.
|
||||
try:
|
||||
structured_llm = llm.with_structured_output(PortfolioDecision)
|
||||
except (NotImplementedError, AttributeError) as exc:
|
||||
logger.warning(
|
||||
"Portfolio Manager: provider does not support with_structured_output (%s); "
|
||||
"falling back to free-text generation",
|
||||
exc,
|
||||
)
|
||||
structured_llm = None
|
||||
structured_llm = bind_structured(llm, PortfolioDecision, "Portfolio Manager")
|
||||
|
||||
def portfolio_manager_node(state) -> dict:
|
||||
instrument_context = build_instrument_context(state["company_of_interest"])
|
||||
@@ -74,7 +63,13 @@ def create_portfolio_manager(llm):
|
||||
|
||||
Be decisive and ground every conclusion in specific evidence from the analysts.{get_language_instruction()}"""
|
||||
|
||||
final_trade_decision = _invoke_pm(structured_llm, llm, prompt)
|
||||
final_trade_decision = invoke_structured_or_freetext(
|
||||
structured_llm,
|
||||
llm,
|
||||
prompt,
|
||||
render_pm_decision,
|
||||
"Portfolio Manager",
|
||||
)
|
||||
|
||||
new_risk_debate_state = {
|
||||
"judge_decision": final_trade_decision,
|
||||
@@ -95,26 +90,3 @@ Be decisive and ground every conclusion in specific evidence from the analysts.{
|
||||
}
|
||||
|
||||
return portfolio_manager_node
|
||||
|
||||
|
||||
def _invoke_pm(structured_llm, plain_llm, prompt: str) -> str:
|
||||
"""Run the PM call and return the markdown-rendered decision.
|
||||
|
||||
Tries the structured-output path first; if it fails for any reason
|
||||
(provider does not support it, model returns malformed JSON, network
|
||||
glitch on the structured endpoint), falls back to the plain free-text
|
||||
invocation so the pipeline still produces a result.
|
||||
"""
|
||||
if structured_llm is not None:
|
||||
try:
|
||||
decision = structured_llm.invoke(prompt)
|
||||
return render_pm_decision(decision)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Portfolio Manager: structured-output invocation failed (%s); "
|
||||
"retrying once as free text",
|
||||
exc,
|
||||
)
|
||||
|
||||
response = plain_llm.invoke(prompt)
|
||||
return response.content
|
||||
|
||||
@@ -1,8 +1,18 @@
|
||||
"""Research Manager: turns the bull/bear debate into a structured investment plan for the trader."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from tradingagents.agents.schemas import ResearchPlan, render_research_plan
|
||||
from tradingagents.agents.utils.agent_utils import build_instrument_context
|
||||
from tradingagents.agents.utils.structured import (
|
||||
bind_structured,
|
||||
invoke_structured_or_freetext,
|
||||
)
|
||||
|
||||
|
||||
def create_research_manager(llm):
|
||||
structured_llm = bind_structured(llm, ResearchPlan, "Research Manager")
|
||||
|
||||
def research_manager_node(state) -> dict:
|
||||
instrument_context = build_instrument_context(state["company_of_interest"])
|
||||
history = state["investment_debate_state"].get("history", "")
|
||||
@@ -24,31 +34,31 @@ def create_research_manager(llm):
|
||||
|
||||
Commit to a clear stance whenever the debate's strongest arguments warrant one; reserve Hold for situations where the evidence on both sides is genuinely balanced.
|
||||
|
||||
**Required Output Structure:**
|
||||
1. **Recommendation**: State one of Buy / Overweight / Hold / Underweight / Sell.
|
||||
2. **Rationale**: Summarise the key points from both sides and explain which arguments led to this recommendation.
|
||||
3. **Strategic Actions**: Concrete steps for the trader to implement the recommendation, including position sizing guidance consistent with the rating.
|
||||
|
||||
Present your analysis conversationally, as if speaking naturally to a teammate.
|
||||
|
||||
---
|
||||
|
||||
**Debate History:**
|
||||
{history}"""
|
||||
response = llm.invoke(prompt)
|
||||
|
||||
investment_plan = invoke_structured_or_freetext(
|
||||
structured_llm,
|
||||
llm,
|
||||
prompt,
|
||||
render_research_plan,
|
||||
"Research Manager",
|
||||
)
|
||||
|
||||
new_investment_debate_state = {
|
||||
"judge_decision": response.content,
|
||||
"judge_decision": investment_plan,
|
||||
"history": investment_debate_state.get("history", ""),
|
||||
"bear_history": investment_debate_state.get("bear_history", ""),
|
||||
"bull_history": investment_debate_state.get("bull_history", ""),
|
||||
"current_response": response.content,
|
||||
"current_response": investment_plan,
|
||||
"count": investment_debate_state["count"],
|
||||
}
|
||||
|
||||
return {
|
||||
"investment_debate_state": new_investment_debate_state,
|
||||
"investment_plan": response.content,
|
||||
"investment_plan": investment_plan,
|
||||
}
|
||||
|
||||
return research_manager_node
|
||||
|
||||
@@ -1,15 +1,16 @@
|
||||
"""Pydantic schemas used by agents that produce structured output.
|
||||
|
||||
The framework's primary artifact is still prose: each agent's natural-language
|
||||
reasoning is what users read, what gets stored in the memory log, and what
|
||||
gets saved as markdown reports. Structured output is layered onto agents
|
||||
whose results have downstream machine-readable consumers (currently only
|
||||
the Portfolio Manager) so that:
|
||||
reasoning is what users read in the saved markdown reports and what the
|
||||
downstream agents read as context. Structured output is layered onto the
|
||||
three decision-making agents (Research Manager, Trader, Portfolio Manager)
|
||||
so that:
|
||||
|
||||
- The rating is type-safe and never has to be regex-extracted
|
||||
- Schema field descriptions become the model's output instructions
|
||||
- Their outputs follow consistent section headers across runs and providers
|
||||
- Each provider's native structured-output mode is used (json_schema for
|
||||
OpenAI/xAI, response_schema for Gemini, tool-use for Anthropic)
|
||||
- Schema field descriptions become the model's output instructions, freeing
|
||||
the prompt body to focus on context and the rating-scale guidance
|
||||
- A render helper turns the parsed Pydantic instance back into the same
|
||||
markdown shape the rest of the system already consumes, so display,
|
||||
memory log, and saved reports keep working unchanged
|
||||
@@ -23,8 +24,13 @@ from typing import Optional
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared rating types
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PortfolioRating(str, Enum):
|
||||
"""5-tier portfolio rating used by the Research Manager and Portfolio Manager."""
|
||||
"""5-tier rating used by the Research Manager and Portfolio Manager."""
|
||||
|
||||
BUY = "Buy"
|
||||
OVERWEIGHT = "Overweight"
|
||||
@@ -33,6 +39,135 @@ class PortfolioRating(str, Enum):
|
||||
SELL = "Sell"
|
||||
|
||||
|
||||
class TraderAction(str, Enum):
|
||||
"""3-tier transaction direction used by the Trader.
|
||||
|
||||
The Trader's job is to translate the Research Manager's investment plan
|
||||
into a concrete transaction proposal: should the desk execute a Buy, a
|
||||
Sell, or sit on Hold this round. Position sizing and the nuanced
|
||||
Overweight / Underweight calls happen later at the Portfolio Manager.
|
||||
"""
|
||||
|
||||
BUY = "Buy"
|
||||
HOLD = "Hold"
|
||||
SELL = "Sell"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Research Manager
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ResearchPlan(BaseModel):
|
||||
"""Structured investment plan produced by the Research Manager.
|
||||
|
||||
Hand-off to the Trader: the recommendation pins the directional view,
|
||||
the rationale captures which side of the bull/bear debate carried the
|
||||
argument, and the strategic actions translate that into concrete
|
||||
instructions the trader can execute against.
|
||||
"""
|
||||
|
||||
recommendation: PortfolioRating = Field(
|
||||
description=(
|
||||
"The investment recommendation. Exactly one of Buy / Overweight / "
|
||||
"Hold / Underweight / Sell. Reserve Hold for situations where the "
|
||||
"evidence on both sides is genuinely balanced; otherwise commit to "
|
||||
"the side with the stronger arguments."
|
||||
),
|
||||
)
|
||||
rationale: str = Field(
|
||||
description=(
|
||||
"Conversational summary of the key points from both sides of the "
|
||||
"debate, ending with which arguments led to the recommendation. "
|
||||
"Speak naturally, as if to a teammate."
|
||||
),
|
||||
)
|
||||
strategic_actions: str = Field(
|
||||
description=(
|
||||
"Concrete steps for the trader to implement the recommendation, "
|
||||
"including position sizing guidance consistent with the rating."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def render_research_plan(plan: ResearchPlan) -> str:
|
||||
"""Render a ResearchPlan to markdown for storage and the trader's prompt context."""
|
||||
return "\n".join([
|
||||
f"**Recommendation**: {plan.recommendation.value}",
|
||||
"",
|
||||
f"**Rationale**: {plan.rationale}",
|
||||
"",
|
||||
f"**Strategic Actions**: {plan.strategic_actions}",
|
||||
])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Trader
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TraderProposal(BaseModel):
|
||||
"""Structured transaction proposal produced by the Trader.
|
||||
|
||||
The trader reads the Research Manager's investment plan and the analyst
|
||||
reports, then turns them into a concrete transaction: what action to
|
||||
take, the reasoning that justifies it, and the practical levels for
|
||||
entry, stop-loss, and sizing.
|
||||
"""
|
||||
|
||||
action: TraderAction = Field(
|
||||
description="The transaction direction. Exactly one of Buy / Hold / Sell.",
|
||||
)
|
||||
reasoning: str = Field(
|
||||
description=(
|
||||
"The case for this action, anchored in the analysts' reports and "
|
||||
"the research plan. Two to four sentences."
|
||||
),
|
||||
)
|
||||
entry_price: Optional[float] = Field(
|
||||
default=None,
|
||||
description="Optional entry price target in the instrument's quote currency.",
|
||||
)
|
||||
stop_loss: Optional[float] = Field(
|
||||
default=None,
|
||||
description="Optional stop-loss price in the instrument's quote currency.",
|
||||
)
|
||||
position_sizing: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Optional sizing guidance, e.g. '5% of portfolio'.",
|
||||
)
|
||||
|
||||
|
||||
def render_trader_proposal(proposal: TraderProposal) -> str:
|
||||
"""Render a TraderProposal to markdown.
|
||||
|
||||
The trailing ``FINAL TRANSACTION PROPOSAL: **BUY/HOLD/SELL**`` line is
|
||||
preserved for backward compatibility with the analyst stop-signal text
|
||||
and any external code that greps for it.
|
||||
"""
|
||||
parts = [
|
||||
f"**Action**: {proposal.action.value}",
|
||||
"",
|
||||
f"**Reasoning**: {proposal.reasoning}",
|
||||
]
|
||||
if proposal.entry_price is not None:
|
||||
parts.extend(["", f"**Entry Price**: {proposal.entry_price}"])
|
||||
if proposal.stop_loss is not None:
|
||||
parts.extend(["", f"**Stop Loss**: {proposal.stop_loss}"])
|
||||
if proposal.position_sizing:
|
||||
parts.extend(["", f"**Position Sizing**: {proposal.position_sizing}"])
|
||||
parts.extend([
|
||||
"",
|
||||
f"FINAL TRANSACTION PROPOSAL: **{proposal.action.value.upper()}**",
|
||||
])
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Portfolio Manager
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PortfolioDecision(BaseModel):
|
||||
"""Structured output produced by the Portfolio Manager.
|
||||
|
||||
|
||||
@@ -1,32 +1,60 @@
|
||||
"""Trader: turns the Research Manager's investment plan into a concrete transaction proposal."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
|
||||
from langchain_core.messages import AIMessage
|
||||
|
||||
from tradingagents.agents.schemas import TraderProposal, render_trader_proposal
|
||||
from tradingagents.agents.utils.agent_utils import build_instrument_context
|
||||
from tradingagents.agents.utils.structured import (
|
||||
bind_structured,
|
||||
invoke_structured_or_freetext,
|
||||
)
|
||||
|
||||
|
||||
def create_trader(llm):
|
||||
structured_llm = bind_structured(llm, TraderProposal, "Trader")
|
||||
|
||||
def trader_node(state, name):
|
||||
company_name = state["company_of_interest"]
|
||||
instrument_context = build_instrument_context(company_name)
|
||||
investment_plan = state["investment_plan"]
|
||||
|
||||
context = {
|
||||
"role": "user",
|
||||
"content": f"Based on a comprehensive analysis by a team of analysts, here is an investment plan tailored for {company_name}. {instrument_context} This plan incorporates insights from current technical market trends, macroeconomic indicators, and social media sentiment. Use this plan as a foundation for evaluating your next trading decision.\n\nProposed Investment Plan: {investment_plan}\n\nLeverage these insights to make an informed and strategic decision.",
|
||||
}
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a trading agent analyzing market data to make investment decisions. Based on your analysis, provide a specific recommendation to buy, sell, or hold. End with a firm decision and always conclude your response with 'FINAL TRANSACTION PROPOSAL: **BUY/HOLD/SELL**' to confirm your recommendation.",
|
||||
"content": (
|
||||
"You are a trading agent analyzing market data to make investment decisions. "
|
||||
"Based on your analysis, provide a specific recommendation to buy, sell, or hold. "
|
||||
"Anchor your reasoning in the analysts' reports and the research plan."
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
f"Based on a comprehensive analysis by a team of analysts, here is an investment "
|
||||
f"plan tailored for {company_name}. {instrument_context} This plan incorporates "
|
||||
f"insights from current technical market trends, macroeconomic indicators, and "
|
||||
f"social media sentiment. Use this plan as a foundation for evaluating your next "
|
||||
f"trading decision.\n\nProposed Investment Plan: {investment_plan}\n\n"
|
||||
f"Leverage these insights to make an informed and strategic decision."
|
||||
),
|
||||
},
|
||||
context,
|
||||
]
|
||||
|
||||
result = llm.invoke(messages)
|
||||
trader_plan = invoke_structured_or_freetext(
|
||||
structured_llm,
|
||||
llm,
|
||||
messages,
|
||||
render_trader_proposal,
|
||||
"Trader",
|
||||
)
|
||||
|
||||
return {
|
||||
"messages": [result],
|
||||
"trader_investment_plan": result.content,
|
||||
"messages": [AIMessage(content=trader_plan)],
|
||||
"trader_investment_plan": trader_plan,
|
||||
"sender": name,
|
||||
}
|
||||
|
||||
|
||||
73
tradingagents/agents/utils/structured.py
Normal file
73
tradingagents/agents/utils/structured.py
Normal file
@@ -0,0 +1,73 @@
|
||||
"""Shared helpers for invoking an agent with structured output and a graceful fallback.
|
||||
|
||||
The Portfolio Manager, Trader, and Research Manager all follow the same
|
||||
canonical pattern:
|
||||
|
||||
1. At agent creation, wrap the LLM with ``with_structured_output(Schema)``
|
||||
so the model returns a typed Pydantic instance. If the provider does
|
||||
not support structured output (rare; mostly older Ollama models), the
|
||||
wrap is skipped and the agent uses free-text generation instead.
|
||||
2. At invocation, run the structured call and render the result back to
|
||||
markdown. If the structured call itself fails for any reason
|
||||
(malformed JSON from a weak model, transient provider issue), fall
|
||||
back to a plain ``llm.invoke`` so the pipeline never blocks.
|
||||
|
||||
Centralising the pattern here keeps the agent factories small and ensures
|
||||
all three agents log the same warnings when fallback fires.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any, Callable, Optional, TypeVar
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
T = TypeVar("T", bound=BaseModel)
|
||||
|
||||
|
||||
def bind_structured(llm: Any, schema: type[T], agent_name: str) -> Optional[Any]:
|
||||
"""Return ``llm.with_structured_output(schema)`` or ``None`` if unsupported.
|
||||
|
||||
Logs a warning when the binding fails so the user understands the agent
|
||||
will use free-text generation for every call instead of one-shot fallback.
|
||||
"""
|
||||
try:
|
||||
return llm.with_structured_output(schema)
|
||||
except (NotImplementedError, AttributeError) as exc:
|
||||
logger.warning(
|
||||
"%s: provider does not support with_structured_output (%s); "
|
||||
"falling back to free-text generation",
|
||||
agent_name, exc,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def invoke_structured_or_freetext(
|
||||
structured_llm: Optional[Any],
|
||||
plain_llm: Any,
|
||||
prompt: Any,
|
||||
render: Callable[[T], str],
|
||||
agent_name: str,
|
||||
) -> str:
|
||||
"""Run the structured call and render to markdown; fall back to free-text on any failure.
|
||||
|
||||
``prompt`` is whatever the underlying LLM accepts (a string for chat
|
||||
invocations, a list of message dicts for chat models that take that
|
||||
shape). The same value is forwarded to the free-text path so the
|
||||
fallback sees the same input the structured call did.
|
||||
"""
|
||||
if structured_llm is not None:
|
||||
try:
|
||||
result = structured_llm.invoke(prompt)
|
||||
return render(result)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"%s: structured-output invocation failed (%s); retrying once as free text",
|
||||
agent_name, exc,
|
||||
)
|
||||
|
||||
response = plain_llm.invoke(prompt)
|
||||
return response.content
|
||||
Reference in New Issue
Block a user