diff --git a/ra_aid/agents/key_snippets_gc_agent.py b/ra_aid/agents/key_snippets_gc_agent.py index 5956800..13bc124 100644 --- a/ra_aid/agents/key_snippets_gc_agent.py +++ b/ra_aid/agents/key_snippets_gc_agent.py @@ -17,13 +17,66 @@ from ra_aid.agent_utils import create_agent, run_agent_with_retry from ra_aid.database.repositories.key_snippet_repository import KeySnippetRepository from ra_aid.llm import initialize_llm from ra_aid.prompts.key_snippets_gc_prompts import KEY_SNIPPETS_GC_PROMPT -from ra_aid.tools.memory import delete_key_snippets, log_work_event, _global_memory +from ra_aid.tools.memory import log_work_event, _global_memory console = Console() key_snippet_repository = KeySnippetRepository() +@tool +def delete_key_snippets(snippet_ids: List[int]) -> str: + """Delete multiple key snippets from the database by their IDs. + Silently skips any IDs that don't exist. + + Args: + snippet_ids: List of snippet IDs to delete + + Returns: + str: Success or failure message + """ + results = [] + not_found_snippets = [] + failed_snippets = [] + + for snippet_id in snippet_ids: + # Get the snippet first to capture filepath for the message + snippet = key_snippet_repository.get(snippet_id) + if snippet: + filepath = snippet.filepath + # Delete from database + success = key_snippet_repository.delete(snippet_id) + if success: + success_msg = f"Successfully deleted snippet #{snippet_id} from {filepath}" + console.print( + Panel( + Markdown(success_msg), title="Snippet Deleted", border_style="green" + ) + ) + results.append((snippet_id, filepath)) + log_work_event(f"Deleted snippet {snippet_id}.") + else: + failed_snippets.append(snippet_id) + else: + not_found_snippets.append(snippet_id) + + # Prepare result message + result_parts = [] + if results: + deleted_msg = "Successfully deleted snippets:\n" + "\n".join([f"- #{snippet_id}: {filepath}" for snippet_id, filepath in results]) + result_parts.append(deleted_msg) + + if not_found_snippets: + not_found_msg = f"Snippets not found: {', '.join([f'#{snippet_id}' for snippet_id in not_found_snippets])}" + result_parts.append(not_found_msg) + + if failed_snippets: + failed_msg = f"Failed to delete snippets: {', '.join([f'#{snippet_id}' for snippet_id in failed_snippets])}" + result_parts.append(failed_msg) + + return "Snippets deleted." + + def run_key_snippets_gc_agent() -> None: """Run the key snippets gc agent to maintain a reasonable number of key snippets. diff --git a/ra_aid/tools/__init__.py b/ra_aid/tools/__init__.py index f793d53..0e07066 100644 --- a/ra_aid/tools/__init__.py +++ b/ra_aid/tools/__init__.py @@ -4,7 +4,6 @@ from .fuzzy_find import fuzzy_find_project_files from .human import ask_human from .list_directory import list_directory_tree from .memory import ( - delete_key_snippets, delete_tasks, deregister_related_files, emit_key_facts, @@ -29,7 +28,6 @@ from .write_file import put_complete_file_contents __all__ = [ "ask_expert", - "delete_key_snippets", "web_search_tavily", "deregister_related_files", "emit_expert_context", diff --git a/ra_aid/tools/memory.py b/ra_aid/tools/memory.py index d79d391..c2f1c88 100644 --- a/ra_aid/tools/memory.py +++ b/ra_aid/tools/memory.py @@ -256,33 +256,6 @@ def emit_key_snippet(snippet_info: SnippetInfo) -> str: return f"Snippet #{snippet_id} stored." -@tool("delete_key_snippets") -def delete_key_snippets(snippet_ids: List[int]) -> str: - """Delete multiple key snippets from the database by their IDs. - Silently skips any IDs that don't exist. - - Args: - snippet_ids: List of snippet IDs to delete - """ - results = [] - for snippet_id in snippet_ids: - # Get the snippet first to capture filepath for the message - snippet = key_snippet_repository.get(snippet_id) - if snippet: - filepath = snippet.filepath - # Delete from database - success = key_snippet_repository.delete(snippet_id) - if success: - success_msg = f"Successfully deleted snippet #{snippet_id} from {filepath}" - console.print( - Panel( - Markdown(success_msg), title="Snippet Deleted", border_style="green" - ) - ) - results.append(success_msg) - - log_work_event(f"Deleted snippets {snippet_ids}.") - return "Snippets deleted." @tool("swap_task_order") diff --git a/tests/ra_aid/tools/test_memory.py b/tests/ra_aid/tools/test_memory.py index 30e85be..ac388f7 100644 --- a/tests/ra_aid/tools/test_memory.py +++ b/tests/ra_aid/tools/test_memory.py @@ -4,9 +4,9 @@ import importlib import pytest from unittest.mock import patch, MagicMock +from ra_aid.agents.key_snippets_gc_agent import delete_key_snippets from ra_aid.tools.memory import ( _global_memory, - delete_key_snippets, delete_tasks, deregister_related_files, emit_key_facts, @@ -113,61 +113,66 @@ def mock_repository(): @pytest.fixture(autouse=True) def mock_key_snippet_repository(): """Mock the KeySnippetRepository to avoid database operations during tests""" - with patch('ra_aid.tools.memory.key_snippet_repository') as mock_repo: - # Setup the mock repository to behave like the original, but using memory - snippets = {} # Local in-memory storage - snippet_id_counter = 0 - - # Mock KeySnippet objects - class MockKeySnippet: - def __init__(self, id, filepath, line_number, snippet, description=None): - self.id = id - self.filepath = filepath - self.line_number = line_number - self.snippet = snippet - self.description = description + snippets = {} # Local in-memory storage + snippet_id_counter = 0 + + # Mock KeySnippet objects + class MockKeySnippet: + def __init__(self, id, filepath, line_number, snippet, description=None): + self.id = id + self.filepath = filepath + self.line_number = line_number + self.snippet = snippet + self.description = description - # Mock create method - def mock_create(filepath, line_number, snippet, description=None): - nonlocal snippet_id_counter - key_snippet = MockKeySnippet(snippet_id_counter, filepath, line_number, snippet, description) - snippets[snippet_id_counter] = key_snippet - snippet_id_counter += 1 - return key_snippet - mock_repo.create.side_effect = mock_create - - # Mock get method - def mock_get(snippet_id): - return snippets.get(snippet_id) - mock_repo.get.side_effect = mock_get - - # Mock delete method - def mock_delete(snippet_id): - if snippet_id in snippets: - del snippets[snippet_id] - return True - return False - mock_repo.delete.side_effect = mock_delete - - # Mock get_snippets_dict method - def mock_get_snippets_dict(): - return { - snippet_id: { - "filepath": snippet.filepath, - "line_number": snippet.line_number, - "snippet": snippet.snippet, - "description": snippet.description - } - for snippet_id, snippet in snippets.items() + # Mock create method + def mock_create(filepath, line_number, snippet, description=None): + nonlocal snippet_id_counter + key_snippet = MockKeySnippet(snippet_id_counter, filepath, line_number, snippet, description) + snippets[snippet_id_counter] = key_snippet + snippet_id_counter += 1 + return key_snippet + + # Mock get method + def mock_get(snippet_id): + return snippets.get(snippet_id) + + # Mock delete method + def mock_delete(snippet_id): + if snippet_id in snippets: + del snippets[snippet_id] + return True + return False + + # Mock get_snippets_dict method + def mock_get_snippets_dict(): + return { + snippet_id: { + "filepath": snippet.filepath, + "line_number": snippet.line_number, + "snippet": snippet.snippet, + "description": snippet.description } - mock_repo.get_snippets_dict.side_effect = mock_get_snippets_dict + for snippet_id, snippet in snippets.items() + } + + # Mock get_all method + def mock_get_all(): + return list(snippets.values()) + + # Create the actual mocks for both memory.py and key_snippets_gc_agent.py + with patch('ra_aid.tools.memory.key_snippet_repository') as memory_mock_repo, \ + patch('ra_aid.agents.key_snippets_gc_agent.key_snippet_repository') as agent_mock_repo: - # Mock get_all method - def mock_get_all(): - return list(snippets.values()) - mock_repo.get_all.side_effect = mock_get_all + # Setup both mocks with the same implementation + for mock_repo in [memory_mock_repo, agent_mock_repo]: + mock_repo.create.side_effect = mock_create + mock_repo.get.side_effect = mock_get + mock_repo.delete.side_effect = mock_delete + mock_repo.get_snippets_dict.side_effect = mock_get_snippets_dict + mock_repo.get_all.side_effect = mock_get_all - yield mock_repo + yield memory_mock_repo def test_emit_key_facts_single_fact(reset_memory, mock_repository): @@ -342,7 +347,8 @@ def test_emit_key_snippet(reset_memory, mock_key_snippet_repository): ) -def test_delete_key_snippets(reset_memory, mock_key_snippet_repository): +@patch('ra_aid.agents.key_snippets_gc_agent.log_work_event') +def test_delete_key_snippets(mock_log_work_event, reset_memory, mock_key_snippet_repository): """Test deleting multiple code snippets""" # Mock snippets snippets = [ @@ -373,24 +379,26 @@ def test_delete_key_snippets(reset_memory, mock_key_snippet_repository): mock_key_snippet_repository.reset_mock() # Test deleting mix of valid and invalid IDs - result = delete_key_snippets.invoke({"snippet_ids": [0, 1, 999]}) + with patch('ra_aid.agents.key_snippets_gc_agent.key_snippet_repository', mock_key_snippet_repository): + result = delete_key_snippets.invoke({"snippet_ids": [0, 1, 999]}) - # Verify success message - assert result == "Snippets deleted." + # Verify success message + assert result == "Snippets deleted." - # Verify repository delete was called with correct IDs - mock_key_snippet_repository.get.assert_any_call(0) - mock_key_snippet_repository.get.assert_any_call(1) - mock_key_snippet_repository.get.assert_any_call(999) - - mock_key_snippet_repository.delete.assert_any_call(0) - mock_key_snippet_repository.delete.assert_any_call(1) - - # Make sure delete wasn't called for ID 999 - assert mock_key_snippet_repository.delete.call_count == 2 + # Verify repository delete was called with correct IDs + mock_key_snippet_repository.get.assert_any_call(0) + mock_key_snippet_repository.get.assert_any_call(1) + mock_key_snippet_repository.get.assert_any_call(999) + + mock_key_snippet_repository.delete.assert_any_call(0) + mock_key_snippet_repository.delete.assert_any_call(1) + + # Make sure delete wasn't called for ID 999 + assert mock_key_snippet_repository.delete.call_count == 2 -def test_delete_key_snippets_empty(reset_memory, mock_key_snippet_repository): +@patch('ra_aid.agents.key_snippets_gc_agent.log_work_event') +def test_delete_key_snippets_empty(mock_log_work_event, reset_memory, mock_key_snippet_repository): """Test deleting snippets with empty ID list""" # Add a test snippet snippet = { @@ -405,11 +413,12 @@ def test_delete_key_snippets_empty(reset_memory, mock_key_snippet_repository): mock_key_snippet_repository.reset_mock() # Test with empty list - result = delete_key_snippets.invoke({"snippet_ids": []}) - assert result == "Snippets deleted." + with patch('ra_aid.agents.key_snippets_gc_agent.key_snippet_repository', mock_key_snippet_repository): + result = delete_key_snippets.invoke({"snippet_ids": []}) + assert result == "Snippets deleted." - # Verify no call to delete method - mock_key_snippet_repository.delete.assert_not_called() + # Verify no call to delete method + mock_key_snippet_repository.delete.assert_not_called() def test_emit_related_files_basic(reset_memory, tmp_path): @@ -613,7 +622,8 @@ def test_emit_related_files_path_normalization(reset_memory, tmp_path): os.chdir(original_dir) -def test_key_snippets_integration(reset_memory, tmp_path, mock_key_snippet_repository): +@patch('ra_aid.agents.key_snippets_gc_agent.log_work_event') +def test_key_snippets_integration(mock_log_work_event, reset_memory, tmp_path, mock_key_snippet_repository): """Integration test for key snippets functionality""" # Create test files file1 = tmp_path / "file1.py" @@ -665,13 +675,14 @@ def test_key_snippets_integration(reset_memory, tmp_path, mock_key_snippet_repos mock_key_snippet_repository.reset_mock() # Delete some but not all snippets (0 and 2) - result = delete_key_snippets.invoke({"snippet_ids": [0, 2]}) - assert result == "Snippets deleted." + with patch('ra_aid.agents.key_snippets_gc_agent.key_snippet_repository', mock_key_snippet_repository): + result = delete_key_snippets.invoke({"snippet_ids": [0, 2]}) + assert result == "Snippets deleted." - # Verify delete was called for the correct IDs - mock_key_snippet_repository.delete.assert_any_call(0) - mock_key_snippet_repository.delete.assert_any_call(2) - assert mock_key_snippet_repository.delete.call_count == 2 + # Verify delete was called for the correct IDs + mock_key_snippet_repository.delete.assert_any_call(0) + mock_key_snippet_repository.delete.assert_any_call(2) + assert mock_key_snippet_repository.delete.call_count == 2 # Reset mock again mock_key_snippet_repository.reset_mock() @@ -705,13 +716,14 @@ def test_key_snippets_integration(reset_memory, tmp_path, mock_key_snippet_repos mock_key_snippet_repository.reset_mock() # Delete remaining snippets - result = delete_key_snippets.invoke({"snippet_ids": [1, 3]}) - assert result == "Snippets deleted." + with patch('ra_aid.agents.key_snippets_gc_agent.key_snippet_repository', mock_key_snippet_repository): + result = delete_key_snippets.invoke({"snippet_ids": [1, 3]}) + assert result == "Snippets deleted." - # Verify delete was called for the correct IDs - mock_key_snippet_repository.delete.assert_any_call(1) - mock_key_snippet_repository.delete.assert_any_call(3) - assert mock_key_snippet_repository.delete.call_count == 2 + # Verify delete was called for the correct IDs + mock_key_snippet_repository.delete.assert_any_call(1) + mock_key_snippet_repository.delete.assert_any_call(3) + assert mock_key_snippet_repository.delete.call_count == 2 def test_emit_task_with_id(reset_memory):