diff --git a/tests/test_portfolio_manager_state_key.py b/tests/test_portfolio_manager_state_key.py new file mode 100644 index 00000000..bf1ec367 --- /dev/null +++ b/tests/test_portfolio_manager_state_key.py @@ -0,0 +1,175 @@ +"""Tests for Portfolio Manager reading the correct state key. + +The Portfolio Manager must use state["trader_investment_plan"] (the Trader's +memory-refined output), not state["investment_plan"] (the Research Manager's +raw output). Using the wrong key means the Trader agent's entire contribution +— including lessons from past trades — is silently bypassed. +""" + +import types +import unittest +from unittest.mock import MagicMock + + +class TestPortfolioManagerStateKey(unittest.TestCase): + """Verify Portfolio Manager reads trader_investment_plan, not investment_plan.""" + + def _make_state(self, trader_plan="TRADER PLAN: BUY AAPL", research_plan="RESEARCH PLAN: HOLD AAPL"): + """Build a minimal AgentState-like dict for testing.""" + return { + "company_of_interest": "AAPL", + "trade_date": "2025-01-15", + "market_report": "Market is bullish with strong momentum.", + "sentiment_report": "Social sentiment is positive.", + "news_report": "No significant negative news.", + "fundamentals_report": "Strong earnings, low P/E ratio.", + "investment_plan": research_plan, # Research Manager's output + "trader_investment_plan": trader_plan, # Trader's refined output + "risk_debate_state": { + "history": "Aggressive: Go all in. Conservative: Be careful.", + "aggressive_history": "Aggressive: Go all in.", + "conservative_history": "Conservative: Be careful.", + "neutral_history": "Neutral: Balance risk.", + "latest_speaker": "Neutral", + "current_aggressive_response": "Aggressive: Go all in.", + "current_conservative_response": "Conservative: Be careful.", + "current_neutral_response": "Neutral: Balance risk.", + "judge_decision": "", + "count": 3, + }, + } + + def test_portfolio_manager_uses_trader_plan_not_research_plan(self): + """The prompt sent to the LLM must contain the Trader's plan, not the Research Manager's.""" + from tradingagents.agents.managers.portfolio_manager import create_portfolio_manager + + mock_llm = MagicMock() + mock_response = MagicMock() + mock_response.content = "Rating: Buy. The trader's analysis is compelling." + mock_llm.invoke.return_value = mock_response + + mock_memory = MagicMock() + mock_memory.get_memories.return_value = [] + + node = create_portfolio_manager(mock_llm, mock_memory) + + state = self._make_state( + trader_plan="UNIQUE_TRADER_MARKER: BUY with 60% position", + research_plan="UNIQUE_RESEARCH_MARKER: HOLD with caution", + ) + + node(state) + + # The LLM should have been called with the Trader's plan + call_args = mock_llm.invoke.call_args + prompt_text = call_args[0][0] if call_args[0] else str(call_args) + + self.assertIn("UNIQUE_TRADER_MARKER", prompt_text, + "Portfolio Manager must use trader_investment_plan (Trader's output), " + "not investment_plan (Research Manager's output)") + self.assertNotIn("UNIQUE_RESEARCH_MARKER", prompt_text, + "Portfolio Manager should NOT contain Research Manager's raw plan") + + def test_portfolio_manager_returns_final_trade_decision(self): + """Portfolio Manager must return final_trade_decision in its output.""" + from tradingagents.agents.managers.portfolio_manager import create_portfolio_manager + + mock_llm = MagicMock() + mock_response = MagicMock() + mock_response.content = "Final decision: BUY AAPL at market open." + mock_llm.invoke.return_value = mock_response + + mock_memory = MagicMock() + mock_memory.get_memories.return_value = [] + + node = create_portfolio_manager(mock_llm, mock_memory) + state = self._make_state() + result = node(state) + + self.assertIn("final_trade_decision", result) + self.assertEqual(result["final_trade_decision"], "Final decision: BUY AAPL at market open.") + + def test_portfolio_manager_preserves_risk_debate_state(self): + """Portfolio Manager must preserve all risk_debate_state fields.""" + from tradingagents.agents.managers.portfolio_manager import create_portfolio_manager + + mock_llm = MagicMock() + mock_response = MagicMock() + mock_response.content = "Decision rendered." + mock_llm.invoke.return_value = mock_response + + mock_memory = MagicMock() + mock_memory.get_memories.return_value = [] + + node = create_portfolio_manager(mock_llm, mock_memory) + state = self._make_state() + result = node(state) + + rds = result["risk_debate_state"] + self.assertEqual(rds["judge_decision"], "Decision rendered.") + self.assertEqual(rds["latest_speaker"], "Judge") + # Original histories must be preserved + self.assertEqual(rds["aggressive_history"], "Aggressive: Go all in.") + self.assertEqual(rds["conservative_history"], "Conservative: Be careful.") + self.assertEqual(rds["neutral_history"], "Neutral: Balance risk.") + self.assertEqual(rds["count"], 3) + + def test_distinct_trader_and_research_plans_both_present_in_state(self): + """Ensure the test state correctly models both plan types as distinct values.""" + state = self._make_state( + trader_plan="Trader says BUY aggressively", + research_plan="Research says HOLD cautiously", + ) + self.assertNotEqual(state["investment_plan"], state["trader_investment_plan"], + "Test setup: plans must be distinct to catch the wrong-key bug") + + def test_portfolio_manager_with_memory_recommendations(self): + """Past memory recommendations should appear in the prompt.""" + from tradingagents.agents.managers.portfolio_manager import create_portfolio_manager + + mock_llm = MagicMock() + mock_response = MagicMock() + mock_response.content = "Decision with memory context." + mock_llm.invoke.return_value = mock_response + + mock_memory = MagicMock() + mock_memory.get_memories.return_value = [ + {"recommendation": "PAST_LESSON_ALPHA: avoid overleveraging", "similarity_score": 0.9}, + {"recommendation": "PAST_LESSON_BETA: check earnings date", "similarity_score": 0.7}, + ] + + node = create_portfolio_manager(mock_llm, mock_memory) + state = self._make_state() + node(state) + + prompt_text = mock_llm.invoke.call_args[0][0] + self.assertIn("PAST_LESSON_ALPHA", prompt_text) + self.assertIn("PAST_LESSON_BETA", prompt_text) + + +class TestStockStatsBulkImport(unittest.TestCase): + """Verify _get_stock_stats_bulk has access to pandas.""" + + def test_pandas_imported_in_y_finance(self): + """y_finance.py must import pandas for _get_stock_stats_bulk to work.""" + import importlib + import tradingagents.dataflows.y_finance as yf_module + + # pd must be accessible in the module namespace + self.assertTrue( + hasattr(yf_module, 'pd'), + "_get_stock_stats_bulk uses pd.isna() but pandas is not imported as 'pd' " + "in y_finance.py, causing NameError on every call" + ) + + def test_pd_isna_callable(self): + """pd.isna must be callable from y_finance module scope.""" + import tradingagents.dataflows.y_finance as yf_module + pd_ref = getattr(yf_module, 'pd', None) + self.assertIsNotNone(pd_ref, "pd not found in y_finance module") + self.assertTrue(callable(getattr(pd_ref, 'isna', None)), + "pd.isna must be callable") + + +if __name__ == "__main__": + unittest.main() diff --git a/tradingagents/agents/managers/portfolio_manager.py b/tradingagents/agents/managers/portfolio_manager.py index acdf940b..dc2da2f4 100644 --- a/tradingagents/agents/managers/portfolio_manager.py +++ b/tradingagents/agents/managers/portfolio_manager.py @@ -12,7 +12,7 @@ def create_portfolio_manager(llm, memory): news_report = state["news_report"] fundamentals_report = state["fundamentals_report"] sentiment_report = state["sentiment_report"] - trader_plan = state["investment_plan"] + trader_plan = state["trader_investment_plan"] curr_situation = f"{market_research_report}\n\n{sentiment_report}\n\n{news_report}\n\n{fundamentals_report}" past_memories = memory.get_memories(curr_situation, n_matches=2) diff --git a/tradingagents/dataflows/y_finance.py b/tradingagents/dataflows/y_finance.py index 3682a01d..d70edb60 100644 --- a/tradingagents/dataflows/y_finance.py +++ b/tradingagents/dataflows/y_finance.py @@ -1,6 +1,7 @@ from typing import Annotated from datetime import datetime from dateutil.relativedelta import relativedelta +import pandas as pd import yfinance as yf import os from .stockstats_utils import StockstatsUtils, _clean_dataframe, yf_retry