localize delete_key_snippets with gc agent

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.
     

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",

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")

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:
+    snippets = {}  # Local in-memory storage
-        # Setup the mock repository to behave like the original, but using memory
+    snippet_id_counter = 0
-        snippets = {}  # Local in-memory storage
+    
-        snippet_id_counter = 0
+    # Mock KeySnippet objects
-        
+    class MockKeySnippet:
-        # Mock KeySnippet objects
+        def __init__(self, id, filepath, line_number, snippet, description=None):
-        class MockKeySnippet:
+            self.id = id
-            def __init__(self, id, filepath, line_number, snippet, description=None):
+            self.filepath = filepath
-                self.id = id
+            self.line_number = line_number
-                self.filepath = filepath
+            self.snippet = snippet
-                self.line_number = line_number
+            self.description = description
-                self.snippet = snippet
-                self.description = description
 
-        # Mock create method
+    # Mock create method
-        def mock_create(filepath, line_number, snippet, description=None):
+    def mock_create(filepath, line_number, snippet, description=None):
-            nonlocal snippet_id_counter
+        nonlocal snippet_id_counter
-            key_snippet = MockKeySnippet(snippet_id_counter, filepath, line_number, snippet, description)
+        key_snippet = MockKeySnippet(snippet_id_counter, filepath, line_number, snippet, description)
-            snippets[snippet_id_counter] = key_snippet
+        snippets[snippet_id_counter] = key_snippet
-            snippet_id_counter += 1
+        snippet_id_counter += 1
-            return key_snippet
+        return key_snippet
-        mock_repo.create.side_effect = mock_create
+    
-        
+    # Mock get method
-        # Mock get method
+    def mock_get(snippet_id):
-        def mock_get(snippet_id):
+        return snippets.get(snippet_id)
-            return snippets.get(snippet_id)
+    
-        mock_repo.get.side_effect = mock_get
+    # Mock delete method
-        
+    def mock_delete(snippet_id):
-        # Mock delete method
+        if snippet_id in snippets:
-        def mock_delete(snippet_id):
+            del snippets[snippet_id]
-            if snippet_id in snippets:
+            return True
-                del snippets[snippet_id]
+        return False
-                return True
+    
-            return False
+    # Mock get_snippets_dict method
-        mock_repo.delete.side_effect = mock_delete
+    def mock_get_snippets_dict():
-        
+        return {
-        # Mock get_snippets_dict method
+            snippet_id: {
-        def mock_get_snippets_dict():
+                "filepath": snippet.filepath,
-            return {
+                "line_number": snippet.line_number,
-                snippet_id: {
+                "snippet": snippet.snippet,
-                    "filepath": snippet.filepath,
+                "description": snippet.description
-                    "line_number": snippet.line_number,
-                    "snippet": snippet.snippet,
-                    "description": snippet.description
-                }
-                for snippet_id, snippet in snippets.items()
             }
-        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
+        # Setup both mocks with the same implementation
-        def mock_get_all():
+        for mock_repo in [memory_mock_repo, agent_mock_repo]:
-            return list(snippets.values())
+            mock_repo.create.side_effect = mock_create
-        mock_repo.get_all.side_effect = mock_get_all
+            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
+        # Verify success message
-    assert result == "Snippets deleted."
+        assert result == "Snippets deleted."
 
-    # Verify repository delete was called with correct IDs
+        # 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(0)
-    mock_key_snippet_repository.get.assert_any_call(1)
+        mock_key_snippet_repository.get.assert_any_call(1)
-    mock_key_snippet_repository.get.assert_any_call(999)
+        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(0)
-    mock_key_snippet_repository.delete.assert_any_call(1)
+        mock_key_snippet_repository.delete.assert_any_call(1)
-    
+        
-    # Make sure delete wasn't called for ID 999
+        # Make sure delete wasn't called for ID 999
-    assert mock_key_snippet_repository.delete.call_count == 2
+        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": []})
+    with patch('ra_aid.agents.key_snippets_gc_agent.key_snippet_repository', mock_key_snippet_repository):
-    assert result == "Snippets deleted."
+        result = delete_key_snippets.invoke({"snippet_ids": []})
+        assert result == "Snippets deleted."
 
-    # Verify no call to delete method
+        # Verify no call to delete method
-    mock_key_snippet_repository.delete.assert_not_called()
+        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]})
+    with patch('ra_aid.agents.key_snippets_gc_agent.key_snippet_repository', mock_key_snippet_repository):
-    assert result == "Snippets deleted."
+        result = delete_key_snippets.invoke({"snippet_ids": [0, 2]})
+        assert result == "Snippets deleted."
 
-    # Verify delete was called for the correct IDs
+        # 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(0)
-    mock_key_snippet_repository.delete.assert_any_call(2)
+        mock_key_snippet_repository.delete.assert_any_call(2)
-    assert mock_key_snippet_repository.delete.call_count == 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]})
+    with patch('ra_aid.agents.key_snippets_gc_agent.key_snippet_repository', mock_key_snippet_repository):
-    assert result == "Snippets deleted."
+        result = delete_key_snippets.invoke({"snippet_ids": [1, 3]})
+        assert result == "Snippets deleted."
 
-    # Verify delete was called for the correct IDs
+        # 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(1)
-    mock_key_snippet_repository.delete.assert_any_call(3)
+        mock_key_snippet_repository.delete.assert_any_call(3)
-    assert mock_key_snippet_repository.delete.call_count == 2
+        assert mock_key_snippet_repository.delete.call_count == 2
 
 
 def test_emit_task_with_id(reset_memory):