diff --git a/ra_aid/tools/__init__.py b/ra_aid/tools/__init__.py index 10d0cb6..7bfa3db 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_files +from .programmer import run_programming_task from .expert import ask_expert, emit_expert_context from .read_file import read_file_tool from .fuzzy_find import fuzzy_find_project_files @@ -9,7 +9,7 @@ from .note_tech_debt import note_tech_debt from .memory import ( emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_facts, request_implementation, skip_implementation, delete_key_facts, emit_research_subtask, - emit_key_snippets, delete_key_snippets + emit_key_snippets, delete_key_snippets, emit_related_files ) __all__ = [ @@ -20,7 +20,7 @@ __all__ = [ 'emit_key_facts', 'emit_key_snippets', 'emit_plan', - 'emit_related_files', + 'emit_related_files', 'emit_research_notes', 'emit_task', 'fuzzy_find_project_files', @@ -32,7 +32,6 @@ __all__ = [ 'run_shell_command', 'skip_implementation', 'emit_research_subtask', - 'fuzzy_find_project_files', 'ripgrep_search', 'note_tech_debt' ] diff --git a/ra_aid/tools/memory.py b/ra_aid/tools/memory.py index cc1d7e1..06f369d 100644 --- a/ra_aid/tools/memory.py +++ b/ra_aid/tools/memory.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Any, Union, TypedDict, Optional, Sequence +from typing import Dict, List, Any, Union, TypedDict, Optional, Sequence, Set from ra_aid.exceptions import TaskCompletedException from rich.console import Console from rich.markdown import Markdown @@ -15,7 +15,7 @@ class SnippetInfo(TypedDict): console = Console() # Global memory store -_global_memory: Dict[str, Union[List[Any], Dict[int, str], Dict[int, SnippetInfo], int]] = { +_global_memory: Dict[str, Union[List[Any], Dict[int, str], Dict[int, SnippetInfo], int, Set[str]]] = { 'research_notes': [], 'plans': [], 'tasks': [], @@ -25,7 +25,8 @@ _global_memory: Dict[str, Union[List[Any], Dict[int, str], Dict[int, SnippetInfo 'key_snippets': {}, # Dict[int, SnippetInfo] - ID to snippet mapping 'key_snippet_id_counter': 0, # Counter for generating unique snippet IDs 'implementation_requested': [], - 'implementation_skipped': [] + 'implementation_skipped': [], + 'related_files': set() } @tool("emit_research_notes") @@ -259,6 +260,44 @@ def one_shot_completed(message: str) -> str: raise ValueError("Cannot complete in one shot - implementation was requested") raise TaskCompletedException(message) +def get_related_files() -> Set[str]: + """Get the current set of related files. + + Returns: + Set of file paths that have been marked as related + """ + return _global_memory['related_files'] + +@tool("emit_related_files") +def emit_related_files(files: List[str]) -> str: + """Store multiple related files that tools should work with. + + Args: + files: List of file paths to add + + Returns: + Confirmation message + """ + results = [] + added_files = [] + + # Process unique files + for file in set(files): # Remove duplicates in input + if file not in _global_memory['related_files']: + _global_memory['related_files'].add(file) + added_files.append(file) + results.append(f"Added related file: {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 Noted:**\n{files_added_md}" + console.print(Panel(Markdown(md_content), + title="📁 Related Files Noted", + border_style="green")) + + return "Files noted." + def get_memory_value(key: str) -> str: """Get a value from global memory. diff --git a/ra_aid/tools/programmer.py b/ra_aid/tools/programmer.py index 0a342c5..a27ea5b 100644 --- a/ra_aid/tools/programmer.py +++ b/ra_aid/tools/programmer.py @@ -8,46 +8,10 @@ from rich.markdown import Markdown from rich.text import Text from ra_aid.proc.interactive import run_interactive_command from pydantic import BaseModel, Field -from .memory import get_memory_value from ra_aid.text.processing import truncate_output console = Console() -# Keep track of related files globally -related_files: List[str] = [] -related_files_set: Set[str] = set() - -@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: - files: List of file paths to add - - Returns: - List of confirmation messages for added files - """ - global related_files, related_files_set - results = [] - added_files = [] - - # 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}") - - # Rich output - single consolidated panel - if added_files: - files_added_md = '\n'.join(f"- `{file}`" for file in added_files) - md_content = f"**Files Noted:**\n{files_added_md}" - console.print(Panel(Markdown(md_content), - title="📁 Related Files Noted", - border_style="green")) - - return results class RunProgrammingTaskInput(BaseModel): instructions: str = Field(description="Instructions for the programming task") @@ -88,13 +52,8 @@ def run_programming_task(input: RunProgrammingTaskInput) -> Dict[str, Union[str, command.append(input.instructions) - # Use both input files and related files - files_to_use = set(related_files) # Start with related files - if input.files: # Add any additional input files - files_to_use.update(input.files) - - if files_to_use: - command.extend(list(files_to_use)) + if input.files: + command.extend(input.files) # Create a pretty display of what we're doing task_display = [ @@ -102,10 +61,10 @@ def run_programming_task(input: RunProgrammingTaskInput) -> Dict[str, Union[str, f"{input.instructions}\n" ] - if files_to_use: + if input.files: task_display.extend([ "\n## Files\n", - *[f"- `{file}`\n" for file in files_to_use] + *[f"- `{file}`\n" for file in input.files] ]) markdown_content = "".join(task_display) @@ -138,4 +97,4 @@ def run_programming_task(input: RunProgrammingTaskInput) -> Dict[str, Union[str, } # Export the functions -__all__ = ['run_programming_task', 'emit_related_files'] +__all__ = ['run_programming_task'] diff --git a/tests/ra_aid/tools/test_memory.py b/tests/ra_aid/tools/test_memory.py index 8e2f5c4..0411f6d 100644 --- a/tests/ra_aid/tools/test_memory.py +++ b/tests/ra_aid/tools/test_memory.py @@ -6,7 +6,9 @@ from ra_aid.tools.memory import ( emit_key_facts, delete_key_facts, emit_key_snippets, - delete_key_snippets + delete_key_snippets, + emit_related_files, + get_related_files ) @pytest.fixture @@ -20,6 +22,7 @@ def reset_memory(): _global_memory['plans'] = [] _global_memory['tasks'] = [] _global_memory['research_subtasks'] = [] + _global_memory['related_files'] = set() yield # Clean up after test _global_memory['key_facts'] = {} @@ -207,6 +210,35 @@ def test_delete_key_snippets_empty(reset_memory): # Verify snippet still exists assert 0 in _global_memory['key_snippets'] +def test_emit_related_files_basic(reset_memory): + """Test basic adding of files""" + # Test adding single file + result = emit_related_files.invoke({"files": ["test.py"]}) + assert result == "Files noted." + assert get_related_files() == {"test.py"} + + # Test adding multiple files + result = emit_related_files.invoke({"files": ["main.py", "utils.py"]}) + assert result == "Files noted." + assert get_related_files() == {"test.py", "main.py", "utils.py"} + +def test_get_related_files_empty(reset_memory): + """Test getting related files when none added""" + assert get_related_files() == set() + +def test_emit_related_files_duplicates(reset_memory): + """Test that duplicate files are handled correctly""" + # Add initial files + result = emit_related_files.invoke({"files": ["test.py", "main.py"]}) + assert result == "Files noted." + assert get_related_files() == {"test.py", "main.py"} + + # Try adding duplicates + result = emit_related_files.invoke({"files": ["test.py", "main.py", "test.py"]}) + assert result == "Files noted." + # Set should still only contain unique entries + assert get_related_files() == {"test.py", "main.py"} + def test_key_snippets_integration(reset_memory): """Integration test for key snippets functionality""" # Initial snippets to add