diff --git a/tests/test_structured_agents.py b/tests/test_structured_agents.py new file mode 100644 index 00000000..ea771a4b --- /dev/null +++ b/tests/test_structured_agents.py @@ -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 diff --git a/tradingagents/agents/managers/portfolio_manager.py b/tradingagents/agents/managers/portfolio_manager.py index 38ab840a..0e7c1823 100644 --- a/tradingagents/agents/managers/portfolio_manager.py +++ b/tradingagents/agents/managers/portfolio_manager.py @@ -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 diff --git a/tradingagents/agents/managers/research_manager.py b/tradingagents/agents/managers/research_manager.py index 020b719e..0e2206b2 100644 --- a/tradingagents/agents/managers/research_manager.py +++ b/tradingagents/agents/managers/research_manager.py @@ -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 diff --git a/tradingagents/agents/schemas.py b/tradingagents/agents/schemas.py index aadcaf0b..55f0e3cf 100644 --- a/tradingagents/agents/schemas.py +++ b/tradingagents/agents/schemas.py @@ -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. diff --git a/tradingagents/agents/trader/trader.py b/tradingagents/agents/trader/trader.py index 0ecae888..ea3f6b23 100644 --- a/tradingagents/agents/trader/trader.py +++ b/tradingagents/agents/trader/trader.py @@ -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, } diff --git a/tradingagents/agents/utils/structured.py b/tradingagents/agents/utils/structured.py new file mode 100644 index 00000000..400e1a82 --- /dev/null +++ b/tradingagents/agents/utils/structured.py @@ -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