diff --git a/ra_aid/__main__.py b/ra_aid/__main__.py index 1d4c1dd..1a1e58e 100644 --- a/ra_aid/__main__.py +++ b/ra_aid/__main__.py @@ -9,7 +9,7 @@ from langgraph.checkpoint.memory import MemorySaver from langgraph.prebuilt import create_react_agent from ra_aid.tools import ( ask_expert, run_shell_command, run_programming_task, - emit_research_notes, emit_plan, emit_related_file, emit_task, + emit_research_notes, emit_plan, emit_related_files, emit_task, emit_expert_context, get_memory_value, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, request_implementation, read_file_tool, emit_research_subtask, @@ -62,9 +62,9 @@ planning_memory = MemorySaver() implementation_memory = MemorySaver() # Define tool sets for each stage -research_tools = [list_directory_tree, emit_research_subtask, run_shell_command, emit_expert_context, ask_expert, emit_research_notes, emit_related_file, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, request_implementation, read_file_tool, fuzzy_find_project_files, ripgrep_search] -planning_tools = [list_directory_tree, emit_expert_context, ask_expert, emit_plan, emit_task, emit_related_file, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, read_file_tool, fuzzy_find_project_files, ripgrep_search] -implementation_tools = [list_directory_tree, run_shell_command, emit_expert_context, ask_expert, run_programming_task, emit_related_file, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, read_file_tool, fuzzy_find_project_files, ripgrep_search] +research_tools = [list_directory_tree, emit_research_subtask, run_shell_command, emit_expert_context, ask_expert, emit_research_notes, emit_related_files, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, request_implementation, read_file_tool, fuzzy_find_project_files, ripgrep_search] +planning_tools = [list_directory_tree, emit_expert_context, ask_expert, emit_plan, emit_task, emit_related_files, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, read_file_tool, fuzzy_find_project_files, ripgrep_search] +implementation_tools = [list_directory_tree, run_shell_command, emit_expert_context, ask_expert, run_programming_task, emit_related_files, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, read_file_tool, fuzzy_find_project_files, ripgrep_search] # Create stage-specific agents with individual memory objects research_agent = create_react_agent(model, research_tools, checkpointer=research_memory) @@ -254,7 +254,7 @@ def main(): # Run research stage print_stage_header("RESEARCH STAGE") - research_prompt = f"""User query: {base_task} + research_prompt = f"""User query: {base_task} --keep it simple {RESEARCH_PROMPT} diff --git a/ra_aid/prompts.py b/ra_aid/prompts.py index 3659625..38e6d17 100644 --- a/ra_aid/prompts.py +++ b/ra_aid/prompts.py @@ -46,7 +46,7 @@ Tools and Methodology Reporting Findings Use emit_research_notes to record detailed, fact-based observations about what currently exists. - For each significant file or directory that is part of the codebase, use emit_related_file to list it. + For each significant file or directory that is part of the codebase, use emit_related_files to list it. Your research notes should be strictly about what you have observed: Document files by their names and locations. Document discovered documentation files and their contents at a high level (e.g., "There is a README.md in the root directory that explains the folder structure"). @@ -83,7 +83,7 @@ If there is a top-level README.md or docs/ folder, always start with that. # Planning stage prompt - guides task breakdown and implementation planning # Includes a directive to scale complexity with request size. PLANNING_PROMPT = """Base Task: -{base_task} +{base_task} --keep it simple Research Notes: @@ -178,7 +178,7 @@ Instructions: # Implementation stage prompt - guides specific task implementation # Added instruction to adjust complexity of implementation to match request. IMPLEMENTATION_PROMPT = """Base-level task (for reference only): -{base_task} +{base_task} --keep it simple Plan Overview: {plan} diff --git a/ra_aid/tools/__init__.py b/ra_aid/tools/__init__.py index f723fd0..3818ea0 100644 --- a/ra_aid/tools/__init__.py +++ b/ra_aid/tools/__init__.py @@ -1,5 +1,5 @@ from .shell import run_shell_command -from .programmer import run_programming_task, emit_related_file +from .programmer import run_programming_task, emit_related_files from .expert import ask_expert, emit_expert_context from .read_file import read_file_tool from .fuzzy_find import fuzzy_find_project_files @@ -19,7 +19,7 @@ __all__ = [ 'emit_key_facts', 'emit_key_snippet', 'emit_plan', - 'emit_related_file', + 'emit_related_files', 'emit_research_notes', 'emit_task', 'fuzzy_find_project_files', diff --git a/ra_aid/tools/programmer.py b/ra_aid/tools/programmer.py index 4d7948a..4fd4129 100644 --- a/ra_aid/tools/programmer.py +++ b/ra_aid/tools/programmer.py @@ -1,5 +1,5 @@ import subprocess -from typing import List, Optional, Dict, Union +from typing import List, Optional, Dict, Union, Set from langchain_core.tools import tool from rich.console import Console from rich.panel import Panel @@ -14,29 +14,40 @@ from ra_aid.text.processing import truncate_output console = Console() # Keep track of related files globally -related_files = [] +related_files: List[str] = [] +related_files_set: Set[str] = set() -@tool -def emit_related_file(file: str) -> str: - """Add a single file to the list of files that the programmer tool should work with. +@tool("emit_related_files") +def emit_related_files(files: List[str]) -> List[str]: + """Store multiple related files that the programmer tool should work with. Args: - file: File path to add + files: List of file paths to add Returns: - A confirmation message with the added file + List of confirmation messages for added files """ - global related_files + global related_files, related_files_set + results = [] + added_files = [] - # Check if file is already in the set - if file not in related_files: - related_files.append(file) + # Process unique files + for file in set(files): # Remove duplicates in input + if file not in related_files_set: + related_files.append(file) + related_files_set.add(file) + added_files.append(file) + results.append(f"Added related file: {file}") - md_content = f"`{file}`" + # Rich output - single consolidated panel + if added_files: + files_added_md = '\n'.join(f"- `{file}`" for file in added_files) + md_content = f"**Files Added:**\n{files_added_md}" + console.print(Panel(Markdown(md_content), + title="📁 Related Files Added", + border_style="green")) - # Display in a panel - console.print(Panel(Markdown(md_content), title="📁 Related File Added", border_style="green")) - return md_content + return results class RunProgrammingTaskInput(BaseModel): instructions: str = Field(description="Instructions for the programming task") @@ -50,9 +61,9 @@ def run_programming_task(input: RunProgrammingTaskInput) -> Dict[str, Union[str, The programmer can edit multiple files at once and is intelligent. - If any new files are created, remember to emit them using the emit_related_file tool once this tool completes. + If any new files are created, remember to emit them using the emit_related_files tool once this tool completes. - Additionally, before invoking this tool, make sure all existing related files have been emitted using the emit_related_file tool. + Additionally, before invoking this tool, make sure all existing related files have been emitted using the emit_related_files tool. Args: instructions: Instructions for the programming task @@ -127,4 +138,4 @@ def run_programming_task(input: RunProgrammingTaskInput) -> Dict[str, Union[str, } # Export the functions -__all__ = ['run_programming_task', 'emit_related_file'] +__all__ = ['run_programming_task', 'emit_related_files']