TradingAgents/cli/main.py

362 lines
13 KiB
Python

from typing import Optional
import datetime
import typer
from pathlib import Path
from rich.console import Console
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
from rich.panel import Panel
from rich.spinner import Spinner
from rich.live import Live
from rich.columns import Columns
from rich.markdown import Markdown
from rich.layout import Layout
from rich.text import Text
from rich.live import Live
from rich.table import Table
from collections import deque
import time
from rich.tree import Tree
from rich import box
from rich.align import Align
from rich.rule import Rule
from tradingagents.graph.trading_graph import TradingAgentsGraph
from tradingagents.default_config import DEFAULT_CONFIG
from cli.helpers import AnalystType, update_research_team_status, extract_content_string
from cli.prompts import *
from cli.message_buffer import MessageBuffer
from cli.ui_display import create_layout, update_display
from cli.report_display import display_complete_report
from cli.asset_detection import detect_asset_class, get_asset_class_display_name
from cli.file_handlers import setup_file_handlers
from cli.stream_processor import process_chunk
console = Console()
app = typer.Typer(
name="Litadel",
help="Litadel CLI: Multi-Agents LLM Financial Trading Framework (successor of TradingAgents)",
add_completion=True, # Enable shell completion
)
# Create a global message buffer instance
message_buffer = MessageBuffer()
def get_user_selections():
"""Get all user selections before starting the analysis display."""
# Load config to check for pre-configured values
from tradingagents.default_config import DEFAULT_CONFIG
# Display ASCII art welcome message
with open("./cli/static/welcome.txt", "r") as f:
welcome_ascii = f.read()
# Create welcome box content
welcome_content = f"{welcome_ascii}\n"
welcome_content += "[bold green]Litadel: Multi-Agents LLM Financial Trading Framework - CLI[/bold green]\n\n"
welcome_content += "[dim]Successor of TradingAgents by TaurusResearch[/dim]\n\n"
welcome_content += "[bold]Workflow Steps:[/bold]\n"
welcome_content += "I. Analyst Team → II. Research Team → III. Trader → IV. Risk Management → V. Portfolio Management"
# Create and center the welcome box
welcome_box = Panel(
welcome_content,
border_style="green",
padding=(1, 2),
title="Welcome to Litadel",
subtitle="Multi-Agents LLM Financial Trading Framework",
)
console.print(Align.center(welcome_box))
console.print() # Add a blank line after the welcome box
# Create a boxed questionnaire for each step
def create_question_box(title, prompt, default=None):
box_content = f"[bold]{title}[/bold]\n"
box_content += f"[dim]{prompt}[/dim]"
if default:
box_content += f"\n[dim]Default: {default}[/dim]"
return Panel(box_content, border_style="blue", padding=(1, 2))
# Step 1: Ticker symbol
console.print(
create_question_box(
"Step 1: Ticker Symbol", "Enter the ticker symbol to analyze", "SPY"
)
)
selected_ticker = get_ticker()
# Auto-detect asset class from ticker
asset_class = detect_asset_class(selected_ticker)
console.print(
f"[dim]→ Detected asset class: [bold]{get_asset_class_display_name(asset_class)}[/bold][/dim]\n"
)
# Step 2: Analysis date
default_date = datetime.datetime.now().strftime("%Y-%m-%d")
console.print(
create_question_box(
"Step 2: Analysis Date",
"Enter the analysis date (YYYY-MM-DD)",
default_date,
)
)
analysis_date = get_analysis_date()
# Step 3: Select analysts (use config default if available)
if "default_analysts" in DEFAULT_CONFIG and DEFAULT_CONFIG["default_analysts"]:
# Convert analyst names to AnalystType
analyst_map = {
"market": AnalystType.MARKET,
"news": AnalystType.NEWS,
"social": AnalystType.SOCIAL,
"fundamentals": AnalystType.FUNDAMENTALS,
}
selected_analysts = [analyst_map[a] for a in DEFAULT_CONFIG["default_analysts"] if a in analyst_map]
# Filter out fundamentals for commodities
if asset_class == "commodity":
selected_analysts = [a for a in selected_analysts if a != AnalystType.FUNDAMENTALS]
console.print(
f"[dim]→ Using configured analysts: [bold]{', '.join(a.value for a in selected_analysts)}[/bold][/dim]\n"
)
else:
console.print(
create_question_box(
"Step 3: Analysts Team", "Select your LLM analyst agents for the analysis"
)
)
selected_analysts = select_analysts(asset_class)
console.print(
f"[green]Selected analysts:[/green] {', '.join(analyst.value for analyst in selected_analysts)}"
)
# Step 4: Research depth (use config default if available)
if "max_debate_rounds" in DEFAULT_CONFIG:
selected_research_depth = DEFAULT_CONFIG["max_debate_rounds"]
depth_labels = {1: "Shallow", 2: "Medium", 3: "Deep"}
console.print(
f"[dim]→ Using configured research depth: [bold]{depth_labels.get(selected_research_depth, selected_research_depth)}[/bold][/dim]\n"
)
else:
console.print(
create_question_box(
"Step 4: Research Depth", "Select your research depth level"
)
)
selected_research_depth = select_research_depth()
# Step 5: LLM backend (use config default if available)
if "llm_provider" in DEFAULT_CONFIG and "backend_url" in DEFAULT_CONFIG:
selected_llm_provider = DEFAULT_CONFIG["llm_provider"]
backend_url = DEFAULT_CONFIG["backend_url"]
console.print(
f"[dim]→ Using configured LLM provider: [bold]{selected_llm_provider.upper()}[/bold] ({backend_url})[/dim]\n"
)
else:
console.print(
create_question_box(
"Step 5: LLM Backend", "Select which service to talk to"
)
)
selected_llm_provider, backend_url = select_llm_provider()
# Step 6: Thinking agents (use config defaults if available)
if "quick_think_llm" in DEFAULT_CONFIG and "deep_think_llm" in DEFAULT_CONFIG:
selected_shallow_thinker = DEFAULT_CONFIG["quick_think_llm"]
selected_deep_thinker = DEFAULT_CONFIG["deep_think_llm"]
console.print(
f"[dim]→ Using configured models:[/dim]\n"
f"[dim] Quick thinking: [bold]{selected_shallow_thinker}[/bold][/dim]\n"
f"[dim] Deep thinking: [bold]{selected_deep_thinker}[/bold][/dim]\n"
)
else:
console.print(
create_question_box(
"Step 6: Thinking Agents", "Select your thinking agents for analysis"
)
)
selected_shallow_thinker = select_shallow_thinking_agent(selected_llm_provider)
selected_deep_thinker = select_deep_thinking_agent(selected_llm_provider)
return {
"ticker": selected_ticker,
"analysis_date": analysis_date,
"asset_class": asset_class,
"analysts": selected_analysts,
"research_depth": selected_research_depth,
"llm_provider": selected_llm_provider.lower(),
"backend_url": backend_url,
"shallow_thinker": selected_shallow_thinker,
"deep_thinker": selected_deep_thinker,
}
def get_ticker():
"""Get ticker symbol from user input."""
return typer.prompt("", default="SPY")
def get_analysis_date():
"""Get the analysis date from user input."""
while True:
date_str = typer.prompt(
"", default=datetime.datetime.now().strftime("%Y-%m-%d")
)
try:
# Validate date format and ensure it's not in the future
analysis_date = datetime.datetime.strptime(date_str, "%Y-%m-%d")
if analysis_date.date() > datetime.datetime.now().date():
console.print("[red]Error: Analysis date cannot be in the future[/red]")
continue
return date_str
except ValueError:
console.print(
"[red]Error: Invalid date format. Please use YYYY-MM-DD[/red]"
)
def setup_config(selections):
"""Create and configure the analysis config from user selections."""
config = DEFAULT_CONFIG.copy()
config["max_debate_rounds"] = selections["research_depth"]
config["max_risk_discuss_rounds"] = selections["research_depth"]
config["quick_think_llm"] = selections["shallow_thinker"]
config["deep_think_llm"] = selections["deep_thinker"]
config["backend_url"] = selections["backend_url"]
config["llm_provider"] = selections["llm_provider"].lower()
config["asset_class"] = selections["asset_class"]
return config
def setup_analysis(selections, config):
"""Initialize the trading graph and file handlers."""
# Initialize the graph
graph = TradingAgentsGraph(
[analyst.value for analyst in selections["analysts"]], config=config, debug=True
)
# Create result directory and setup file handlers
results_dir = Path(config["results_dir"]) / selections["ticker"] / selections["analysis_date"]
log_file, report_dir = setup_file_handlers(message_buffer, results_dir)
return graph, log_file, report_dir
def initialize_display(layout, selections):
"""Initialize the display with startup messages and agent statuses."""
# Initial display
update_display(layout, message_buffer)
# Add initial messages
message_buffer.add_message("System", f"Selected ticker: {selections['ticker']}")
message_buffer.add_message("System", f"Analysis date: {selections['analysis_date']}")
message_buffer.add_message(
"System",
f"Selected analysts: {', '.join(analyst.value for analyst in selections['analysts'])}",
)
update_display(layout, message_buffer)
# Reset agent statuses
for agent in message_buffer.agent_status:
message_buffer.update_agent_status(agent, "pending")
# Reset report sections
for section in message_buffer.report_sections:
message_buffer.report_sections[section] = None
message_buffer.current_report = None
message_buffer.final_report = None
# Update agent status to in_progress for the first analyst
first_analyst = f"{selections['analysts'][0].value.capitalize()} Analyst"
message_buffer.update_agent_status(first_analyst, "in_progress")
update_display(layout, message_buffer)
# Create spinner text
spinner_text = f"Analyzing {selections['ticker']} on {selections['analysis_date']}..."
update_display(layout, message_buffer, spinner_text)
def run_stream_analysis(graph, selections, layout):
"""Stream the analysis and process chunks in real-time."""
# Initialize state and get graph args
init_agent_state = graph.propagator.create_initial_state(
selections["ticker"], selections["analysis_date"]
)
# CRITICAL: Add asset_class to state so market analyst can branch correctly
init_agent_state["asset_class"] = selections["asset_class"]
args = graph.propagator.get_graph_args()
# Stream the analysis
trace = []
for chunk in graph.graph.stream(init_agent_state, **args):
# Process the chunk and update message buffer
if process_chunk(chunk, message_buffer, selections["analysts"]):
# Update the display after successful chunk processing
update_display(layout, message_buffer)
trace.append(chunk)
return trace
def finalize_analysis(trace, graph, selections, layout):
"""Process final results and display the complete report."""
# Get final state and decision
final_state = trace[-1]
decision = graph.process_signal(final_state["final_trade_decision"])
# Update all agent statuses to completed
for agent in message_buffer.agent_status:
message_buffer.update_agent_status(agent, "completed")
message_buffer.add_message(
"Analysis", f"Completed analysis for {selections['analysis_date']}"
)
# Update final report sections
for section in message_buffer.report_sections.keys():
if section in final_state:
message_buffer.update_report_section(section, final_state[section])
# Display the complete final report
display_complete_report(final_state)
update_display(layout, message_buffer)
def run_analysis():
"""Main analysis orchestrator - coordinates the entire trading analysis workflow."""
# Get user selections
selections = get_user_selections()
# Setup configuration and initialize components
config = setup_config(selections)
graph, log_file, report_dir = setup_analysis(selections, config)
# Create display layout and run analysis
layout = create_layout()
with Live(layout, refresh_per_second=4) as live:
# Initialize display with startup messages
initialize_display(layout, selections)
# Run the streaming analysis
trace = run_stream_analysis(graph, selections, layout)
# Process and display final results
finalize_analysis(trace, graph, selections, layout)
@app.command()
def analyze():
run_analysis()
if __name__ == "__main__":
app()