From d44d309028dc9420ad66ccc045afecd900db2dc6 Mon Sep 17 00:00:00 2001 From: user Date: Sat, 21 Dec 2024 13:40:17 -0500 Subject: [PATCH] Allow planning agent to direct implementation of tasks. --- ra_aid/__main__.py | 42 --------------------------------- ra_aid/prompts.py | 3 ++- ra_aid/tool_configs.py | 5 ++-- ra_aid/tools/agent.py | 53 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 58 insertions(+), 45 deletions(-) diff --git a/ra_aid/__main__.py b/ra_aid/__main__.py index 60874a9..f20aec3 100644 --- a/ra_aid/__main__.py +++ b/ra_aid/__main__.py @@ -12,12 +12,9 @@ from ra_aid.agent_utils import run_agent_with_retry, run_task_implementation_age from ra_aid.agent_utils import run_research_agent from ra_aid.prompts import ( PLANNING_PROMPT, - IMPLEMENTATION_PROMPT, CHAT_PROMPT, EXPERT_PROMPT_SECTION_PLANNING, - EXPERT_PROMPT_SECTION_IMPLEMENTATION, HUMAN_PROMPT_SECTION_PLANNING, - HUMAN_PROMPT_SECTION_IMPLEMENTATION ) from ra_aid.llm import initialize_llm @@ -125,35 +122,6 @@ def is_stage_requested(stage: str) -> bool: return _global_memory.get('implementation_requested', False) return False -def run_implementation_stage(base_task, tasks, plan, related_files, model, expert_enabled: bool): - """Run implementation stage with a distinct agent for each task.""" - if not is_stage_requested('implementation'): - print_stage_header("Implementation Stage Skipped") - return - - print_stage_header("Implementation Stage") - - # Get tasks directly from memory, maintaining order by ID - task_list = [task for _, task in sorted(_global_memory['tasks'].items())] - - print_task_header(f"Found {len(task_list)} tasks to implement") - - for i, task in enumerate(task_list, 1): - print_task_header(task) - - # Run implementation agent for this task - run_task_implementation_agent( - base_task=base_task, - tasks=task_list, - task=task, - plan=plan, - related_files=related_files, - model=model, - expert_enabled=expert_enabled - ) - - - def main(): """Main entry point for the ra-aid command line tool.""" try: @@ -261,16 +229,6 @@ def main(): # Run planning agent run_agent_with_retry(planning_agent, planning_prompt, config) - # Run implementation stage with task-specific agents - run_implementation_stage( - base_task, - get_memory_value('tasks'), - get_memory_value('plan'), - get_related_files(), - model, - expert_enabled=expert_enabled - ) - except KeyboardInterrupt: console.print("\n[red]Operation cancelled by user[/red]") sys.exit(1) diff --git a/ra_aid/prompts.py b/ra_aid/prompts.py index bab5be5..e8e1623 100644 --- a/ra_aid/prompts.py +++ b/ra_aid/prompts.py @@ -238,8 +238,9 @@ Guidelines: Use emit_plan to store the high-level implementation plan. For each sub-task, use emit_task to store a step-by-step description. The description should be only as detailed as warranted by the complexity of the request. + You may use delete_tasks or swap_task_order to adjust the task list/order as you plan. - Do not implement anything yet. + Once you are absolutely sure you are completed planning, you may begin to call request_task_implementation one-by-one for each task to implement the plan. {expert_section} {human_section} diff --git a/ra_aid/tool_configs.py b/ra_aid/tool_configs.py index 17f9641..17c4c9f 100644 --- a/ra_aid/tool_configs.py +++ b/ra_aid/tool_configs.py @@ -9,7 +9,7 @@ from ra_aid.tools import ( swap_task_order, monorepo_detected, existing_project_detected, ui_detected ) from ra_aid.tools.memory import one_shot_completed -from ra_aid.tools.agent import request_research +from ra_aid.tools.agent import request_research, request_task_implementation # Read-only tools that don't modify system state def get_read_only_tools(human_interaction: bool = False) -> list: @@ -76,7 +76,8 @@ def get_planning_tools(expert_enabled: bool = True) -> list: delete_tasks, emit_plan, emit_task, - swap_task_order + swap_task_order, + request_task_implementation ] tools.extend(planning_tools) diff --git a/ra_aid/tools/agent.py b/ra_aid/tools/agent.py index d8460b7..97087f8 100644 --- a/ra_aid/tools/agent.py +++ b/ra_aid/tools/agent.py @@ -6,6 +6,7 @@ from rich.console import Console from ra_aid.tools.memory import _global_memory from .memory import get_memory_value, get_related_files from ..llm import initialize_llm +from ..console import print_task_header console = Console() @@ -59,3 +60,55 @@ def request_research(query: str) -> Dict[str, Any]: "success": success, "reason": reason } + +@tool("request_task_implementation") +def request_task_implementation(task_spec: str) -> Dict[str, Any]: + """Spawn an implementation agent to execute the given task. + + Args: + task_spec: The full task specification + """ + # Initialize model from config + config = _global_memory.get('config', {}) + model = initialize_llm(config.get('provider', 'anthropic'), config.get('model', 'claude-3-5-sonnet-20241022')) + + # Get required parameters + tasks = [_global_memory['tasks'][task_id] for task_id in sorted(_global_memory['tasks'])] + plan = _global_memory.get('plan', '') + related_files = list(get_related_files()) + + try: + print_task_header(task_spec) + # Run implementation agent + from ..agent_utils import run_task_implementation_agent + result = run_task_implementation_agent( + base_task=_global_memory.get('base_task', ''), + tasks=tasks, + task=task_spec, + plan=plan, + related_files=related_files, + model=model, + expert_enabled=True + ) + + success = True + reason = None + except KeyboardInterrupt: + console.print("\n[yellow]Task implementation interrupted by user[/yellow]") + success = False + reason = "cancelled_by_user" + except Exception as e: + console.print(f"\n[red]Error during task implementation: {str(e)}[/red]") + success = False + reason = f"error: {str(e)}" + + # Get completion message if available + completion_message = _global_memory.get('completion_message', 'Task was completed successfully.' if success else None) + + return { + "facts": get_memory_value("key_facts"), + "files": list(get_related_files()), + "completion_message": completion_message, + "success": success, + "reason": reason + }