diff --git a/ra_aid/__main__.py b/ra_aid/__main__.py index a7dfce5..e735564 100644 --- a/ra_aid/__main__.py +++ b/ra_aid/__main__.py @@ -98,10 +98,10 @@ def is_stage_requested(stage: str) -> bool: def run_implementation_stage(base_task, tasks, plan, related_files): """Run implementation stage with a distinct agent for each task.""" if not is_stage_requested('implementation'): - print_stage_header("SKIPPING IMPLEMENTATION STAGE (not requested)") + print_stage_header("Implementation Stage Skipped") return - print_stage_header("IMPLEMENTATION STAGE") + print_stage_header("Implementation Stage") # Get tasks directly from memory instead of using get_memory_value which joins with newlines task_list = _global_memory['tasks'] @@ -150,7 +150,7 @@ def summarize_research_findings(base_task: str, config: dict) -> None: base_task: The original user query config: Configuration dictionary for the agent """ - print_stage_header("RESEARCH SUMMARY") + print_stage_header("Research Summary") # Create dedicated memory for research summarization summary_memory = MemorySaver() @@ -183,7 +183,7 @@ def run_research_subtasks(base_task: str, config: dict): if not subtasks: return - print_stage_header("RESEARCH SUBTASKS") + print_stage_header("Research Subtasks") # Create tools for subtask agents (excluding spawn_research_subtask and request_implementation) subtask_tools = [ @@ -255,7 +255,7 @@ def main(): _global_memory['config'] = config # Run research stage - print_stage_header("RESEARCH STAGE") + print_stage_header("Research Stage") research_prompt = f"""User query: {base_task} --keep it simple {RESEARCH_PROMPT} @@ -275,7 +275,7 @@ Be very thorough in your research and emit lots of snippets, key facts. If you t print(f"Encountered Anthropic Internal Server Error: {e}. Retrying...") continue except TaskCompletedException as e: - print_stage_header("TASK COMPLETED") + print_stage_header("Task Completed") raise # Re-raise to be caught by outer handler # Run any research subtasks @@ -286,7 +286,7 @@ Be very thorough in your research and emit lots of snippets, key facts. If you t summarize_research_findings(base_task, config) else: # Only proceed with planning and implementation if not an informational query - print_stage_header("PLANNING STAGE") + print_stage_header("Planning Stage") planning_prompt = PLANNING_PROMPT.format( research_notes=get_memory_value('research_notes'), key_facts=get_memory_value('key_facts'), diff --git a/ra_aid/console/formatting.py b/ra_aid/console/formatting.py index de47070..6678ec2 100644 --- a/ra_aid/console/formatting.py +++ b/ra_aid/console/formatting.py @@ -1,16 +1,39 @@ from rich.console import Console +from rich.rule import Rule from rich.panel import Panel from rich.markdown import Markdown console = Console() def print_stage_header(stage: str) -> None: - """Print a stage header with green styling and rocket emoji. Content is rendered as Markdown. + """Print a stage header with stage-specific styling and icons. Args: - stage: The stage name to print (supports Markdown formatting) + stage: The stage name to print (automatically formatted to Title Case) """ - console.print(Panel(Markdown(stage), title="๐Ÿš€ Stage", style="green bold")) + # Define stage icons mapping + icons = { + 'research stage': '๐Ÿ”Ž', + 'planning stage': '๐Ÿ“', + 'implementation stage': '๐Ÿ› ๏ธ', + 'task completed': 'โœ…', + 'debug stage': '๐Ÿ›', + 'testing stage': '๐Ÿงช', + 'research summary': '๐Ÿ“‹', + 'research subtasks': '๐Ÿ“š', + 'skipping implementation stage': 'โญ๏ธ' + } + + # Format stage name to Title Case and normalize for mapping lookup + stage_title = stage.title() + stage_key = stage.lower() + + # Get appropriate icon with fallback + icon = icons.get(stage_key, '๐Ÿš€') + + # Create styled rule with icon + rule_content = f"{icon} {stage_title} {icon}" + console.print(Rule(rule_content, style="green bold")) def print_task_header(task: str) -> None: """Print a task header with yellow styling and wrench emoji. Content is rendered as Markdown.