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,8 +113,6 @@ 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
|
||||
|
||||
|
|
@ -134,12 +132,10 @@ def mock_key_snippet_repository():
|
|||
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):
|
||||
|
|
@ -147,7 +143,6 @@ def mock_key_snippet_repository():
|
|||
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():
|
||||
|
|
@ -160,14 +155,24 @@ def mock_key_snippet_repository():
|
|||
}
|
||||
for snippet_id, snippet in snippets.items()
|
||||
}
|
||||
mock_repo.get_snippets_dict.side_effect = mock_get_snippets_dict
|
||||
|
||||
# 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:
|
||||
|
||||
# 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,6 +379,7 @@ 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
|
||||
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
|
||||
|
|
@ -390,7 +397,8 @@ def test_delete_key_snippets(reset_memory, mock_key_snippet_repository):
|
|||
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,6 +413,7 @@ def test_delete_key_snippets_empty(reset_memory, mock_key_snippet_repository):
|
|||
mock_key_snippet_repository.reset_mock()
|
||||
|
||||
# Test with empty list
|
||||
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."
|
||||
|
||||
|
|
@ -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,6 +675,7 @@ 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)
|
||||
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."
|
||||
|
||||
|
|
@ -705,6 +716,7 @@ def test_key_snippets_integration(reset_memory, tmp_path, mock_key_snippet_repos
|
|||
mock_key_snippet_repository.reset_mock()
|
||||
|
||||
# Delete remaining snippets
|
||||
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."
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue