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
This commit is contained in:
user
parent
497c24c894
commit
80486ec533
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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]})
|
||||
|
|
|
|||
Loading…
Reference in New Issue