mirror of
https://github.com/TauricResearch/TradingAgents.git
synced 2026-09-27 06:56:39 +03:00
refactor(cli): split the interactive choices and the run loop out of cli/main.py
- selections.py asks what to run; run.py builds the graph, streams it to the live view and saves the report - main.py keeps the Typer app and its two commands
This commit is contained in:
Yijia-Xiao
+2
-712
@@ -1,60 +1,12 @@
|
||||
import datetime
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from functools import wraps
|
||||
from pathlib import Path
|
||||
|
||||
import typer
|
||||
from rich.align import Align
|
||||
from rich.live import Live
|
||||
from rich.panel import Panel
|
||||
|
||||
from cli.announcements import display_announcements, fetch_announcements
|
||||
from cli.display import (
|
||||
ANALYST_ORDER,
|
||||
AnalystWallTimeTracker,
|
||||
classify_message_type,
|
||||
console,
|
||||
create_layout,
|
||||
display_complete_report,
|
||||
message_buffer,
|
||||
update_analyst_statuses,
|
||||
update_display,
|
||||
update_research_team_status,
|
||||
)
|
||||
from cli.prefs import load_last_run, sanitize, save_last_run
|
||||
from cli.prompts import (
|
||||
ask_anthropic_effort,
|
||||
ask_gemini_thinking_config,
|
||||
ask_glm_region,
|
||||
ask_minimax_region,
|
||||
ask_openai_reasoning_effort,
|
||||
ask_output_language,
|
||||
ask_qwen_region,
|
||||
confirm_ollama_endpoint,
|
||||
detect_asset_type,
|
||||
ensure_api_key,
|
||||
get_ticker,
|
||||
prompt_openai_compatible_url,
|
||||
resolve_backend_url,
|
||||
select_analysts,
|
||||
select_deep_thinking_agent,
|
||||
select_llm_provider,
|
||||
select_research_depth,
|
||||
select_shallow_thinking_agent,
|
||||
)
|
||||
from cli.stats_handler import StatsCallbackHandler
|
||||
from tradingagents.agents.rating import is_review
|
||||
from cli.display import console
|
||||
from cli.run import run_analysis
|
||||
from tradingagents.backtest import iter_grid, run_backtest, summarize
|
||||
from tradingagents.dataflows.symbols import safe_ticker_component
|
||||
from tradingagents.default_config import DEFAULT_CONFIG
|
||||
from tradingagents.graph.analyst_execution import (
|
||||
build_analyst_execution_plan,
|
||||
)
|
||||
from tradingagents.graph.trading_graph import TradingAgentsGraph
|
||||
from tradingagents.portfolio import load_portfolio
|
||||
from tradingagents.reporting import write_report_tree
|
||||
|
||||
# prompt_toolkit's win32 output module is importable only on Windows (it asserts
|
||||
# the platform at import time), so gate on the platform rather than catching the
|
||||
@@ -75,667 +27,6 @@ app = typer.Typer(
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def get_user_selections():
|
||||
"""Ask for the run's settings, offering the previous run's answers."""
|
||||
selections = _prompt_selections(load_last_run())
|
||||
save_last_run(selections)
|
||||
return selections
|
||||
|
||||
|
||||
def _prompt_selections(prefs):
|
||||
"""Walk the selection steps. ``prefs`` prefills, the environment skips."""
|
||||
# Display ASCII art welcome message
|
||||
with open(Path(__file__).parent / "static" / "welcome.txt", encoding="utf-8") as f:
|
||||
welcome_ascii = f.read()
|
||||
|
||||
# Create welcome box content
|
||||
welcome_content = f"{welcome_ascii}\n"
|
||||
welcome_content += "[bold green]TradingAgents: Multi-Agents LLM Financial Trading Framework - CLI[/bold green]\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\n\n"
|
||||
welcome_content += (
|
||||
"[dim]Built by [Tauric Research](https://github.com/TauricResearch)[/dim]"
|
||||
)
|
||||
|
||||
# Create and center the welcome box
|
||||
welcome_box = Panel(
|
||||
welcome_content,
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
title="Welcome to TradingAgents",
|
||||
subtitle="Multi-Agents LLM Financial Trading Framework",
|
||||
)
|
||||
console.print(Align.center(welcome_box))
|
||||
console.print()
|
||||
console.print() # Add vertical space before announcements
|
||||
|
||||
# Fetch and display announcements (silent on failure)
|
||||
announcements = fetch_announcements()
|
||||
display_announcements(console, announcements)
|
||||
|
||||
# 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))
|
||||
|
||||
def thinking_value_or_prompt(env_var, config_key, label, box_title, box_body, prompt_fn):
|
||||
"""Return the env-configured reasoning/thinking value, or prompt for it.
|
||||
|
||||
When ``env_var`` is set the interactive choice is skipped and the value
|
||||
the env overlay placed on DEFAULT_CONFIG is used — mirroring the
|
||||
env-precedence rule applied to the other selection steps.
|
||||
"""
|
||||
if os.environ.get(env_var):
|
||||
value = DEFAULT_CONFIG[config_key]
|
||||
console.print(f"[green]✓ {label} from environment:[/green] {value}")
|
||||
return value
|
||||
console.print(create_question_box(box_title, box_body))
|
||||
return prompt_fn()
|
||||
|
||||
# Step 1: Ticker symbol
|
||||
console.print(
|
||||
create_question_box(
|
||||
"Step 1: Ticker Symbol",
|
||||
"Enter the ticker, with exchange suffix when needed (e.g. SPY, 0700.HK, BTC-USD)",
|
||||
"SPY",
|
||||
)
|
||||
)
|
||||
selected_ticker = get_ticker()
|
||||
asset_type = detect_asset_type(selected_ticker)
|
||||
# Only announce when it's not the default stock path, to avoid printing
|
||||
# "stock" on every run.
|
||||
if asset_type.value != "stock":
|
||||
console.print(
|
||||
f"[green]Detected asset type:[/green] {asset_type.value}"
|
||||
)
|
||||
|
||||
# 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: Output language (skipped when set via TRADINGAGENTS_OUTPUT_LANGUAGE)
|
||||
if os.environ.get("TRADINGAGENTS_OUTPUT_LANGUAGE"):
|
||||
output_language = DEFAULT_CONFIG["output_language"]
|
||||
console.print(
|
||||
f"[green]✓ Output language from environment:[/green] {output_language}"
|
||||
)
|
||||
else:
|
||||
console.print(
|
||||
create_question_box(
|
||||
"Step 3: Output Language",
|
||||
"Select the language for analyst reports and final decision"
|
||||
)
|
||||
)
|
||||
output_language = ask_output_language(prefs.get("output_language"))
|
||||
|
||||
# Step 4: Select analysts
|
||||
console.print(
|
||||
create_question_box(
|
||||
"Step 4: Analysts Team", "Select your LLM analyst agents for the analysis"
|
||||
)
|
||||
)
|
||||
prefs = sanitize(prefs, asset_type.value)
|
||||
selected_analysts = select_analysts(asset_type, prefs.get("analysts"))
|
||||
console.print(
|
||||
f"[green]Selected analysts:[/green] {', '.join(analyst.value for analyst in selected_analysts)}"
|
||||
)
|
||||
|
||||
# Step 5: Research depth (skipped when both round counts are set via env).
|
||||
# Research depth maps to the debate + risk round counts; when both are
|
||||
# supplied through TRADINGAGENTS_MAX_DEBATE_ROUNDS / _MAX_RISK_ROUNDS we keep
|
||||
# the run non-interactive and honor the env values (#977).
|
||||
depth_from_env = bool(os.environ.get("TRADINGAGENTS_MAX_DEBATE_ROUNDS")) and bool(
|
||||
os.environ.get("TRADINGAGENTS_MAX_RISK_ROUNDS")
|
||||
)
|
||||
if depth_from_env:
|
||||
selected_research_depth = DEFAULT_CONFIG["max_debate_rounds"]
|
||||
console.print(
|
||||
f"[green]✓ Research depth from environment:[/green] "
|
||||
f"{DEFAULT_CONFIG['max_debate_rounds']} debate / "
|
||||
f"{DEFAULT_CONFIG['max_risk_discuss_rounds']} risk rounds"
|
||||
)
|
||||
else:
|
||||
console.print(
|
||||
create_question_box(
|
||||
"Step 5: Research Depth", "Select your research depth level"
|
||||
)
|
||||
)
|
||||
selected_research_depth = select_research_depth(prefs.get("research_depth"))
|
||||
|
||||
# Step 6: LLM Provider (skipped when set via TRADINGAGENTS_LLM_PROVIDER).
|
||||
# The backend URL comes from TRADINGAGENTS_LLM_BACKEND_URL when set,
|
||||
# otherwise the provider's default endpoint — the same value the menu
|
||||
# would have picked.
|
||||
provider_from_env = bool(os.environ.get("TRADINGAGENTS_LLM_PROVIDER"))
|
||||
if provider_from_env:
|
||||
selected_llm_provider = DEFAULT_CONFIG["llm_provider"].lower()
|
||||
backend_url = resolve_backend_url(
|
||||
selected_llm_provider, env_url=DEFAULT_CONFIG["backend_url"]
|
||||
)
|
||||
console.print(f"[green]✓ LLM provider from environment:[/green] {selected_llm_provider}")
|
||||
console.print(f"[green]✓ Backend URL:[/green] {backend_url}")
|
||||
# Still confirm/persist the API key so the run doesn't fail later.
|
||||
ensure_api_key(selected_llm_provider)
|
||||
else:
|
||||
console.print(
|
||||
create_question_box(
|
||||
"Step 6: LLM Provider", "Select your LLM provider"
|
||||
)
|
||||
)
|
||||
selected_llm_provider, backend_url = select_llm_provider(prefs.get("llm_provider"))
|
||||
|
||||
# Providers with regional endpoints prompt for the region as a secondary
|
||||
# step so the main dropdown stays clean (mainland China and international
|
||||
# accounts cannot share API keys).
|
||||
if selected_llm_provider == "qwen":
|
||||
selected_llm_provider, backend_url = ask_qwen_region()
|
||||
elif selected_llm_provider == "minimax":
|
||||
selected_llm_provider, backend_url = ask_minimax_region()
|
||||
elif selected_llm_provider == "glm":
|
||||
selected_llm_provider, backend_url = ask_glm_region()
|
||||
|
||||
# Honor an explicit env backend URL even when the provider was chosen
|
||||
# interactively, so it isn't overwritten by the menu default (#978).
|
||||
backend_url = resolve_backend_url(
|
||||
selected_llm_provider, backend_url, env_url=DEFAULT_CONFIG["backend_url"]
|
||||
)
|
||||
|
||||
# The generic OpenAI-compatible endpoint has no default; ask for it if
|
||||
# neither the menu nor the environment supplied one.
|
||||
if selected_llm_provider == "openai_compatible" and not backend_url:
|
||||
remembered_url = (prefs.get("backend_url")
|
||||
if prefs.get("llm_provider") == selected_llm_provider else None)
|
||||
backend_url = prompt_openai_compatible_url(remembered_url)
|
||||
|
||||
# For Ollama, surface the resolved endpoint (OLLAMA_BASE_URL vs default)
|
||||
# before model selection so it's obvious where we're connecting.
|
||||
if selected_llm_provider == "ollama":
|
||||
confirm_ollama_endpoint(backend_url)
|
||||
|
||||
# Confirm the provider's API key is present; prompt the user to paste
|
||||
# one and persist it to .env if it's missing, so the analysis run
|
||||
# doesn't fail later at the first API call.
|
||||
ensure_api_key(selected_llm_provider)
|
||||
|
||||
# Step 7: Thinking agents (skipped when either model is set via environment)
|
||||
if os.environ.get("TRADINGAGENTS_QUICK_THINK_LLM") or os.environ.get("TRADINGAGENTS_DEEP_THINK_LLM"):
|
||||
selected_shallow_thinker = DEFAULT_CONFIG["quick_think_llm"]
|
||||
selected_deep_thinker = DEFAULT_CONFIG["deep_think_llm"]
|
||||
console.print(
|
||||
f"[green]✓ Thinking agents from environment:[/green] "
|
||||
f"quick={selected_shallow_thinker}, deep={selected_deep_thinker}"
|
||||
)
|
||||
else:
|
||||
console.print(
|
||||
create_question_box(
|
||||
"Step 7: Thinking Agents", "Select your thinking agents for analysis"
|
||||
)
|
||||
)
|
||||
remembered = prefs if prefs.get("llm_provider") == selected_llm_provider else {}
|
||||
selected_shallow_thinker = select_shallow_thinking_agent(
|
||||
selected_llm_provider, remembered.get("quick_think_llm")
|
||||
)
|
||||
selected_deep_thinker = select_deep_thinking_agent(
|
||||
selected_llm_provider, remembered.get("deep_think_llm")
|
||||
)
|
||||
|
||||
# Step 8: Provider-specific reasoning/thinking configuration. Each knob is
|
||||
# settable via its TRADINGAGENTS_* env var; when that var is set (or the
|
||||
# provider itself came from env) the prompt is skipped and the configured
|
||||
# value is used — same env-precedence rule as the steps above. None = each
|
||||
# provider's own default.
|
||||
thinking_level = None
|
||||
reasoning_effort = None
|
||||
anthropic_effort = None
|
||||
|
||||
provider_lower = selected_llm_provider.lower()
|
||||
if provider_from_env:
|
||||
thinking_level = DEFAULT_CONFIG["google_thinking_level"]
|
||||
reasoning_effort = DEFAULT_CONFIG["openai_reasoning_effort"]
|
||||
anthropic_effort = DEFAULT_CONFIG["anthropic_effort"]
|
||||
elif provider_lower == "google":
|
||||
thinking_level = thinking_value_or_prompt(
|
||||
"TRADINGAGENTS_GOOGLE_THINKING_LEVEL", "google_thinking_level",
|
||||
"Gemini thinking mode", "Step 8: Thinking Mode",
|
||||
"Configure Gemini thinking mode", ask_gemini_thinking_config,
|
||||
)
|
||||
elif provider_lower == "openai":
|
||||
reasoning_effort = thinking_value_or_prompt(
|
||||
"TRADINGAGENTS_OPENAI_REASONING_EFFORT", "openai_reasoning_effort",
|
||||
"Reasoning effort", "Step 8: Reasoning Effort",
|
||||
"Configure OpenAI reasoning effort level", ask_openai_reasoning_effort,
|
||||
)
|
||||
elif provider_lower == "anthropic":
|
||||
anthropic_effort = thinking_value_or_prompt(
|
||||
"TRADINGAGENTS_ANTHROPIC_EFFORT", "anthropic_effort",
|
||||
"Claude effort", "Step 8: Effort Level",
|
||||
"Configure Claude effort level", ask_anthropic_effort,
|
||||
)
|
||||
|
||||
return {
|
||||
"ticker": selected_ticker,
|
||||
"asset_type": asset_type.value,
|
||||
"analysis_date": analysis_date,
|
||||
"analysts": selected_analysts,
|
||||
"research_depth": selected_research_depth,
|
||||
"llm_provider": selected_llm_provider.lower(),
|
||||
"backend_url": backend_url,
|
||||
"quick_think_llm": selected_shallow_thinker,
|
||||
"deep_think_llm": selected_deep_thinker,
|
||||
"google_thinking_level": thinking_level,
|
||||
"openai_reasoning_effort": reasoning_effort,
|
||||
"anthropic_effort": anthropic_effort,
|
||||
"output_language": output_language,
|
||||
}
|
||||
|
||||
|
||||
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 _run_directory(config: dict, ticker: str, trade_date: str) -> Path:
|
||||
"""Where this run writes, with the ticker validated as a path component.
|
||||
|
||||
Every other path that interpolates a ticker checks it first; a value of
|
||||
".." here would place the run outside the results directory.
|
||||
"""
|
||||
return Path(config["results_dir"]) / safe_ticker_component(ticker) / trade_date
|
||||
|
||||
|
||||
def _announce_checkpoint_state(graph, ticker: str, trade_date: str) -> None:
|
||||
"""Say whether this run resumed a saved one, where the user can see it.
|
||||
|
||||
The graph logs this, but nothing in the CLI configures logging and the live
|
||||
view owns the screen, so a resume was invisible.
|
||||
"""
|
||||
if getattr(graph, "_resuming", False):
|
||||
message_buffer.add_message(
|
||||
"System", f"Resuming the saved run for {ticker} on {trade_date}"
|
||||
)
|
||||
else:
|
||||
message_buffer.add_message("System", f"Starting fresh for {ticker} on {trade_date}")
|
||||
|
||||
|
||||
def _build_run_config(selections: dict, checkpoint: bool | None) -> dict:
|
||||
"""Assemble the run config from interactive selections, honoring env precedence.
|
||||
|
||||
Round counts and checkpoint follow "explicit env/flag wins": an env-applied
|
||||
value on DEFAULT_CONFIG is preserved unless the user overrode it on the CLI.
|
||||
"""
|
||||
config = DEFAULT_CONFIG.copy()
|
||||
# Research depth sets both round counts, but an explicit env override
|
||||
# (TRADINGAGENTS_MAX_DEBATE_ROUNDS / _MAX_RISK_ROUNDS) wins over the
|
||||
# interactive selection — leave the env-applied value in place (#977).
|
||||
for env_var, key in (("TRADINGAGENTS_MAX_DEBATE_ROUNDS", "max_debate_rounds"),
|
||||
("TRADINGAGENTS_MAX_RISK_ROUNDS", "max_risk_discuss_rounds")):
|
||||
if os.environ.get(env_var):
|
||||
# The depth prompt still appeared (it is skipped only when both are
|
||||
# set), so say which half of the answer the environment overrode.
|
||||
console.print(
|
||||
f"[green]✓ {key} from environment:[/green] {config[key]} "
|
||||
f"(set by {env_var}, so the research depth you chose does not apply to it)"
|
||||
)
|
||||
else:
|
||||
config[key] = selections["research_depth"]
|
||||
config["quick_think_llm"] = selections["quick_think_llm"]
|
||||
config["deep_think_llm"] = selections["deep_think_llm"]
|
||||
config["backend_url"] = selections["backend_url"]
|
||||
config["llm_provider"] = selections["llm_provider"].lower()
|
||||
# Provider-specific thinking configuration
|
||||
config["google_thinking_level"] = selections.get("google_thinking_level")
|
||||
config["openai_reasoning_effort"] = selections.get("openai_reasoning_effort")
|
||||
config["anthropic_effort"] = selections.get("anthropic_effort")
|
||||
config["output_language"] = selections.get("output_language", "English")
|
||||
# --checkpoint/--no-checkpoint overrides only when explicitly given; omitting
|
||||
# the flag preserves TRADINGAGENTS_CHECKPOINT_ENABLED / the default (#976).
|
||||
if checkpoint is not None:
|
||||
config["checkpoint_enabled"] = checkpoint
|
||||
return config
|
||||
|
||||
|
||||
def run_analysis(checkpoint: bool | None = None, portfolio=None):
|
||||
# First get all user selections
|
||||
selections = get_user_selections()
|
||||
|
||||
config = _build_run_config(selections, checkpoint)
|
||||
|
||||
# Create stats callback handler for tracking LLM/tool calls
|
||||
stats_handler = StatsCallbackHandler()
|
||||
|
||||
# Normalize analyst selection to predefined order (selection is a 'set', order is fixed)
|
||||
selected_set = {analyst.value for analyst in selections["analysts"]}
|
||||
selected_analyst_keys = [a for a in ANALYST_ORDER if a in selected_set]
|
||||
analyst_execution_plan = build_analyst_execution_plan(selected_analyst_keys)
|
||||
analyst_wall_time_tracker = AnalystWallTimeTracker(analyst_execution_plan)
|
||||
|
||||
# Initialize the graph with callbacks bound to LLMs
|
||||
graph = TradingAgentsGraph(
|
||||
selected_analyst_keys,
|
||||
config=config,
|
||||
debug=True,
|
||||
callbacks=[stats_handler],
|
||||
)
|
||||
|
||||
# Initialize message buffer with selected analysts
|
||||
message_buffer.init_for_analysis(selected_analyst_keys)
|
||||
|
||||
# Track start time for elapsed display
|
||||
start_time = time.time()
|
||||
|
||||
# Create result directory
|
||||
results_dir = _run_directory(config, selections["ticker"], selections["analysis_date"])
|
||||
results_dir.mkdir(parents=True, exist_ok=True)
|
||||
report_dir = results_dir / "reports"
|
||||
report_dir.mkdir(parents=True, exist_ok=True)
|
||||
log_file = results_dir / "message_tool.log"
|
||||
log_file.touch(exist_ok=True)
|
||||
|
||||
def save_message_decorator(obj, func_name):
|
||||
func = getattr(obj, func_name)
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
func(*args, **kwargs)
|
||||
timestamp, message_type, content = obj.messages[-1]
|
||||
content = content.replace("\n", " ") # Replace newlines with spaces
|
||||
with open(log_file, "a", encoding="utf-8") as f:
|
||||
f.write(f"{timestamp} [{message_type}] {content}\n")
|
||||
return wrapper
|
||||
|
||||
def save_tool_call_decorator(obj, func_name):
|
||||
func = getattr(obj, func_name)
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
func(*args, **kwargs)
|
||||
timestamp, tool_name, args = obj.tool_calls[-1]
|
||||
args_str = ", ".join(f"{k}={v}" for k, v in args.items())
|
||||
with open(log_file, "a", encoding="utf-8") as f:
|
||||
f.write(f"{timestamp} [Tool Call] {tool_name}({args_str})\n")
|
||||
return wrapper
|
||||
|
||||
def save_report_section_decorator(obj, func_name):
|
||||
func = getattr(obj, func_name)
|
||||
@wraps(func)
|
||||
def wrapper(section_name, content):
|
||||
func(section_name, content)
|
||||
if section_name in obj.report_sections and obj.report_sections[section_name] is not None:
|
||||
content = obj.report_sections[section_name]
|
||||
if content:
|
||||
file_name = f"{section_name}.md"
|
||||
text = "\n".join(str(item) for item in content) if isinstance(content, list) else content
|
||||
with open(report_dir / file_name, "w", encoding="utf-8") as f:
|
||||
f.write(text)
|
||||
return wrapper
|
||||
|
||||
message_buffer.add_message = save_message_decorator(message_buffer, "add_message")
|
||||
message_buffer.add_tool_call = save_tool_call_decorator(message_buffer, "add_tool_call")
|
||||
message_buffer.update_report_section = save_report_section_decorator(message_buffer, "update_report_section")
|
||||
|
||||
# Now start the display layout
|
||||
layout = create_layout()
|
||||
|
||||
# The alternate screen keeps a layout taller than the window from redrawing
|
||||
# by scrolling; the final report prints after this block, on the normal screen.
|
||||
with Live(layout, refresh_per_second=4, screen=True):
|
||||
# Initial display
|
||||
update_display(layout, stats_handler=stats_handler, start_time=start_time)
|
||||
|
||||
# Add initial messages
|
||||
message_buffer.add_message("System", f"Selected ticker: {selections['ticker']}")
|
||||
if selections["asset_type"] != "stock":
|
||||
message_buffer.add_message("System", f"Detected asset type: {selections['asset_type']}")
|
||||
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, stats_handler=stats_handler, start_time=start_time)
|
||||
|
||||
# Update agent status to in_progress for the first analyst
|
||||
first_analyst = analyst_execution_plan.specs[0].agent_node
|
||||
message_buffer.update_agent_status(first_analyst, "in_progress")
|
||||
analyst_wall_time_tracker.mark_started(selected_analyst_keys[0])
|
||||
update_display(layout, stats_handler=stats_handler, start_time=start_time)
|
||||
|
||||
# Create spinner text
|
||||
spinner_text = (
|
||||
f"Analyzing {selections['ticker']} on {selections['analysis_date']}..."
|
||||
)
|
||||
update_display(layout, spinner_text, stats_handler=stats_handler, start_time=start_time)
|
||||
|
||||
# The same initial state propagate() builds: settled decision log, past
|
||||
# context and resolved instrument identity.
|
||||
init_agent_state = graph.create_run_state(
|
||||
selections["ticker"], selections["analysis_date"], selections["asset_type"], portfolio
|
||||
)
|
||||
# Pass callbacks to graph config for tool execution tracking
|
||||
# (LLM tracking is handled separately via LLM constructor)
|
||||
args = graph.propagator.get_graph_args(callbacks=[stats_handler])
|
||||
|
||||
# Recompile with a checkpointer and inject the thread_id so --checkpoint
|
||||
# actually saves and resumes on the CLI path (#1249); a no-op when
|
||||
# checkpointing is disabled. Torn down in the finally below.
|
||||
checkpoint_tid = graph.begin_checkpoint(
|
||||
selections["ticker"], selections["analysis_date"], selections["asset_type"], portfolio
|
||||
)
|
||||
if checkpoint_tid is not None:
|
||||
args.setdefault("config", {}).setdefault("configurable", {})["thread_id"] = checkpoint_tid
|
||||
_announce_checkpoint_state(graph, selections["ticker"], selections["analysis_date"])
|
||||
|
||||
# Stream the analysis. On resume, feed None so LangGraph continues the
|
||||
# interrupted run instead of re-appending the initial state (#1249); the
|
||||
# try/finally tears the checkpointer down even if the stream raises.
|
||||
trace = []
|
||||
try:
|
||||
for chunk in graph.graph.stream(graph.checkpoint_input(init_agent_state), **args):
|
||||
# Process all messages in chunk, deduplicating by message ID
|
||||
for message in chunk.get("messages", []):
|
||||
msg_id = getattr(message, "id", None)
|
||||
if msg_id is not None:
|
||||
if msg_id in message_buffer._processed_message_ids:
|
||||
continue
|
||||
message_buffer._processed_message_ids.add(msg_id)
|
||||
|
||||
msg_type, content = classify_message_type(message)
|
||||
if content and content.strip():
|
||||
message_buffer.add_message(msg_type, content)
|
||||
|
||||
if hasattr(message, "tool_calls") and message.tool_calls:
|
||||
for tool_call in message.tool_calls:
|
||||
if isinstance(tool_call, dict):
|
||||
message_buffer.add_tool_call(tool_call["name"], tool_call["args"])
|
||||
else:
|
||||
message_buffer.add_tool_call(tool_call.name, tool_call.args)
|
||||
|
||||
# Update analyst statuses based on report state (runs on every chunk)
|
||||
update_analyst_statuses(
|
||||
message_buffer,
|
||||
chunk,
|
||||
wall_time_tracker=analyst_wall_time_tracker,
|
||||
)
|
||||
|
||||
# Research Team - Handle Investment Debate State
|
||||
if chunk.get("investment_debate_state"):
|
||||
debate_state = chunk["investment_debate_state"]
|
||||
bull_hist = debate_state.get("bull_history", "").strip()
|
||||
bear_hist = debate_state.get("bear_history", "").strip()
|
||||
judge = debate_state.get("judge_decision", "").strip()
|
||||
|
||||
# Only update status when there's actual content
|
||||
if bull_hist or bear_hist:
|
||||
update_research_team_status("in_progress")
|
||||
if bull_hist:
|
||||
message_buffer.update_report_section(
|
||||
"investment_plan", f"### Bull Researcher Analysis\n{bull_hist}"
|
||||
)
|
||||
if bear_hist:
|
||||
message_buffer.update_report_section(
|
||||
"investment_plan", f"### Bear Researcher Analysis\n{bear_hist}"
|
||||
)
|
||||
if judge:
|
||||
message_buffer.update_report_section(
|
||||
"investment_plan", f"### Research Manager Decision\n{judge}"
|
||||
)
|
||||
update_research_team_status("completed")
|
||||
message_buffer.update_agent_status("Trader", "in_progress")
|
||||
|
||||
# Trading Team
|
||||
if chunk.get("trader_investment_plan"):
|
||||
message_buffer.update_report_section(
|
||||
"trader_investment_plan", chunk["trader_investment_plan"]
|
||||
)
|
||||
if message_buffer.agent_status.get("Trader") != "completed":
|
||||
message_buffer.update_agent_status("Trader", "completed")
|
||||
message_buffer.update_agent_status("Aggressive Analyst", "in_progress")
|
||||
|
||||
# Risk Management Team - Handle Risk Debate State
|
||||
if chunk.get("risk_debate_state"):
|
||||
risk_state = chunk["risk_debate_state"]
|
||||
agg_hist = risk_state.get("aggressive_history", "").strip()
|
||||
con_hist = risk_state.get("conservative_history", "").strip()
|
||||
neu_hist = risk_state.get("neutral_history", "").strip()
|
||||
judge = risk_state.get("judge_decision", "").strip()
|
||||
|
||||
if agg_hist:
|
||||
if message_buffer.agent_status.get("Aggressive Analyst") != "completed":
|
||||
message_buffer.update_agent_status("Aggressive Analyst", "in_progress")
|
||||
message_buffer.update_report_section(
|
||||
"final_trade_decision", f"### Aggressive Analyst Analysis\n{agg_hist}"
|
||||
)
|
||||
if con_hist:
|
||||
if message_buffer.agent_status.get("Conservative Analyst") != "completed":
|
||||
message_buffer.update_agent_status("Conservative Analyst", "in_progress")
|
||||
message_buffer.update_report_section(
|
||||
"final_trade_decision", f"### Conservative Analyst Analysis\n{con_hist}"
|
||||
)
|
||||
if neu_hist:
|
||||
if message_buffer.agent_status.get("Neutral Analyst") != "completed":
|
||||
message_buffer.update_agent_status("Neutral Analyst", "in_progress")
|
||||
message_buffer.update_report_section(
|
||||
"final_trade_decision", f"### Neutral Analyst Analysis\n{neu_hist}"
|
||||
)
|
||||
if judge and message_buffer.agent_status.get("Portfolio Manager") != "completed":
|
||||
message_buffer.update_agent_status("Portfolio Manager", "in_progress")
|
||||
message_buffer.update_report_section(
|
||||
"final_trade_decision", f"### Portfolio Manager Decision\n{judge}"
|
||||
)
|
||||
message_buffer.update_agent_status("Aggressive Analyst", "completed")
|
||||
message_buffer.update_agent_status("Conservative Analyst", "completed")
|
||||
message_buffer.update_agent_status("Neutral Analyst", "completed")
|
||||
message_buffer.update_agent_status("Portfolio Manager", "completed")
|
||||
|
||||
# Update the display
|
||||
update_display(layout, stats_handler=stats_handler, start_time=start_time)
|
||||
|
||||
trace.append(chunk)
|
||||
|
||||
# Streamed chunks are per-node deltas, not full state. Merge them
|
||||
# so every report field populated across the run is present.
|
||||
final_state = {}
|
||||
for chunk in trace:
|
||||
final_state.update(chunk)
|
||||
|
||||
# Clean run: log the decision, then drop this run's checkpoint so a
|
||||
# later run starts fresh. A mid-stream failure skips both, keeping
|
||||
# the checkpoint for resume.
|
||||
graph.record_decision(selections["ticker"], selections["analysis_date"], final_state)
|
||||
graph.clear_checkpoint_on_success(
|
||||
selections["ticker"], selections["analysis_date"], selections["asset_type"], portfolio
|
||||
)
|
||||
finally:
|
||||
# Always restore the plain uncheckpointed graph, even on failure.
|
||||
graph.end_checkpoint()
|
||||
|
||||
# Update all agent statuses to completed
|
||||
for agent in message_buffer.agent_status:
|
||||
message_buffer.update_agent_status(agent, "completed")
|
||||
|
||||
message_buffer.add_message(
|
||||
"System", f"Completed analysis for {selections['analysis_date']}"
|
||||
)
|
||||
message_buffer.add_message("System", analyst_wall_time_tracker.format_summary())
|
||||
|
||||
# Update final report sections
|
||||
for section in message_buffer.report_sections:
|
||||
if section in final_state:
|
||||
message_buffer.update_report_section(section, final_state[section])
|
||||
|
||||
update_display(layout, stats_handler=stats_handler, start_time=start_time)
|
||||
|
||||
# Post-analysis prompts (outside Live context for clean interaction)
|
||||
console.print("\n[bold cyan]Analysis Complete![/bold cyan]\n")
|
||||
|
||||
# A decision nobody can read is not a position. Say so here rather than
|
||||
# leaving the run to look like a normal result.
|
||||
if is_review(graph.process_signal(final_state.get("final_trade_decision", ""))):
|
||||
console.print(
|
||||
"[yellow]No rating could be read from the final decision, so this run "
|
||||
"is recorded for review rather than as a position. Re-run, or read the "
|
||||
"decision text below and judge it yourself.[/yellow]\n"
|
||||
)
|
||||
console.print(f"[dim]{analyst_wall_time_tracker.format_summary()}[/dim]")
|
||||
|
||||
# Prompt to save report
|
||||
save_choice = typer.prompt("Save report?", default="Y").strip().upper()
|
||||
if save_choice in ("Y", "YES", ""):
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
# Under results_dir, not the working directory: in Docker the working
|
||||
# directory is inside the container and the report goes with it, while
|
||||
# results_dir is the mounted volume the rest of the run already writes to.
|
||||
default_path = (Path(config["results_dir"]) / "reports"
|
||||
/ f"{safe_ticker_component(selections['ticker'])}_{timestamp}")
|
||||
save_path_str = typer.prompt(
|
||||
"Save path (press Enter for default)",
|
||||
default=str(default_path)
|
||||
).strip()
|
||||
save_path = Path(save_path_str)
|
||||
try:
|
||||
report_file = write_report_tree(final_state, selections["ticker"], save_path)
|
||||
console.print(f"\n[green]✓ Report saved to:[/green] {save_path.resolve()}")
|
||||
console.print(f" [dim]Complete report:[/dim] {report_file.name}")
|
||||
except Exception as e:
|
||||
console.print(f"[red]Error saving report: {e}[/red]")
|
||||
|
||||
# Prompt to display full report
|
||||
display_choice = typer.prompt("\nDisplay full report on screen?", default="Y").strip().upper()
|
||||
if display_choice in ("Y", "YES", ""):
|
||||
display_complete_report(final_state)
|
||||
|
||||
|
||||
@app.callback(invoke_without_command=True)
|
||||
def analyze(
|
||||
ctx: typer.Context,
|
||||
@@ -766,7 +57,6 @@ def analyze(
|
||||
console.print(f"[yellow]Cleared {n} checkpoint(s).[/yellow]")
|
||||
portfolio_context = None
|
||||
if portfolio:
|
||||
from tradingagents.portfolio import load_portfolio
|
||||
try:
|
||||
portfolio_context = load_portfolio(portfolio)
|
||||
except ValueError as exc:
|
||||
|
||||
Reference in New Issue
Block a user