Fix bugs.

This commit is contained in:
user 2024-12-24 08:31:08 -05:00
parent c0fa79fe77
commit 72f8c33e43
4 changed files with 63 additions and 23 deletions

View File

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

View File

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

View File

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

View File

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