diff --git a/ra_aid/__main__.py b/ra_aid/__main__.py index 6a22028..93c2450 100644 --- a/ra_aid/__main__.py +++ b/ra_aid/__main__.py @@ -6,10 +6,9 @@ from rich.console import Console from langgraph.checkpoint.memory import MemorySaver from langgraph.prebuilt import create_react_agent from ra_aid.env import validate_environment -from ra_aid.tools.memory import _global_memory, get_related_files, get_memory_value +from ra_aid.tools.memory import _global_memory from ra_aid.tools.human import ask_human from ra_aid import print_stage_header, print_error, print_interrupt -from ra_aid.tools.agent import CANCELLED_BY_USER_REASON from ra_aid.tools.human import ask_human from ra_aid.agent_utils import ( AgentInterrupt, @@ -18,15 +17,12 @@ from ra_aid.agent_utils import ( run_planning_agent ) from ra_aid.prompts import ( - PLANNING_PROMPT, CHAT_PROMPT, - EXPERT_PROMPT_SECTION_PLANNING, - HUMAN_PROMPT_SECTION_PLANNING, + WEB_RESEARCH_PROMPT_SECTION_CHAT ) from ra_aid.llm import initialize_llm from ra_aid.tool_configs import ( - get_planning_tools, get_chat_tools ) diff --git a/ra_aid/agent_utils.py b/ra_aid/agent_utils.py index 13b6a55..a05c6c4 100644 --- a/ra_aid/agent_utils.py +++ b/ra_aid/agent_utils.py @@ -24,7 +24,8 @@ from ra_aid.console.output import print_agent_output from ra_aid.tool_configs import ( get_implementation_tools, get_research_tools, - get_planning_tools + get_planning_tools, + get_web_research_tools ) from ra_aid.prompts import ( IMPLEMENTATION_PROMPT, @@ -200,11 +201,8 @@ def run_web_research_agent( if thread_id is None: thread_id = str(uuid.uuid4()) - # Configure tools - tools = get_research_tools( - expert_enabled=expert_enabled, - human_interaction=hil - ) + # Configure tools using restricted web research toolset + tools = get_web_research_tools(expert_enabled=expert_enabled) # Create agent agent = create_react_agent(model, tools, checkpointer=memory) diff --git a/ra_aid/prompts.py b/ra_aid/prompts.py index 747ce2b..b731fac 100644 --- a/ra_aid/prompts.py +++ b/ra_aid/prompts.py @@ -632,3 +632,10 @@ You have often been criticized for: NEVER ANNOUNCE WHAT YOU ARE DOING, JUST DO IT! """ +WEB_RESEARCH_PROMPT_SECTION_CHAT = """ +Web Research: + If you need to obtain additional context from online sources during chat: + - Use request_web_research to gather relevant information + - Wait for web research results before proceeding + - Web research can help provide up-to-date information and best practices +""" diff --git a/ra_aid/tool_configs.py b/ra_aid/tool_configs.py index bbba83d..d4ce9da 100644 --- a/ra_aid/tool_configs.py +++ b/ra_aid/tool_configs.py @@ -6,7 +6,7 @@ from ra_aid.tools import ( emit_key_snippets, delete_key_snippets, deregister_related_files, delete_tasks, read_file_tool, fuzzy_find_project_files, ripgrep_search, list_directory_tree, swap_task_order, monorepo_detected, existing_project_detected, ui_detected, - task_completed, plan_implementation_completed + task_completed, plan_implementation_completed, web_search_tavily ) from ra_aid.tools.memory import one_shot_completed from ra_aid.tools.agent import request_research, request_implementation, request_research_and_implementation, request_task_implementation, request_web_research @@ -107,6 +107,25 @@ def get_implementation_tools(expert_enabled: bool = True) -> list: return tools +def get_web_research_tools(expert_enabled: bool = True) -> list: + """Get the list of tools available for web research. + + Args: + expert_enabled: Whether expert tools should be included + + Returns: + list: List of tools configured for web research + """ + tools = [ + web_search_tavily + ] + + if expert_enabled: + tools.append(emit_expert_context) + tools.append(ask_expert) + + return tools + def get_chat_tools(expert_enabled: bool = True) -> list: """Get the list of tools available in chat mode. diff --git a/tests/ra_aid/test_tool_configs.py b/tests/ra_aid/test_tool_configs.py index 47f96a0..a70f31b 100644 --- a/tests/ra_aid/test_tool_configs.py +++ b/tests/ra_aid/test_tool_configs.py @@ -3,7 +3,8 @@ from ra_aid.tool_configs import ( get_read_only_tools, get_research_tools, get_planning_tools, - get_implementation_tools + get_implementation_tools, + get_web_research_tools ) def test_get_read_only_tools(): @@ -49,3 +50,23 @@ def test_get_implementation_tools(): # Test without expert tools_no_expert = get_implementation_tools(expert_enabled=False) assert len(tools_no_expert) < len(tools) + +def test_get_web_research_tools(): + # Test with expert enabled + tools = get_web_research_tools(expert_enabled=True) + assert len(tools) == 3 + assert all(callable(tool) for tool in tools) + + # Get tool names and verify exact matches + tool_names = [tool.name for tool in tools] + expected_names = ['emit_expert_context', 'ask_expert', 'web_search_tavily'] + assert sorted(tool_names) == sorted(expected_names) + + # Test without expert enabled + tools_no_expert = get_web_research_tools(expert_enabled=False) + assert len(tools_no_expert) == 1 + assert all(callable(tool) for tool in tools_no_expert) + + # Verify exact tool name when expert is disabled + tool_names_no_expert = [tool.name for tool in tools_no_expert] + assert tool_names_no_expert == ['web_search_tavily']