Extricate research subtasks.

This commit is contained in:
AI Christianson 2024-12-21 11:27:14 -05:00
parent a9c6ea93b4
commit 2ac796cf2b
6 changed files with 3 additions and 79 deletions

View File

@ -166,43 +166,6 @@ def run_implementation_stage(base_task, tasks, plan, related_files, model, exper
# Run agent for this task # Run agent for this task
run_agent_with_retry(task_agent, task_prompt, {"configurable": {"thread_id": "abc123"}, "recursion_limit": 100}) run_agent_with_retry(task_agent, task_prompt, {"configurable": {"thread_id": "abc123"}, "recursion_limit": 100})
def run_research_subtasks(base_task: str, config: dict, model, expert_enabled: bool):
"""Run research subtasks with separate agents."""
subtasks = _global_memory.get('research_subtasks', [])
if not subtasks:
return
print_stage_header("Research Subtasks")
# Get tools for subtask agents (excluding request_research_subtask and implementation)
research_only = _global_memory.get('config', {}).get('research_only', False)
subtask_tools = [
t for t in get_research_tools(research_only=research_only, expert_enabled=expert_enabled)
if t.name not in ['request_research_subtask']
]
for i, subtask in enumerate(subtasks, 1):
print_task_header(f"Research Subtask {i}/{len(subtasks)}")
# Create fresh memory and agent for each subtask
subtask_memory = MemorySaver()
subtask_agent = create_react_agent(
model,
subtask_tools,
checkpointer=subtask_memory
)
# Run the subtask agent
expert_section = EXPERT_PROMPT_SECTION_RESEARCH if expert_enabled else ""
human_section = HUMAN_PROMPT_SECTION_RESEARCH if config.get('hil', False) else ""
subtask_prompt = f"Base Task: {base_task}\nResearch Subtask: {subtask}\n\n{RESEARCH_PROMPT.format(
base_task=base_task,
research_only_note='',
expert_section=expert_section,
human_section=human_section
)}"
run_agent_with_retry(subtask_agent, subtask_prompt, config)
def main(): def main():
@ -300,9 +263,6 @@ def main():
# Run research agent # Run research agent
run_agent_with_retry(research_agent, research_prompt, config) run_agent_with_retry(research_agent, research_prompt, config)
# Run any research subtasks
run_research_subtasks(base_task, config, model, expert_enabled=expert_enabled)
# Proceed with planning and implementation if not an informational query # Proceed with planning and implementation if not an informational query
if not is_informational_query(): if not is_informational_query():
print_stage_header("Planning Stage") print_stage_header("Planning Stage")

View File

