# TradingAgents/graph/trading_graph.py import os from pathlib import Path import json from datetime import date from typing import Dict, Any, Tuple, List, Optional from langgraph.prebuilt import ToolNode from tradingagents.llm_clients import create_llm_client from tradingagents.agents import * from tradingagents.default_config import DEFAULT_CONFIG from tradingagents.agents.utils.memory import FinancialSituationMemory from tradingagents.agents.utils.agent_states import ( AgentState, InvestDebateState, RiskDebateState, ) from tradingagents.dataflows.config import set_config # Import the new abstract tool methods from agent_utils from tradingagents.agents.utils.agent_utils import ( get_stock_data, get_indicators, get_fundamentals, get_balance_sheet, get_cashflow, get_income_statement, get_news, get_insider_transactions, get_global_news ) from .conditional_logic import ConditionalLogic from .setup import GraphSetup from .propagation import Propagator from .reflection import Reflector from .signal_processing import SignalProcessor class TradingAgentsGraph: """Main class that orchestrates the trading agents framework.""" def __init__( self, selected_analysts=["market", "social", "news", "fundamentals"], debug=False, config: Dict[str, Any] = None, callbacks: Optional[List] = None, ): """Initialize the trading agents graph and components. Args: selected_analysts: List of analyst types to include debug: Whether to run in debug mode config: Configuration dictionary. If None, uses default config callbacks: Optional list of callback handlers (e.g., for tracking LLM/tool stats) """ self.debug = debug self.config = config or DEFAULT_CONFIG self.callbacks = callbacks or [] # Update the interface's config set_config(self.config) # Create necessary directories os.makedirs(self.config["data_cache_dir"], exist_ok=True) os.makedirs(self.config["results_dir"], exist_ok=True) # Initialize LLMs with provider-specific thinking configuration llm_kwargs = self._get_provider_kwargs() # Add callbacks to kwargs if provided (passed to LLM constructor) if self.callbacks: llm_kwargs["callbacks"] = self.callbacks deep_client = create_llm_client( provider=self.config["llm_provider"], model=self.config["deep_think_llm"], base_url=self.config.get("backend_url"), **llm_kwargs, ) quick_client = create_llm_client( provider=self.config["llm_provider"], model=self.config["quick_think_llm"], base_url=self.config.get("backend_url"), **llm_kwargs, ) self.deep_thinking_llm = deep_client.get_llm() self.quick_thinking_llm = quick_client.get_llm() # Initialize memories self.bull_memory = FinancialSituationMemory("bull_memory", self.config) self.bear_memory = FinancialSituationMemory("bear_memory", self.config) self.trader_memory = FinancialSituationMemory("trader_memory", self.config) self.invest_judge_memory = FinancialSituationMemory("invest_judge_memory", self.config) self.portfolio_manager_memory = FinancialSituationMemory("portfolio_manager_memory", self.config) # Create tool nodes self.tool_nodes = self._create_tool_nodes() # Initialize components self.conditional_logic = ConditionalLogic( max_debate_rounds=self.config["max_debate_rounds"], max_risk_discuss_rounds=self.config["max_risk_discuss_rounds"], ) self.graph_setup = GraphSetup( self.quick_thinking_llm, self.deep_thinking_llm, self.tool_nodes, self.bull_memory, self.bear_memory, self.trader_memory, self.invest_judge_memory, self.portfolio_manager_memory, self.conditional_logic, ) self.propagator = Propagator() self.reflector = Reflector(self.quick_thinking_llm) self.signal_processor = SignalProcessor(self.quick_thinking_llm) # State tracking self.curr_state = None self.ticker = None self.log_states_dict = {} # date to full state dict # Set up the graph self.graph = self.graph_setup.setup_graph(selected_analysts) def _get_provider_kwargs(self) -> Dict[str, Any]: """Get provider-specific kwargs for LLM client creation.""" kwargs = {} provider = self.config.get("llm_provider", "").lower() if provider == "google": thinking_level = self.config.get("google_thinking_level") if thinking_level: kwargs["thinking_level"] = thinking_level elif provider == "openai": reasoning_effort = self.config.get("openai_reasoning_effort") if reasoning_effort: kwargs["reasoning_effort"] = reasoning_effort elif provider == "anthropic": effort = self.config.get("anthropic_effort") if effort: kwargs["effort"] = effort return kwargs def _create_tool_nodes(self) -> Dict[str, ToolNode]: """Create tool nodes for different data sources using abstract methods.""" return { "market": ToolNode( [ # Core stock data tools get_stock_data, # Technical indicators get_indicators, ] ), "social": ToolNode( [ # News tools for social media analysis get_news, ] ), "news": ToolNode( [ # News and insider information get_news, get_global_news, get_insider_transactions, ] ), "fundamentals": ToolNode( [ # Fundamental analysis tools get_fundamentals, get_balance_sheet, get_cashflow, get_income_statement, ] ), } def propagate(self, company_name, trade_date): """Run the trading agents graph for a company on a specific date. Returns: Tuple of (final_state, signal). If budget is exceeded mid-run, returns partial state with ``final_trade_decision`` set to ``"BUDGET_EXCEEDED"`` and signal ``"hold"``. """ from tradingagents.spend_tracker import BudgetExceededError, SpendTracker self.ticker = company_name # Notify spend trackers of ticker start for cb in self.callbacks: if isinstance(cb, SpendTracker): cb.begin_ticker(company_name) # Initialize state init_agent_state = self.propagator.create_initial_state( company_name, trade_date ) args = self.propagator.get_graph_args() try: if self.debug: # Debug mode with tracing trace = [] for chunk in self.graph.stream(init_agent_state, **args): if len(chunk["messages"]) == 0: pass else: chunk["messages"][-1].pretty_print() trace.append(chunk) final_state = trace[-1] if trace else dict(init_agent_state) else: # Standard mode without tracing final_state = self.graph.invoke(init_agent_state, **args) except BudgetExceededError: # Graceful abort: build partial state from what we have final_state = dict(init_agent_state) if self.debug and trace: final_state.update(trace[-1]) final_state.setdefault("final_trade_decision", "BUDGET_EXCEEDED") # Store current state for reflection self.curr_state = final_state # Log state (partial or full) self._log_state(trade_date, final_state) # Log per-ticker and cumulative spend to stderr for cb in self.callbacks: if isinstance(cb, SpendTracker): cb.log_ticker_spend(company_name) cb.log_audit_trail() decision_text = final_state.get("final_trade_decision", "BUDGET_EXCEEDED") return final_state, self.process_signal(decision_text) def _log_state(self, trade_date, final_state): """Log the final state to a JSON file. Tolerates partial state from budget abort.""" invest_debate = final_state.get("investment_debate_state") or {} risk_debate = final_state.get("risk_debate_state") or {} self.log_states_dict[str(trade_date)] = { "company_of_interest": final_state.get("company_of_interest", ""), "trade_date": final_state.get("trade_date", ""), "market_report": final_state.get("market_report", ""), "sentiment_report": final_state.get("sentiment_report", ""), "news_report": final_state.get("news_report", ""), "fundamentals_report": final_state.get("fundamentals_report", ""), "investment_debate_state": { "bull_history": invest_debate.get("bull_history", ""), "bear_history": invest_debate.get("bear_history", ""), "history": invest_debate.get("history", ""), "current_response": invest_debate.get("current_response", ""), "judge_decision": invest_debate.get("judge_decision", ""), }, "trader_investment_decision": final_state.get("trader_investment_plan", ""), "risk_debate_state": { "aggressive_history": risk_debate.get("aggressive_history", ""), "conservative_history": risk_debate.get("conservative_history", ""), "neutral_history": risk_debate.get("neutral_history", ""), "history": risk_debate.get("history", ""), "judge_decision": risk_debate.get("judge_decision", ""), }, "investment_plan": final_state.get("investment_plan", ""), "final_trade_decision": final_state.get("final_trade_decision", ""), } # Save to file directory = Path(self.config["results_dir"]) / self.ticker / "TradingAgentsStrategy_logs" directory.mkdir(parents=True, exist_ok=True) log_path = directory / f"full_states_log_{trade_date}.json" with open(log_path, "w", encoding="utf-8") as f: json.dump(self.log_states_dict[str(trade_date)], f, indent=4) def reflect_and_remember(self, returns_losses): """Reflect on decisions and update memory based on returns.""" self.reflector.reflect_bull_researcher( self.curr_state, returns_losses, self.bull_memory ) self.reflector.reflect_bear_researcher( self.curr_state, returns_losses, self.bear_memory ) self.reflector.reflect_trader( self.curr_state, returns_losses, self.trader_memory ) self.reflector.reflect_invest_judge( self.curr_state, returns_losses, self.invest_judge_memory ) self.reflector.reflect_portfolio_manager( self.curr_state, returns_losses, self.portfolio_manager_memory ) def process_signal(self, full_signal): """Process a signal to extract the core decision.""" return self.signal_processor.process_signal(full_signal)