Added web research tool config.

This commit is contained in:
user 2024-12-23 17:08:07 -05:00
parent 48f4fbc41d
commit f5482dce2a
5 changed files with 55 additions and 14 deletions

View File

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

View File

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

View File

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

View File

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

View File

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