localize delete_key_snippets with gc agent
This commit is contained in:
AI Christianson
parent
332cbec826
commit
d4353b1824
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in New Issue