diff --git a/cli/main.py b/cli/main.py index 3f42f2e..64616ee 100644 --- a/cli/main.py +++ b/cli/main.py @@ -1,6 +1,8 @@ from typing import Optional import datetime import typer +from pathlib import Path +from functools import wraps from rich.console import Console from rich.panel import Panel from rich.spinner import Spinner @@ -747,6 +749,53 @@ def run_analysis(): [analyst.value for analyst in selections["analysts"]], config=config, debug=True ) + # Create result directory + results_dir = Path(config["results_dir"]) / 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") 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") 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" + with open(report_dir / file_name, "w") as f: + f.write(content) + 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() diff --git a/tradingagents/default_config.py b/tradingagents/default_config.py index 2cf15b8..089e9c2 100644 --- a/tradingagents/default_config.py +++ b/tradingagents/default_config.py @@ -2,6 +2,7 @@ import os DEFAULT_CONFIG = { "project_dir": os.path.abspath(os.path.join(os.path.dirname(__file__), ".")), + "results_dir": os.getenv("TRADINGAGENTS_RESULTS_DIR", "./results"), "data_dir": "/Users/yluo/Documents/Code/ScAI/FR1-data", "data_cache_dir": os.path.join( os.path.abspath(os.path.join(os.path.dirname(__file__), ".")),