@ -4,7 +4,7 @@ from ra_aid.tools import (
emit_research_notes, emit_plan, emit_related_files, emit_task, emit_research_notes, emit_plan, emit_related_files, emit_task,
emit_expert_context, emit_key_facts, delete_key_facts, emit_expert_context, emit_key_facts, delete_key_facts,
emit_key_snippets, delete_key_snippets, delete_tasks, emit_key_snippets, delete_key_snippets, delete_tasks,
request_research_subtask, request_implementation, read_file_tool, request_implementation, read_file_tool,
fuzzy_find_project_files, ripgrep_search, list_directory_tree, fuzzy_find_project_files, ripgrep_search, list_directory_tree,
swap_task_order, monorepo_detected, existing_project_detected, ui_detected swap_task_order, monorepo_detected, existing_project_detected, ui_detected
) )
@ -38,7 +38,6 @@ MODIFICATION_TOOLS = [run_programming_task]
COMMON_TOOLS = READ_ONLY_TOOLS + [] COMMON_TOOLS = READ_ONLY_TOOLS + []
EXPERT_TOOLS = [emit_expert_context, ask_expert] EXPERT_TOOLS = [emit_expert_context, ask_expert]
RESEARCH_TOOLS = [ RESEARCH_TOOLS = [
request_research_subtask,
emit_research_notes, emit_research_notes,
one_shot_completed, one_shot_completed,
monorepo_detected, monorepo_detected,

View File

@ -11,7 +11,7 @@ from .list_directory import list_directory_tree
from .ripgrep import ripgrep_search from .ripgrep import ripgrep_search
from .memory import ( from .memory import (
delete_tasks, emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_facts, delete_tasks, emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_facts,
request_implementation, delete_key_facts, request_research_subtask, request_implementation, delete_key_facts,
emit_key_snippets, delete_key_snippets, emit_related_files, swap_task_order emit_key_snippets, delete_key_snippets, emit_related_files, swap_task_order
) )
@ -34,7 +34,6 @@ __all__ = [
'run_programming_task', 'run_programming_task',
'run_shell_command', 'run_shell_command',
'write_file_tool', 'write_file_tool',
'request_research_subtask',
'ripgrep_search', 'ripgrep_search',
'file_str_replace', 'file_str_replace',
'delete_tasks', 'delete_tasks',

View File

@ -22,7 +22,6 @@ _global_memory: Dict[str, Union[List[Any], Dict[int, str], Dict[int, SnippetInfo
'task_completed': False, # Flag indicating if task is complete 'task_completed': False, # Flag indicating if task is complete
'completion_message': '', # Message explaining completion 'completion_message': '', # Message explaining completion
'task_id_counter': 0, # Counter for generating unique task IDs 'task_id_counter': 0, # Counter for generating unique task IDs
'research_subtasks': [],
'key_facts': {}, # Dict[int, str] - ID to fact mapping 'key_facts': {}, # Dict[int, str] - ID to fact mapping
'key_fact_id_counter': 0, # Counter for generating unique fact IDs 'key_fact_id_counter': 0, # Counter for generating unique fact IDs
'key_snippets': {}, # Dict[int, SnippetInfo] - ID to snippet mapping 'key_snippets': {}, # Dict[int, SnippetInfo] - ID to snippet mapping
@ -79,22 +78,6 @@ def emit_task(task: str) -> str:
console.print(Panel(Markdown(task), title=f"✅ Task #{task_id}")) console.print(Panel(Markdown(task), title=f"✅ Task #{task_id}"))
return f"Task #{task_id} stored." return f"Task #{task_id} stored."
@tool("request_research_subtask")
def request_research_subtask(subtask: str) -> str:
"""Spawn a research subtask for investigation of a specific topic.
Use this anytime you can to offload your work to specific things that need to be looked into.
Use this only when it's necessary to dig deeper into a specific topic.
Args:
subtask: Detailed description of the research subtask
Returns:
Confirmation message
"""
_global_memory['research_subtasks'].append(subtask)
console.print(Panel(Markdown(subtask), title="🔬 Research Subtask"))
return "Subtask added."
@tool("emit_key_facts") @tool("emit_key_facts")
@ -309,8 +292,6 @@ def one_shot_completed(message: str) -> str:
Original message if task can be completed, or error message if there are Original message if task can be completed, or error message if there are
pending subtasks or implementation requests pending subtasks or implementation requests
""" """
if len(_global_memory['research_subtasks']) > 0:
return "Cannot complete in one shot - research subtasks pending"
if _global_memory.get('implementation_requested', False): if _global_memory.get('implementation_requested', False):
return "Cannot complete in one shot - implementation was requested" return "Cannot complete in one shot - implementation was requested"

View File

@ -40,7 +40,7 @@ def monorepo_detected() -> dict:
"- Find and note existing versioning and release management practices already in place.\n\n" "- Find and note existing versioning and release management practices already in place.\n\n"
"- Pay extra attention to integration nuances such as authentication, authorization, examples of how APIs are called, etc.\n\n" "- Pay extra attention to integration nuances such as authentication, authorization, examples of how APIs are called, etc.\n\n"
"- Find and note specific examples of all of the above.\n\n" "- Find and note specific examples of all of the above.\n\n"
"- Because you are in a monorepo, you will likely need to call request_research_subtask multiple times.\n\n" "- Because you are in a monorepo, you will need to carefully organize your research into focused areas.\n\n"
"Your goal is to enhance the entire codebase without disrupting its well-established, unified structure." "Your goal is to enhance the entire codebase without disrupting its well-established, unified structure."
) )
} }

View File

@ -2,7 +2,6 @@ import pytest
from ra_aid.tools.memory import ( from ra_aid.tools.memory import (
_global_memory, _global_memory,
get_memory_value, get_memory_value,
request_research_subtask,
emit_key_facts, emit_key_facts,
delete_key_facts, delete_key_facts,
emit_key_snippets, emit_key_snippets,
@ -24,7 +23,6 @@ def reset_memory():
_global_memory['research_notes'] = [] _global_memory['research_notes'] = []
_global_memory['plans'] = [] _global_memory['plans'] = []
_global_memory['tasks'] = [] _global_memory['tasks'] = []
_global_memory['research_subtasks'] = []
_global_memory['related_files'] = set() _global_memory['related_files'] = set()
_global_memory['tasks'] = {} _global_memory['tasks'] = {}
_global_memory['task_id_counter'] = 0 _global_memory['task_id_counter'] = 0
@ -38,7 +36,6 @@ def reset_memory():
_global_memory['plans'] = [] _global_memory['plans'] = []
_global_memory['tasks'] = {} _global_memory['tasks'] = {}
_global_memory['task_id_counter'] = 0 _global_memory['task_id_counter'] = 0
_global_memory['research_subtasks'] = []
def test_emit_key_facts_single_fact(reset_memory): def test_emit_key_facts_single_fact(reset_memory):
"""Test emitting a single key fact using emit_key_facts""" """Test emitting a single key fact using emit_key_facts"""
@ -431,15 +428,3 @@ def test_swap_task_order_after_delete(reset_memory):
assert _global_memory['tasks'][0] == "Task 3" assert _global_memory['tasks'][0] == "Task 3"
assert _global_memory['tasks'][2] == "Task 1" assert _global_memory['tasks'][2] == "Task 1"
def test_request_research_subtask(reset_memory):
"""Test requesting research subtasks"""
# Test adding a research subtask
subtask = "Research Python async patterns"
result = request_research_subtask(subtask)
# Verify return message
assert result == "Subtask added."
# Verify it was stored in memory
assert len(_global_memory['research_subtasks']) == 1
assert _global_memory['research_subtasks'][0] == subtask