From bfb4453f11c3f767340e0986ea419b8a8f4b4f6b Mon Sep 17 00:00:00 2001 From: AI Christianson Date: Thu, 19 Dec 2024 14:08:53 -0500 Subject: [PATCH] Refactor: extract tool configs. --- ra_aid/__main__.py | 103 ++---------------------------- ra_aid/tool_configs.py | 98 ++++++++++++++++++++++++++++ tests/ra_aid/test_tool_configs.py | 51 +++++++++++++++ 3 files changed, 155 insertions(+), 97 deletions(-) create mode 100644 ra_aid/tool_configs.py create mode 100644 tests/ra_aid/test_tool_configs.py diff --git a/ra_aid/__main__.py b/ra_aid/__main__.py index f0a13c2..888e36b 100644 --- a/ra_aid/__main__.py +++ b/ra_aid/__main__.py @@ -32,52 +32,12 @@ import time from anthropic import APIError, APITimeoutError, RateLimitError, InternalServerError from ra_aid.llm import initialize_llm -# Read-only tools that don't modify system state -def get_read_only_tools(human_interaction: bool = False) -> list: - """Get the list of read-only tools, optionally including human interaction tools.""" - tools = [ - emit_related_files, - emit_key_facts, - delete_key_facts, - emit_key_snippets, - delete_key_snippets, - list_directory_tree, - read_file_tool, - fuzzy_find_project_files, - ripgrep_search, - run_shell_command # can modify files, but we still need it for read-only tasks. - ] - - if human_interaction: - tools.append(ask_human) - - return tools - -READ_ONLY_TOOLS = get_read_only_tools() - -# Tools that can modify files or system state -MODIFICATION_TOOLS = [ - run_programming_task -] - -# Common tools used across multiple agents -COMMON_TOOLS = READ_ONLY_TOOLS + [] - -# Expert-specific tools -EXPERT_TOOLS = [ - emit_expert_context, - ask_expert -] - -# Research-specific tools -RESEARCH_TOOLS = [ - request_research_subtask, - emit_research_notes, - one_shot_completed, - monorepo_detected, - existing_project_detected, - ui_detected -] +from ra_aid.tool_configs import ( + get_read_only_tools, + get_research_tools, + get_planning_tools, + get_implementation_tools +) def parse_arguments(): parser = argparse.ArgumentParser( @@ -157,57 +117,6 @@ research_memory = MemorySaver() planning_memory = MemorySaver() implementation_memory = MemorySaver() -def get_research_tools(research_only: bool = False, expert_enabled: bool = True, human_interaction: bool = False) -> list: - """Get the list of research tools based on mode and whether expert is enabled.""" - # Start with read-only tools - tools = get_read_only_tools(human_interaction).copy() - - tools.extend(RESEARCH_TOOLS) - - # Add modification tools if not research_only - if not research_only: - tools.extend(MODIFICATION_TOOLS) - tools.append(request_implementation) - - # Add expert tools if enabled - if expert_enabled: - tools.extend(EXPERT_TOOLS) - - return tools - -def get_planning_tools(expert_enabled: bool = True) -> list: - """Get the list of planning tools based on whether expert is enabled.""" - # Start with common tools - tools = COMMON_TOOLS.copy() - - # Add planning-specific tools - planning_tools = [ - delete_tasks, - emit_plan, - emit_task, - swap_task_order - ] - tools.extend(planning_tools) - - # Add expert tools if enabled - if expert_enabled: - tools.extend(EXPERT_TOOLS) - - return tools - -def get_implementation_tools(expert_enabled: bool = True) -> list: - """Get the list of implementation tools based on whether expert is enabled.""" - # Start with common tools - tools = COMMON_TOOLS.copy() - - # Add modification tools since it's not research-only - tools.extend(MODIFICATION_TOOLS) - - # Add expert tools if enabled - if expert_enabled: - tools.extend(EXPERT_TOOLS) - - return tools def is_informational_query() -> bool: """Determine if the current query is informational based on implementation_requested state.""" diff --git a/ra_aid/tool_configs.py b/ra_aid/tool_configs.py new file mode 100644 index 0000000..f6397a3 --- /dev/null +++ b/ra_aid/tool_configs.py @@ -0,0 +1,98 @@ +from typing import List +from ra_aid.tools import ( + ask_expert, ask_human, run_shell_command, run_programming_task, + emit_research_notes, emit_plan, emit_related_files, emit_task, + emit_expert_context, emit_key_facts, delete_key_facts, + emit_key_snippets, delete_key_snippets, delete_tasks, + request_research_subtask, request_implementation, read_file_tool, + fuzzy_find_project_files, ripgrep_search, list_directory_tree, + swap_task_order, monorepo_detected, existing_project_detected, ui_detected +) +from ra_aid.tools.memory import one_shot_completed + +# Read-only tools that don't modify system state +def get_read_only_tools(human_interaction: bool = False) -> list: + """Get the list of read-only tools, optionally including human interaction tools.""" + tools = [ + emit_related_files, + emit_key_facts, + delete_key_facts, + emit_key_snippets, + delete_key_snippets, + list_directory_tree, + read_file_tool, + fuzzy_find_project_files, + ripgrep_search, + run_shell_command # can modify files, but we still need it for read-only tasks. + ] + + if human_interaction: + tools.append(ask_human) + + return tools + +# Define constant tool groups +READ_ONLY_TOOLS = get_read_only_tools() +MODIFICATION_TOOLS = [run_programming_task] +COMMON_TOOLS = READ_ONLY_TOOLS + [] +EXPERT_TOOLS = [emit_expert_context, ask_expert] +RESEARCH_TOOLS = [ + request_research_subtask, + emit_research_notes, + one_shot_completed, + monorepo_detected, + existing_project_detected, + ui_detected +] + +def get_research_tools(research_only: bool = False, expert_enabled: bool = True, human_interaction: bool = False) -> list: + """Get the list of research tools based on mode and whether expert is enabled.""" + # Start with read-only tools + tools = get_read_only_tools(human_interaction).copy() + + tools.extend(RESEARCH_TOOLS) + + # Add modification tools if not research_only + if not research_only: + tools.extend(MODIFICATION_TOOLS) + tools.append(request_implementation) + + # Add expert tools if enabled + if expert_enabled: + tools.extend(EXPERT_TOOLS) + + return tools + +def get_planning_tools(expert_enabled: bool = True) -> list: + """Get the list of planning tools based on whether expert is enabled.""" + # Start with common tools + tools = COMMON_TOOLS.copy() + + # Add planning-specific tools + planning_tools = [ + delete_tasks, + emit_plan, + emit_task, + swap_task_order + ] + tools.extend(planning_tools) + + # Add expert tools if enabled + if expert_enabled: + tools.extend(EXPERT_TOOLS) + + return tools + +def get_implementation_tools(expert_enabled: bool = True) -> list: + """Get the list of implementation tools based on whether expert is enabled.""" + # Start with common tools + tools = COMMON_TOOLS.copy() + + # Add modification tools since it's not research-only + tools.extend(MODIFICATION_TOOLS) + + # Add expert tools if enabled + if expert_enabled: + tools.extend(EXPERT_TOOLS) + + return tools diff --git a/tests/ra_aid/test_tool_configs.py b/tests/ra_aid/test_tool_configs.py new file mode 100644 index 0000000..47f96a0 --- /dev/null +++ b/tests/ra_aid/test_tool_configs.py @@ -0,0 +1,51 @@ +import pytest +from ra_aid.tool_configs import ( + get_read_only_tools, + get_research_tools, + get_planning_tools, + get_implementation_tools +) + +def test_get_read_only_tools(): + # Test without human interaction + tools = get_read_only_tools(human_interaction=False) + assert len(tools) > 0 + assert all(callable(tool) for tool in tools) + + # Test with human interaction + tools_with_human = get_read_only_tools(human_interaction=True) + assert len(tools_with_human) == len(tools) + 1 + +def test_get_research_tools(): + # Test basic research tools + tools = get_research_tools() + assert len(tools) > 0 + assert all(callable(tool) for tool in tools) + + # Test without expert + tools_no_expert = get_research_tools(expert_enabled=False) + assert len(tools_no_expert) < len(tools) + + # Test research-only mode + tools_research_only = get_research_tools(research_only=True) + assert len(tools_research_only) < len(tools) + +def test_get_planning_tools(): + # Test with expert enabled + tools = get_planning_tools(expert_enabled=True) + assert len(tools) > 0 + assert all(callable(tool) for tool in tools) + + # Test without expert + tools_no_expert = get_planning_tools(expert_enabled=False) + assert len(tools_no_expert) < len(tools) + +def test_get_implementation_tools(): + # Test with expert enabled + tools = get_implementation_tools(expert_enabled=True) + assert len(tools) > 0 + assert all(callable(tool) for tool in tools) + + # Test without expert + tools_no_expert = get_implementation_tools(expert_enabled=False) + assert len(tools_no_expert) < len(tools)