Refactor: extract tool configs.

This commit is contained in:
AI Christianson 2024-12-19 14:08:53 -05:00
parent 86d5fae30f
commit bfb4453f11
3 changed files with 155 additions and 97 deletions

View File

@ -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."""

98
ra_aid/tool_configs.py Normal file
View File

@ -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

View File

@ -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)