TradingAgents/tradingagents/agents/utils/agent_states.py

71 lines
3.4 KiB
Python

from typing import Annotated
from langgraph.graph import MessagesState
from typing_extensions import TypedDict
from tradingagents.agents import *
# Researcher team state
class InvestDebateState(TypedDict):
bull_history: Annotated[str, "Bullish Conversation history"] # Bullish Conversation history
bear_history: Annotated[str, "Bearish Conversation history"] # Bullish Conversation history
history: Annotated[str, "Conversation history"] # Conversation history
current_response: Annotated[str, "Latest response"] # Last response
judge_decision: Annotated[str, "Final judge decision"] # Last response
count: Annotated[int, "Length of the current conversation"] # Conversation length
# Risk management team state
class RiskDebateState(TypedDict):
risky_history: Annotated[str, "Risky Agent's Conversation history"] # Conversation history
safe_history: Annotated[str, "Safe Agent's Conversation history"] # Conversation history
neutral_history: Annotated[str, "Neutral Agent's Conversation history"] # Conversation history
history: Annotated[str, "Conversation history"] # Conversation history
latest_speaker: Annotated[str, "Analyst that spoke last"]
current_risky_response: Annotated[str, "Latest response by the risky analyst"] # Last response
current_safe_response: Annotated[str, "Latest response by the safe analyst"] # Last response
current_neutral_response: Annotated[
str, "Latest response by the neutral analyst"
] # Last response
judge_decision: Annotated[str, "Judge's decision"]
count: Annotated[int, "Length of the current conversation"] # Conversation length
class AgentState(MessagesState):
company_of_interest: Annotated[str, "Company that we are interested in trading"]
trade_date: Annotated[str, "What date we are trading at"]
sender: Annotated[str, "Agent that sent this message"]
# research step
market_report: Annotated[str, "Report from the Market Analyst"]
sentiment_report: Annotated[str, "Report from the Social Media Analyst"]
news_report: Annotated[str, "Report from the News Researcher of current world affairs"]
fundamentals_report: Annotated[str, "Report from the Fundamentals Researcher"]
# researcher team discussion step
investment_debate_state: Annotated[
InvestDebateState, "Current state of the debate on if to invest or not"
]
investment_plan: Annotated[str, "Plan generated by the Analyst"]
trader_investment_plan: Annotated[str, "Plan generated by the Trader"]
# risk management team discussion step
risk_debate_state: Annotated[RiskDebateState, "Current state of the debate on evaluating risk"]
final_trade_decision: Annotated[str, "Final decision made by the Risk Analysts"]
class DiscoveryState(TypedDict):
trade_date: Annotated[str, "Current trading date (yyyy-mm-dd format)"]
tickers: Annotated[list[str], "List of tickers found"]
candidate_metadata: Annotated[list[dict], "Metadata for candidates (source, strategy)"]
filtered_tickers: Annotated[list[str], "List of tickers after filtering"]
opportunities: Annotated[list[dict], "List of final opportunities with rationale"]
final_ranking: Annotated[str, "Final ranking from LLM"]
status: Annotated[str, "Current status of discovery"]
tool_logs: Annotated[
list[dict], "Detailed logs of all tool calls across all nodes (scanner, filter, deep_dive)"
]