TradingAgents/tradingagents/agents/utils/agent_states.py

77 lines
2.7 KiB
Python

from typing import Annotated, Sequence
from datetime import date, timedelta, datetime
from typing_extensions import TypedDict, Optional
from langchain_openai import ChatOpenAI
from tradingagents.agents import *
from langgraph.prebuilt import ToolNode
from langgraph.graph import END, StateGraph, START, MessagesState
class InvestDebateState(TypedDict):
"""Researcher team state"""
bull_history: Annotated[
str, "Bullish Conversation history"
]
bear_history: Annotated[
str, "Bearish Conversation history"
]
history: Annotated[str, "Conversation history"]
current_response: Annotated[str, "Latest response"]
judge_decision: Annotated[str, "Final judge decision"]
count: Annotated[int, "Length of the current conversation"]
class RiskDebateState(TypedDict):
"""Risk management team state"""
risky_history: Annotated[
str, "Risky Agent's Conversation history"
]
safe_history: Annotated[
str, "Safe Agent's Conversation history"
]
neutral_history: Annotated[
str, "Neutral Agent's Conversation history"
]
history: Annotated[str, "Conversation history"]
latest_speaker: Annotated[str, "Analyst that spoke last"]
current_risky_response: Annotated[
str, "Latest response by the risky analyst"
]
current_safe_response: Annotated[
str, "Latest response by the safe analyst"
]
current_neutral_response: Annotated[
str, "Latest response by the neutral analyst"
]
judge_decision: Annotated[str, "Judge's decision"]
count: Annotated[int, "Length of the current conversation"]
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
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"]
# research
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 mgmt
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"]