fix one_shot_completed and allow reordering of tasks

This commit is contained in:
AI Christianson 2024-12-17 10:36:32 -05:00
parent db5ac38a5a
commit 3be54fac2f
6 changed files with 122 additions and 16 deletions

View File

@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- When key snippets are emitted, snippet files are auto added to related files. - When key snippets are emitted, snippet files are auto added to related files.
- Add base task to research subtask prompt. - Add base task to research subtask prompt.
- Adjust research prompt to make sure related files are related to the base task, not just the research subtask. - Adjust research prompt to make sure related files are related to the base task, not just the research subtask.
- Track tasks by ID and allow them to be deleted.
- Make one_shot_completed tool available to research agent.
## [0.6.0] - 2024-12-17 ## [0.6.0] - 2024-12-17

View File

@ -15,10 +15,10 @@ 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, get_memory_value, emit_key_facts, delete_key_facts, emit_expert_context, get_memory_value, emit_key_facts, delete_key_facts,
emit_key_snippets, delete_key_snippets, emit_key_snippets, delete_key_snippets,
emit_research_subtask, request_implementation, read_file_tool, write_file_tool, fuzzy_find_project_files, ripgrep_search, list_directory_tree, emit_research_subtask, request_complex_implementation, read_file_tool, write_file_tool, fuzzy_find_project_files, ripgrep_search, list_directory_tree,
file_str_replace file_str_replace, swap_task_order
) )
from ra_aid.tools.memory import _global_memory, get_related_files from ra_aid.tools.memory import _global_memory, get_related_files, one_shot_completed
from ra_aid import print_agent_output, print_stage_header, print_task_header, print_error from ra_aid import print_agent_output, print_stage_header, print_task_header, print_error
from ra_aid.prompts import ( from ra_aid.prompts import (
RESEARCH_PROMPT, RESEARCH_PROMPT,
@ -55,7 +55,8 @@ RESEARCH_TOOLS = [
list_directory_tree, list_directory_tree,
emit_research_subtask, emit_research_subtask,
run_shell_command, run_shell_command,
emit_research_notes emit_research_notes,
one_shot_completed
] ]
def parse_arguments(): def parse_arguments():
@ -142,7 +143,7 @@ def get_research_tools(research_only: bool = False, expert_enabled: bool = True)
tools.extend(EXPERT_TOOLS) tools.extend(EXPERT_TOOLS)
if not research_only: if not research_only:
tools.append(request_implementation) tools.append(request_complex_implementation)
return tools return tools
@ -151,6 +152,7 @@ def get_planning_tools(expert_enabled: bool = True) -> list:
list_directory_tree, list_directory_tree,
emit_plan, emit_plan,
emit_task, emit_task,
swap_task_order,
emit_related_files, emit_related_files,
emit_key_facts, emit_key_facts,
delete_key_facts, delete_key_facts,

View File

@ -96,7 +96,7 @@ Thoroughness and Completeness
Decision on Implementation Decision on Implementation
After completing your factual enumeration and description, decide: After completing your factual enumeration and description, decide:
If you see reasons that implementation changes will be required in the future, after documenting all findings, call request_implementation and specify why. If you see reasons that implementation changes will be required in the future, after documenting all findings, call request_complex_implementation and specify why.
If no changes are needed, simply state that no changes are required. If no changes are needed, simply state that no changes are required.
Be thorough on locating all potential change sites/gauging blast radius. Be thorough on locating all potential change sites/gauging blast radius.
@ -154,8 +154,6 @@ Guidelines:
API contracts, endpoints, or protocols it requires or provides API contracts, endpoints, or protocols it requires or provides
Testing strategies appropriate to the complexity of that sub-task Testing strategies appropriate to the complexity of that sub-task
You may include pseudocode, but not full code. You may include pseudocode, but not full code.
If you need to consult with the expert, do that *BEFORE* emitting any tasks.
After finalizing the overall approach: After finalizing the overall approach:
Use emit_plan to store the high-level implementation plan. Use emit_plan to store the high-level implementation plan.
@ -171,7 +169,7 @@ Guidelines:
IMPLEMENTATION_PROMPT = """Base-level task (for reference only): IMPLEMENTATION_PROMPT = """Base-level task (for reference only):
{base_task} --keep it simple {base_task} --keep it simple
Plan Overview: Plan Overview (for reference only, remember you are only implementing your specific task):
{plan} {plan}
Key Facts: Key Facts:

View File

@ -9,8 +9,8 @@ from .list_directory import list_directory_tree
from .ripgrep import ripgrep_search from .ripgrep import ripgrep_search
from .memory import ( from .memory import (
emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_facts, emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_facts,
request_implementation, skip_implementation, delete_key_facts, emit_research_subtask, request_complex_implementation, skip_implementation, delete_key_facts, emit_research_subtask,
emit_key_snippets, delete_key_snippets, emit_related_files emit_key_snippets, delete_key_snippets, emit_related_files, swap_task_order
) )
__all__ = [ __all__ = [
@ -28,12 +28,13 @@ __all__ = [
'get_memory_value', 'get_memory_value',
'list_directory_tree', 'list_directory_tree',
'read_file_tool', 'read_file_tool',
'request_implementation', 'request_complex_implementation',
'run_programming_task', 'run_programming_task',
'run_shell_command', 'run_shell_command',
'skip_implementation', 'skip_implementation',
'write_file_tool', 'write_file_tool',
'emit_research_subtask', 'emit_research_subtask',
'ripgrep_search', 'ripgrep_search',
'file_str_replace' 'file_str_replace',
'swap_task_order'
] ]

View File

@ -170,10 +170,11 @@ def delete_tasks(task_ids: List[int]) -> str:
return "Tasks deleted." return "Tasks deleted."
@tool("request_implementation") @tool("request_complex_implementation")
def request_implementation(reason: str) -> str: def request_complex_implementation(reason: str) -> str:
"""Request that implementation proceed after research/planning. """Request that implementation proceed after research/planning.
Used to indicate the agent should move to implementation stage. Used to indicate the agent should move to implementation stage.
Should be called when the implementation is more complex than a one-shot task.
Args: Args:
reason: Why implementation should proceed reason: Why implementation should proceed
@ -275,9 +276,43 @@ def delete_key_snippets(snippet_ids: List[int]) -> str:
return "Snippets deleted." return "Snippets deleted."
@tool("swap_task_order")
def swap_task_order(id1: int, id2: int) -> str:
"""Swap the order of two tasks in global memory by their IDs.
Args:
id1: First task ID
id2: Second task ID
Returns:
Success or error message depending on outcome
"""
# Validate IDs are different
if id1 == id2:
return "Cannot swap task with itself"
# Validate both IDs exist
if id1 not in _global_memory['tasks'] or id2 not in _global_memory['tasks']:
return "Invalid task ID(s)"
# Swap the tasks
_global_memory['tasks'][id1], _global_memory['tasks'][id2] = \
_global_memory['tasks'][id2], _global_memory['tasks'][id1]
# Display what was swapped
console.print(Panel(
Markdown(f"Swapped:\n- Task #{id1} ↔️ Task #{id2}"),
title="🔄 Tasks Reordered",
border_style="green"
))
return "Tasks swapped."
@tool("one_shot_completed") @tool("one_shot_completed")
def one_shot_completed(message: str) -> str: def one_shot_completed(message: str) -> str:
"""Signal that a one-shot task has been completed and execution should stop. """Signal that a one-shot task has been completed and execution should stop.
Only call this if you have already **fully** completed the original request.
Args: Args:
message: Completion message to display message: Completion message to display

View File

@ -10,7 +10,8 @@ from ra_aid.tools.memory import (
emit_related_files, emit_related_files,
get_related_files, get_related_files,
emit_task, emit_task,
delete_tasks delete_tasks,
swap_task_order
) )
@pytest.fixture @pytest.fixture
@ -363,6 +364,73 @@ def test_delete_tasks(reset_memory):
# Counter should remain unchanged after deletions # Counter should remain unchanged after deletions
assert _global_memory['task_id_counter'] == 3 assert _global_memory['task_id_counter'] == 3
def test_swap_task_order_valid_ids(reset_memory):
"""Test basic task swapping functionality"""
# Add test tasks
tasks = ["Task 1", "Task 2", "Task 3"]
for task in tasks:
emit_task.invoke({"task": task})
# Swap tasks 0 and 2
result = swap_task_order.invoke({"id1": 0, "id2": 2})
assert result == "Tasks swapped."
# Verify tasks were swapped
assert _global_memory['tasks'][0] == "Task 3"
assert _global_memory['tasks'][2] == "Task 1"
assert _global_memory['tasks'][1] == "Task 2" # Unchanged
def test_swap_task_order_invalid_ids(reset_memory):
"""Test error handling for invalid task IDs"""
# Add a test task
emit_task.invoke({"task": "Task 1"})
# Try to swap with non-existent ID
result = swap_task_order.invoke({"id1": 0, "id2": 999})
assert result == "Invalid task ID(s)"
# Verify original task unchanged
assert _global_memory['tasks'][0] == "Task 1"
def test_swap_task_order_same_id(reset_memory):
"""Test handling of attempt to swap a task with itself"""
# Add test task
emit_task.invoke({"task": "Task 1"})
# Try to swap task with itself
result = swap_task_order.invoke({"id1": 0, "id2": 0})
assert result == "Cannot swap task with itself"
# Verify task unchanged
assert _global_memory['tasks'][0] == "Task 1"
def test_swap_task_order_empty_tasks(reset_memory):
"""Test swapping behavior with empty tasks dictionary"""
result = swap_task_order.invoke({"id1": 0, "id2": 1})
assert result == "Invalid task ID(s)"
def test_swap_task_order_after_delete(reset_memory):
"""Test swapping after deleting a task"""
# Add test tasks
tasks = ["Task 1", "Task 2", "Task 3"]
for task in tasks:
emit_task.invoke({"task": task})
# Delete middle task
delete_tasks.invoke({"task_ids": [1]})
# Try to swap with deleted task
result = swap_task_order.invoke({"id1": 0, "id2": 1})
assert result == "Invalid task ID(s)"
# Try to swap remaining valid tasks
result = swap_task_order.invoke({"id1": 0, "id2": 2})
assert result == "Tasks swapped."
# Verify swap worked
assert _global_memory['tasks'][0] == "Task 3"
assert _global_memory['tasks'][2] == "Task 1"
def test_emit_research_subtask(reset_memory): def test_emit_research_subtask(reset_memory):
"""Test emitting research subtasks""" """Test emitting research subtasks"""
# Test adding a research subtask # Test adding a research subtask