Refactor: extract tool configs.
This commit is contained in:
parent
86d5fae30f
commit
bfb4453f11
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
Loading…
Reference in New Issue