zach
/
RA.Aid
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

stage ux

This commit is contained in:
AI Christianson 2024-12-11 16:29:38 -05:00
parent 4ea696e32a
commit 6586874e63
2 changed files with 33 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
ra_aid/__main__.py
View File

@ -98,10 +98,10 @@ def is_stage_requested(stage: str) -> bool:
def run_implementation_stage(base_task, tasks, plan, related_files): def run_implementation_stage(base_task, tasks, plan, related_files):
"""Run implementation stage with a distinct agent for each task.""" """Run implementation stage with a distinct agent for each task."""
if not is_stage_requested('implementation'): if not is_stage_requested('implementation'):
print_stage_header("SKIPPING IMPLEMENTATION STAGE (not requested)") print_stage_header("Implementation Stage Skipped")
return 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 # Get tasks directly from memory instead of using get_memory_value which joins with newlines
task_list = _global_memory['tasks'] 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 base_task: The original user query
config: Configuration dictionary for the agent config: Configuration dictionary for the agent
""" """
print_stage_header("RESEARCH SUMMARY") print_stage_header("Research Summary")
# Create dedicated memory for research summarization # Create dedicated memory for research summarization
summary_memory = MemorySaver() summary_memory = MemorySaver()
@ -183,7 +183,7 @@ def run_research_subtasks(base_task: str, config: dict):
if not subtasks: if not subtasks:
return return
print_stage_header("RESEARCH SUBTASKS") print_stage_header("Research Subtasks")
# Create tools for subtask agents (excluding spawn_research_subtask and request_implementation) # Create tools for subtask agents (excluding spawn_research_subtask and request_implementation)
subtask_tools = [ subtask_tools = [
@ -255,7 +255,7 @@ def main():
_global_memory['config'] = config _global_memory['config'] = config
# Run research stage # 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 = f"""User query: {base_task} --keep it simple
{RESEARCH_PROMPT} {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...") print(f"Encountered Anthropic Internal Server Error: {e}. Retrying...")
continue continue
except TaskCompletedException as e: except TaskCompletedException as e:
print_stage_header("TASK COMPLETED") print_stage_header("Task Completed")
raise # Re-raise to be caught by outer handler raise # Re-raise to be caught by outer handler
# Run any research subtasks # 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) summarize_research_findings(base_task, config)
else: else:
# Only proceed with planning and implementation if not an informational query # 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( planning_prompt = PLANNING_PROMPT.format(
research_notes=get_memory_value('research_notes'), research_notes=get_memory_value('research_notes'),
key_facts=get_memory_value('key_facts'), key_facts=get_memory_value('key_facts'),

29
ra_aid/console/formatting.py
View File

@ -1,16 +1,39 @@
from rich.console import Console from rich.console import Console
from rich.rule import Rule
from rich.panel import Panel from rich.panel import Panel
from rich.markdown import Markdown from rich.markdown import Markdown
console = Console() console = Console()
def print_stage_header(stage: str) -> None: 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: 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: def print_task_header(task: str) -> None:
"""Print a task header with yellow styling and wrench emoji. Content is rendered as Markdown. """Print a task header with yellow styling and wrench emoji. Content is rendered as Markdown.