From 72f8c33e4300f54fc1d7616940b61aa7674cec80 Mon Sep 17 00:00:00 2001 From: user Date: Tue, 24 Dec 2024 08:31:08 -0500 Subject: [PATCH] Fix bugs. --- ra_aid/__main__.py | 2 +- ra_aid/agent_utils.py | 13 +++++-- ra_aid/tool_configs.py | 69 ++++++++++++++++++++++++++++---------- ra_aid/tools/programmer.py | 2 +- 4 files changed, 63 insertions(+), 23 deletions(-) diff --git a/ra_aid/__main__.py b/ra_aid/__main__.py index 93c2450..4c7246d 100644 --- a/ra_aid/__main__.py +++ b/ra_aid/__main__.py @@ -161,7 +161,7 @@ def main(): # Create chat agent with appropriate tools chat_agent = create_react_agent( model, - get_chat_tools(expert_enabled=expert_enabled), + get_chat_tools(expert_enabled=expert_enabled, web_research_enabled=web_research_enabled), checkpointer=MemorySaver() ) diff --git a/ra_aid/agent_utils.py b/ra_aid/agent_utils.py index 90514d8..2354d18 100644 --- a/ra_aid/agent_utils.py +++ b/ra_aid/agent_utils.py @@ -76,6 +76,7 @@ def run_research_agent( expert_enabled: bool = False, research_only: bool = False, hil: bool = False, + web_research_enabled: bool = False, memory: Optional[Any] = None, config: Optional[dict] = None, thread_id: Optional[str] = None, @@ -89,6 +90,7 @@ def run_research_agent( expert_enabled: Whether expert mode is enabled research_only: Whether this is a research-only task hil: Whether human-in-the-loop mode is enabled + web_research_enabled: Whether web research is enabled memory: Optional memory instance to use config: Optional configuration dictionary thread_id: Optional thread ID (defaults to new UUID) @@ -117,7 +119,8 @@ def run_research_agent( tools = get_research_tools( research_only=research_only, expert_enabled=expert_enabled, - human_interaction=hil + human_interaction=hil, + web_research_enabled=config.get('web_research', False) ) # Create agent @@ -166,6 +169,7 @@ def run_web_research_agent( *, expert_enabled: bool = False, hil: bool = False, + web_research_enabled: bool = False, memory: Optional[Any] = None, config: Optional[dict] = None, thread_id: Optional[str] = None, @@ -178,6 +182,7 @@ def run_web_research_agent( model: The LLM model to use expert_enabled: Whether expert mode is enabled hil: Whether human-in-the-loop mode is enabled + web_research_enabled: Whether web research is enabled memory: Optional memory instance to use config: Optional configuration dictionary thread_id: Optional thread ID (defaults to new UUID) @@ -274,7 +279,7 @@ def run_planning_agent( thread_id = str(uuid.uuid4()) # Configure tools - tools = get_planning_tools(expert_enabled=expert_enabled) + tools = get_planning_tools(expert_enabled=expert_enabled, web_research_enabled=config.get('web_research', False)) # Create agent agent = create_react_agent(model, tools, checkpointer=memory) @@ -318,6 +323,7 @@ def run_task_implementation_agent( model, *, expert_enabled: bool = False, + web_research_enabled: bool = False, memory: Optional[Any] = None, config: Optional[dict] = None, thread_id: Optional[str] = None @@ -331,6 +337,7 @@ def run_task_implementation_agent( related_files: List of related files model: The LLM model to use expert_enabled: Whether expert mode is enabled + web_research_enabled: Whether web research is enabled memory: Optional memory instance to use config: Optional configuration dictionary thread_id: Optional thread ID (defaults to new UUID) @@ -347,7 +354,7 @@ def run_task_implementation_agent( thread_id = str(uuid.uuid4()) # Configure tools - tools = get_implementation_tools(expert_enabled=expert_enabled) + tools = get_implementation_tools(expert_enabled=expert_enabled, web_research_enabled=config.get('web_research', False)) # Create agent agent = create_react_agent(model, tools, checkpointer=memory) diff --git a/ra_aid/tool_configs.py b/ra_aid/tool_configs.py index bd25746..9a6fc2e 100644 --- a/ra_aid/tool_configs.py +++ b/ra_aid/tool_configs.py @@ -12,8 +12,16 @@ 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 # 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.""" +def get_read_only_tools(human_interaction: bool = False, web_research_enabled: bool = False) -> list: + """Get the list of read-only tools, optionally including human interaction tools. + + Args: + human_interaction: Whether to include human interaction tools + web_research_enabled: Whether to include web research tools + + Returns: + List of tool functions + """ tools = [ emit_related_files, emit_key_facts, @@ -25,9 +33,11 @@ def get_read_only_tools(human_interaction: bool = False) -> list: read_file_tool, fuzzy_find_project_files, ripgrep_search, - run_shell_command, # can modify files, but we still need it for read-only tasks. - request_web_research + run_shell_command # can modify files, but we still need it for read-only tasks. ] + + if web_research_enabled: + tools.append(request_web_research) if human_interaction: tools.append(ask_human) @@ -37,7 +47,7 @@ def get_read_only_tools(human_interaction: bool = False) -> list: # Define constant tool groups READ_ONLY_TOOLS = get_read_only_tools() MODIFICATION_TOOLS = [run_programming_task] -COMMON_TOOLS = READ_ONLY_TOOLS + [] +COMMON_TOOLS = get_read_only_tools() EXPERT_TOOLS = [emit_expert_context, ask_expert] RESEARCH_TOOLS = [ emit_research_notes, @@ -47,10 +57,17 @@ RESEARCH_TOOLS = [ 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.""" +def get_research_tools(research_only: bool = False, expert_enabled: bool = True, human_interaction: bool = False, web_research_enabled: bool = False) -> list: + """Get the list of research tools based on mode and whether expert is enabled. + + Args: + research_only: Whether to exclude modification tools + expert_enabled: Whether to include expert tools + human_interaction: Whether to include human interaction tools + web_research_enabled: Whether to include web research tools + """ # Start with read-only tools - tools = get_read_only_tools(human_interaction).copy() + tools = get_read_only_tools(human_interaction, web_research_enabled).copy() tools.extend(RESEARCH_TOOLS) @@ -68,10 +85,15 @@ def get_research_tools(research_only: bool = False, expert_enabled: bool = True, 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() +def get_planning_tools(expert_enabled: bool = True, web_research_enabled: bool = False) -> list: + """Get the list of planning tools based on whether expert is enabled. + + Args: + expert_enabled: Whether to include expert tools + web_research_enabled: Whether to include web research tools + """ + # Start with read-only tools + tools = get_read_only_tools(web_research_enabled=web_research_enabled).copy() # Add planning-specific tools planning_tools = [ @@ -90,10 +112,15 @@ def get_planning_tools(expert_enabled: bool = True) -> list: 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() +def get_implementation_tools(expert_enabled: bool = True, web_research_enabled: bool = False) -> list: + """Get the list of implementation tools based on whether expert is enabled. + + Args: + expert_enabled: Whether to include expert tools + web_research_enabled: Whether to include web research tools + """ + # Start with read-only tools + tools = get_read_only_tools(web_research_enabled=web_research_enabled).copy() # Add modification tools since it's not research-only tools.extend(MODIFICATION_TOOLS) @@ -127,21 +154,27 @@ def get_web_research_tools(expert_enabled: bool = True) -> list: return tools -def get_chat_tools(expert_enabled: bool = True) -> list: +def get_chat_tools(expert_enabled: bool = True, web_research_enabled: bool = False) -> list: """Get the list of tools available in chat mode. Chat mode includes research and implementation capabilities but excludes complex planning tools. Human interaction is always enabled. + + Args: + expert_enabled: Whether to include expert tools + web_research_enabled: Whether to include web research tools """ tools = [ ask_human, request_research, request_research_and_implementation, - request_web_research, emit_key_facts, delete_key_facts, delete_key_snippets, deregister_related_files ] + if web_research_enabled: + tools.append(request_web_research) + return tools diff --git a/ra_aid/tools/programmer.py b/ra_aid/tools/programmer.py index 1656a6d..ce6843c 100644 --- a/ra_aid/tools/programmer.py +++ b/ra_aid/tools/programmer.py @@ -1,4 +1,4 @@ -import subprocess +import os from typing import List, Optional, Dict, Union, Set from langchain_core.tools import tool from rich.console import Console