From 80486ec533f2dcf7a53b71a2059b16a151cb33b4 Mon Sep 17 00:00:00 2001 From: user Date: Sun, 22 Dec 2024 14:39:16 -0500 Subject: [PATCH] Convert related files system to ID-based storage - Refactored related files storage to use dictionary with integer IDs as keys - Added related_file_id_counter for generating unique file IDs - Updated get_related_files() to return formatted strings with IDs - Added delete_related_files() function for removing files by ID - Maintained duplicate file path handling to preserve original IDs - All tests passing --- ra_aid/tools/memory.py | 72 +++++++++++++++---- tests/ra_aid/tools/test_memory.py | 116 +++++++++++++++++++++++++----- 2 files changed, 156 insertions(+), 32 deletions(-) diff --git a/ra_aid/tools/memory.py b/ra_aid/tools/memory.py index efc1a1b..bc4afd6 100644 --- a/ra_aid/tools/memory.py +++ b/ra_aid/tools/memory.py @@ -27,7 +27,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': False, - 'related_files': set(), + 'related_files': {}, # Dict[int, str] - ID to filepath mapping + 'related_file_id_counter': 0, # Counter for generating unique file IDs 'plan_completed': False, 'research_depth': 0 } @@ -191,7 +192,7 @@ def emit_key_snippets(snippets: List[SnippetInfo]) -> str: List of stored snippet confirmation messages """ # First collect unique filepaths to add as related files - _global_memory['related_files'].update(snippet_info['filepath'] for snippet_info in snippets) + emit_related_files.invoke({"files": [snippet_info['filepath'] for snippet_info in snippets]}) results = [] for snippet_info in snippets: @@ -328,13 +329,14 @@ def plan_implementation_completed(message: str) -> str: console.print(Panel(Markdown(message), title="✅ Plan Executed")) return "Plan completion noted." -def get_related_files() -> Set[str]: - """Get the current set of related files. +def get_related_files() -> List[str]: + """Get the current list of related files. Returns: - Set of file paths that have been marked as related + List of formatted strings in the format 'ID#X path/to/file.py' """ - return _global_memory['related_files'] + files = _global_memory['related_files'] + return [f"ID#{file_id} {filepath}" for file_id, filepath in sorted(files.items())] @tool("emit_related_files") def emit_related_files(files: List[str]) -> str: @@ -344,27 +346,67 @@ def emit_related_files(files: List[str]) -> str: files: List of file paths to add Returns: - Confirmation message + Formatted string containing file IDs and paths for all processed files """ 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}") + # Process files + for file in files: + # Check if file path already exists in values + existing_id = None + for fid, fpath in _global_memory['related_files'].items(): + if fpath == file: + existing_id = fid + break + + if existing_id is not None: + # File exists, use existing ID + results.append(f"File ID #{existing_id}: {file}") + else: + # New file, assign new ID + file_id = _global_memory['related_file_id_counter'] + _global_memory['related_file_id_counter'] += 1 + + # Store file with ID + _global_memory['related_files'][file_id] = file + added_files.append((file_id, file)) + results.append(f"File ID #{file_id}: {file}") # Rich output - single consolidated panel if added_files: - files_added_md = '\n'.join(f"- `{file}`" for file in added_files) + files_added_md = '\n'.join(f"- ID#{id}: `{file}`" for id, 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." + return '\n'.join(results) + + +@tool("delete_related_files") +def delete_related_files(file_ids: List[int]) -> str: + """Delete multiple related files from global memory by their IDs. + Silently skips any IDs that don't exist. + + Args: + file_ids: List of file IDs to delete + + Returns: + Success message string + """ + results = [] + for file_id in file_ids: + if file_id in _global_memory['related_files']: + # Delete the file reference + deleted_file = _global_memory['related_files'].pop(file_id) + success_msg = f"Successfully removed related file #{file_id}: {deleted_file}" + console.print(Panel(Markdown(success_msg), + title="File Reference Removed", + border_style="green")) + results.append(success_msg) + + return "File references removed." def get_memory_value(key: str) -> str: """Get a value from global memory. diff --git a/tests/ra_aid/tools/test_memory.py b/tests/ra_aid/tools/test_memory.py index cb33ea4..e69d5ed 100644 --- a/tests/ra_aid/tools/test_memory.py +++ b/tests/ra_aid/tools/test_memory.py @@ -8,6 +8,7 @@ from ra_aid.tools.memory import ( delete_key_snippets, emit_related_files, get_related_files, + delete_related_files, emit_task, delete_tasks, swap_task_order @@ -23,7 +24,8 @@ def reset_memory(): _global_memory['research_notes'] = [] _global_memory['plans'] = [] _global_memory['tasks'] = [] - _global_memory['related_files'] = set() + _global_memory['related_files'] = {} + _global_memory['related_file_id_counter'] = 0 _global_memory['tasks'] = {} _global_memory['task_id_counter'] = 0 yield @@ -34,6 +36,8 @@ def reset_memory(): _global_memory['key_snippet_id_counter'] = 0 _global_memory['research_notes'] = [] _global_memory['plans'] = [] + _global_memory['related_files'] = {} + _global_memory['related_file_id_counter'] = 0 _global_memory['tasks'] = {} _global_memory['task_id_counter'] = 0 @@ -214,33 +218,104 @@ def test_delete_key_snippets_empty(reset_memory): assert 0 in _global_memory['key_snippets'] def test_emit_related_files_basic(reset_memory): - """Test basic adding of files""" + """Test basic adding of files with ID tracking""" # Test adding single file result = emit_related_files.invoke({"files": ["test.py"]}) - assert result == "Files noted." - assert get_related_files() == {"test.py"} + assert result == "File ID #0: test.py" + assert _global_memory['related_files'][0] == "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"} + assert result == "File ID #1: main.py\nFile ID #2: utils.py" + # Verify both files exist in related_files + values = list(_global_memory['related_files'].values()) + assert "main.py" in values + assert "utils.py" in values def test_get_related_files_empty(reset_memory): """Test getting related files when none added""" - assert get_related_files() == set() + assert get_related_files() == [] def test_emit_related_files_duplicates(reset_memory): - """Test that duplicate files are handled correctly""" + """Test that duplicate files return existing IDs with proper formatting""" # 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"} + assert result == "File ID #0: test.py\nFile ID #1: main.py" + first_id = 0 # ID of test.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"} + result = emit_related_files.invoke({"files": ["test.py"]}) + assert result == "File ID #0: test.py" # Should return same ID + assert len(_global_memory['related_files']) == 2 # Count should not increase + + # Try mix of new and duplicate files + result = emit_related_files.invoke({"files": ["test.py", "new.py"]}) + assert result == "File ID #0: test.py\nFile ID #2: new.py" + assert len(_global_memory['related_files']) == 3 + +def test_related_files_id_tracking(reset_memory): + """Test ID assignment and counter functionality for related files""" + # Add first file + result = emit_related_files.invoke({"files": ["file1.py"]}) + assert result == "File ID #0: file1.py" + assert _global_memory['related_file_id_counter'] == 1 + + # Add second file + result = emit_related_files.invoke({"files": ["file2.py"]}) + assert result == "File ID #1: file2.py" + assert _global_memory['related_file_id_counter'] == 2 + + # Verify all files stored correctly + assert _global_memory['related_files'][0] == "file1.py" + assert _global_memory['related_files'][1] == "file2.py" + +def test_delete_related_files(reset_memory): + """Test deleting related files""" + # Add test files + emit_related_files.invoke({"files": ["file1.py", "file2.py", "file3.py"]}) + + # Delete middle file + result = delete_related_files.invoke({"file_ids": [1]}) + assert result == "File references removed." + assert 1 not in _global_memory['related_files'] + assert len(_global_memory['related_files']) == 2 + + # Delete multiple files including non-existent ID + result = delete_related_files.invoke({"file_ids": [0, 2, 999]}) + assert result == "File references removed." + assert len(_global_memory['related_files']) == 0 + + # Counter should remain unchanged after deletions + assert _global_memory['related_file_id_counter'] == 3 + +def test_related_files_duplicates(reset_memory): + """Test duplicate file handling returns same ID""" + # Add initial file + result1 = emit_related_files.invoke({"files": ["test.py"]}) + assert result1 == "File ID #0: test.py" + + # Add same file again + result2 = emit_related_files.invoke({"files": ["test.py"]}) + assert result2 == "File ID #0: test.py" + + # Verify only one entry exists + assert len(_global_memory['related_files']) == 1 + assert _global_memory['related_file_id_counter'] == 1 + +def test_related_files_formatting(reset_memory): + """Test related files output string formatting""" + # Add some files + emit_related_files.invoke({"files": ["file1.py", "file2.py"]}) + + # Get formatted output + output = get_memory_value('related_files') + # Expect just the IDs on separate lines + expected = "0\n1" + assert output == expected + + # Test empty case + _global_memory['related_files'] = {} + assert get_memory_value('related_files') == "" def test_key_snippets_integration(reset_memory): """Integration test for key snippets functionality""" @@ -270,8 +345,13 @@ def test_key_snippets_integration(reset_memory): result = emit_key_snippets.invoke({"snippets": snippets}) assert result == "Snippets stored." assert _global_memory['key_snippet_id_counter'] == 3 - # Verify related files were tracked - assert _global_memory['related_files'] == {"file1.py", "file2.py", "file3.py"} + # Verify related files were tracked with IDs + assert len(_global_memory['related_files']) == 3 + # Check files are stored with proper IDs + file_values = _global_memory['related_files'].values() + assert "file1.py" in file_values + assert "file2.py" in file_values + assert "file3.py" in file_values # Verify all snippets were stored correctly assert len(_global_memory['key_snippets']) == 3 @@ -302,7 +382,9 @@ def test_key_snippets_integration(reset_memory): assert result == "Snippets stored." assert _global_memory['key_snippet_id_counter'] == 4 # Verify new file was added to related files - assert _global_memory['related_files'] == {"file1.py", "file2.py", "file3.py", "file4.py"} + file_values = _global_memory['related_files'].values() + assert "file4.py" in file_values + assert len(_global_memory['related_files']) == 4 # Delete remaining snippets result = delete_key_snippets.invoke({"snippet_ids": [1, 3]})