From 417b09712ce1ea316e22f054f257888a985fa5f1 Mon Sep 17 00:00:00 2001 From: ZeroAct Date: Thu, 12 Jun 2025 13:53:28 +0900 Subject: [PATCH] refactor --- cli/main.py | 63 ++++++++++++++++++++++++++++++++++++++--------------- 1 file changed, 45 insertions(+), 18 deletions(-) diff --git a/cli/main.py b/cli/main.py index 462ad49..92cfc92 100644 --- a/cli/main.py +++ b/cli/main.py @@ -2,6 +2,7 @@ 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 @@ -702,8 +703,51 @@ def run_analysis(): ) # Create result directory - results_dir = Path(config["results_dir"]) / selections["ticker"] / datetime.datetime.now().strftime("%Y-%m-%d") + 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() @@ -999,23 +1043,6 @@ def run_analysis(): if section in final_state: message_buffer.update_report_section(section, final_state[section]) - # Save results to file - report_dir = results_dir / "reports" - report_dir.mkdir(parents=True, exist_ok=True) - - for section, content in message_buffer.report_sections.items(): - if content: - with open(report_dir / f"{section}.md", "w") as f: - f.write(content) - - for (timestamp, msg_type, content) in message_buffer.messages: - with open(results_dir / "messages.log", "a") as f: - f.write(f"{timestamp} [{msg_type}]: {content}\n") - - for (timestamp, tool_name, args) in message_buffer.tool_calls: - with open(results_dir / "tool_calls.log", "a") as f: - f.write(f"{timestamp} [Tool: {tool_name}]: {args}\n") - # Display the complete final report display_complete_report(final_state)