Allow planning agent to direct implementation of tasks.

This commit is contained in:
user 2024-12-21 13:40:17 -05:00
parent 37e36967ee
commit d44d309028
4 changed files with 58 additions and 45 deletions

View File

@ -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.agent_utils import run_research_agent
from ra_aid.prompts import ( from ra_aid.prompts import (
PLANNING_PROMPT, PLANNING_PROMPT,
IMPLEMENTATION_PROMPT,
CHAT_PROMPT, CHAT_PROMPT,
EXPERT_PROMPT_SECTION_PLANNING, EXPERT_PROMPT_SECTION_PLANNING,
EXPERT_PROMPT_SECTION_IMPLEMENTATION,
HUMAN_PROMPT_SECTION_PLANNING, HUMAN_PROMPT_SECTION_PLANNING,
HUMAN_PROMPT_SECTION_IMPLEMENTATION
) )
from ra_aid.llm import initialize_llm 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 _global_memory.get('implementation_requested', False)
return 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(): def main():
"""Main entry point for the ra-aid command line tool.""" """Main entry point for the ra-aid command line tool."""
try: try:
@ -261,16 +229,6 @@ def main():
# Run planning agent # Run planning agent
run_agent_with_retry(planning_agent, planning_prompt, config) 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: except KeyboardInterrupt:
console.print("\n[red]Operation cancelled by user[/red]") console.print("\n[red]Operation cancelled by user[/red]")
sys.exit(1) sys.exit(1)

View File

@ -238,8 +238,9 @@ Guidelines:
Use emit_plan to store the high-level implementation plan. Use emit_plan to store the high-level implementation plan.
For each sub-task, use emit_task to store a step-by-step description. 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. 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} {expert_section}
{human_section} {human_section}

View File

@ -9,7 +9,7 @@ from ra_aid.tools import (
swap_task_order, monorepo_detected, existing_project_detected, ui_detected swap_task_order, monorepo_detected, existing_project_detected, ui_detected
) )
from ra_aid.tools.memory import one_shot_completed 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 # Read-only tools that don't modify system state
def get_read_only_tools(human_interaction: bool = False) -> list: 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, delete_tasks,
emit_plan, emit_plan,
emit_task, emit_task,
swap_task_order swap_task_order,
request_task_implementation
] ]
tools.extend(planning_tools) tools.extend(planning_tools)

View File

@ -6,6 +6,7 @@ from rich.console import Console
from ra_aid.tools.memory import _global_memory from ra_aid.tools.memory import _global_memory
from .memory import get_memory_value, get_related_files from .memory import get_memory_value, get_related_files
from ..llm import initialize_llm from ..llm import initialize_llm
from ..console import print_task_header
console = Console() console = Console()
@ -59,3 +60,55 @@ def request_research(query: str) -> Dict[str, Any]:
"success": success, "success": success,
"reason": reason "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
}