emit snippets in batches

2024-12-11 13:52:08 -05:00 · 2024-12-11 13:52:08 -05:00 · cb8ee556e5
parent b02b125d06
commit cb8ee556e5
5 changed files with 243 additions and 70 deletions
--- a/ra_aid/main.py
+++ b/ra_aid/main.py
@ -11,7 +11,7 @@ from ra_aid.tools import (
    ask_expert, run_shell_command, run_programming_task,
    emit_research_notes, emit_plan, emit_related_files, emit_task,
    emit_expert_context, get_memory_value, emit_key_facts, delete_key_facts,
-    emit_key_snippet, delete_key_snippet,
+    emit_key_snippets, delete_key_snippets,
    request_implementation, read_file_tool, emit_research_subtask,
    fuzzy_find_project_files, ripgrep_search, list_directory_tree
 )
@ -62,9 +62,9 @@ planning_memory = MemorySaver()
 implementation_memory = MemorySaver()

 # Define tool sets for each stage
-research_tools = [list_directory_tree, emit_research_subtask, run_shell_command, emit_expert_context, ask_expert, emit_research_notes, emit_related_files, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, request_implementation, read_file_tool, fuzzy_find_project_files, ripgrep_search]
-planning_tools = [list_directory_tree, emit_expert_context, ask_expert, emit_plan, emit_task, emit_related_files, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, read_file_tool, fuzzy_find_project_files, ripgrep_search]
-implementation_tools = [list_directory_tree, run_shell_command, emit_expert_context, ask_expert, run_programming_task, emit_related_files, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, read_file_tool, fuzzy_find_project_files, ripgrep_search]
+research_tools = [list_directory_tree, emit_research_subtask, run_shell_command, emit_expert_context, ask_expert, emit_research_notes, emit_related_files, emit_key_facts, delete_key_facts, emit_key_snippets, delete_key_snippets, request_implementation, read_file_tool, fuzzy_find_project_files, ripgrep_search]
+planning_tools = [list_directory_tree, emit_expert_context, ask_expert, emit_plan, emit_task, emit_related_files, emit_key_facts, delete_key_facts, emit_key_snippets, delete_key_snippets, read_file_tool, fuzzy_find_project_files, ripgrep_search]
+implementation_tools = [list_directory_tree, run_shell_command, emit_expert_context, ask_expert, run_programming_task, emit_related_files, emit_key_facts, delete_key_facts, emit_key_snippets, delete_key_snippets, read_file_tool, fuzzy_find_project_files, ripgrep_search]

 # Create stage-specific agents with individual memory objects
 research_agent = create_react_agent(model, research_tools, checkpointer=research_memory)
--- a/ra_aid/prompts.py
+++ b/ra_aid/prompts.py
@ -77,6 +77,8 @@ Decision on Implementation
        If you see reasons that implementation changes will be required in the future, after documenting all findings, call request_implementation and specify why.
        If no changes are needed, simply state that no changes are required.

+Be thorough on locating all potential change sites/gauging blast radius.
+
 If there is a top-level README.md or docs/ folder, always start with that.
 """

@ -108,8 +110,8 @@ Snippet Management:
    Each snippet is identified with [Snippet ID: X].
    Snippets include file path, line number, and source code.
    Snippets may have optional descriptions explaining their significance.
-    Delete snippets with delete_key_snippet if they become outdated or irrelevant.
-    Use emit_key_snippet to store important code sections needed for reference.
+    Delete snippets with delete_key_snippets([id1, id2, ...]) to remove outdated or irrelevant ones.
+    Use emit_key_snippets to store important code sections needed for reference in batches.

 Guidelines:

@ -165,8 +167,8 @@ Snippet Management:
    Each snippet is identified with [Snippet ID: X].
    Snippets include file path, line number, and source code.
    Snippets may have optional descriptions explaining their significance.
-    Delete snippets with delete_key_snippet if they become outdated or irrelevant.
-    Use emit_key_snippet to store important code sections needed for reference.
+    Delete snippets with delete_key_snippets([id1, id2, ...]) to remove outdated or irrelevant ones.
+    Use emit_key_snippets to store important code sections needed for reference in batches.

 Instructions:
 - **Stay Within Provided Information**: Do not include any information not present in the Research Notes or Key Facts.
@ -196,8 +198,8 @@ Important Notes:
 - Focus solely on the given task and implement it as described.
 - Scale the complexity of your solution to the complexity of the request. For simple requests, keep it straightforward and minimal. For complex requests, maintain the previously planned depth.
 - Use delete_key_facts to remove facts that become outdated, irrelevant, or duplicated.
- Use emit_key_snippet to manage code sections before and after modifications as needed.
- Regularly remove outdated snippets with delete_key_snippet.
+- Use emit_key_snippets to manage code sections before and after modifications in batches.
+- Regularly remove outdated snippets with delete_key_snippets.

 Instructions:
 1. Review the provided base task, plan, and key facts.
--- a/ra_aid/tools/init.py
+++ b/ra_aid/tools/init.py
@ -8,16 +8,16 @@ from .ripgrep import ripgrep_search
 from .memory import (
    emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_facts,
    request_implementation, skip_implementation, delete_key_facts, emit_research_subtask,
-    emit_key_snippet, delete_key_snippet
+    emit_key_snippets, delete_key_snippets
 )

 __all__ = [
    'ask_expert',
    'delete_key_facts',
-    'delete_key_snippet', 
+    'delete_key_snippets',
    'emit_expert_context', 
    'emit_key_facts',
-    'emit_key_snippet',
+    'emit_key_snippets',
    'emit_plan',
    'emit_related_files',
    'emit_research_notes',
--- a/ra_aid/tools/memory.py
+++ b/ra_aid/tools/memory.py
@ -168,71 +168,75 @@ def skip_implementation(reason: str) -> str:
    console.print(Panel(Markdown(reason), title="⏭️ Implementation Skipped"))
    return reason

-@tool("emit_key_snippet")
-def emit_key_snippet(filepath: str, line_number: int, snippet: str, description: Optional[str] = None) -> str:
-    """Store a key source code snippet in global memory.
+@tool("emit_key_snippets")
+def emit_key_snippets(snippets: List[SnippetInfo]) -> List[str]:
+    """Store multiple key source code snippets in global memory.
    
    Args:
-        filepath: Path to the source file
-        line_number: Line number where the snippet starts
-        snippet: The source code snippet text
-        description: Optional description of the snippet's significance
+        snippets: List of snippet information dictionaries containing:
+                 - filepath: Path to the source file
+                 - line_number: Line number where the snippet starts  
+                 - snippet: The source code snippet text
+                 - description: Optional description of the significance
                 
    Returns:
-        The stored snippet information
+        List of stored snippet confirmation messages
    """
-    # Get and increment snippet ID
-    snippet_id = _global_memory['key_snippet_id_counter']
-    _global_memory['key_snippet_id_counter'] += 1
+    results = []
+    for snippet_info in snippets:
+        # Get and increment snippet ID 
+        snippet_id = _global_memory['key_snippet_id_counter']
+        _global_memory['key_snippet_id_counter'] += 1
        
-    # Store snippet info
-    snippet_info: SnippetInfo = {
-        'filepath': filepath,
-        'line_number': line_number,
-        'snippet': snippet,
-        'description': description
-    }
-    _global_memory['key_snippets'][snippet_id] = snippet_info
+        # Store snippet info
+        _global_memory['key_snippets'][snippet_id] = snippet_info
        
-    # Format display text as markdown
-    display_text = [
-        f"**Source Location**:",
-        f"- File: `{filepath}`",
-        f"- Line: `{line_number}`",
-        "",  # Empty line before code block
-        "**Code**:",
-        "```python",
-        snippet.rstrip(),  # Remove trailing whitespace
-        "```"
-    ]
-    if description:
-        display_text.extend(["", "**Description**:", description])
+        # Format display text as markdown
+        display_text = [
+            f"**Source Location**:",
+            f"- File: `{snippet_info['filepath']}`",
+            f"- Line: `{snippet_info['line_number']}`",
+            "",  # Empty line before code block
+            "**Code**:",
+            "```python",
+            snippet_info['snippet'].rstrip(),  # Remove trailing whitespace 
+            "```"
+        ]
+        if snippet_info['description']:
+            display_text.extend(["", "**Description**:", snippet_info['description']])
            
-    # Display panel
-    console.print(Panel(Markdown("\n".join(display_text)), title=f"📝 Key Snippet #{snippet_id}", border_style="bright_cyan"))
+        # Display panel
+        console.print(Panel(Markdown("\n".join(display_text)), 
+                          title=f"📝 Key Snippet #{snippet_id}", 
+                          border_style="bright_cyan"))
        
-    return f"Stored snippet #{snippet_id}"
+        results.append(f"Stored snippet #{snippet_id}")
        
-@tool("delete_key_snippet")
-def delete_key_snippet(snippet_id: int) -> str:
-    """Delete a key snippet from global memory by its ID.
+    return results
+
+@tool("delete_key_snippets") 
+def delete_key_snippets(snippet_ids: List[int]) -> List[str]:
+    """Delete multiple key snippets from global memory by their IDs.
+    Silently skips any IDs that don't exist.
    
    Args:
-        snippet_id: The ID of the snippet to delete
+        snippet_ids: List of snippet IDs to delete
        
    Returns:
-        A message indicating success or failure
+        List of success messages for deleted snippets
    """
-    if snippet_id not in _global_memory['key_snippets']:
-        error_msg = f"Error: No snippet found with ID #{snippet_id}"
-        console.print(Panel(Markdown(error_msg), title="❌ Delete Failed", border_style="red"))
-        return error_msg
+    results = []
+    for snippet_id in snippet_ids:
+        if snippet_id in _global_memory['key_snippets']:
+            # Delete the snippet
+            deleted_snippet = _global_memory['key_snippets'].pop(snippet_id)
+            success_msg = f"Successfully deleted snippet #{snippet_id} from {deleted_snippet['filepath']}"
+            console.print(Panel(Markdown(success_msg), 
+                              title="🗑️ Snippet Deleted", 
+                              border_style="green"))
+            results.append(success_msg)
            
-    # Delete the snippet
-    deleted_snippet = _global_memory['key_snippets'].pop(snippet_id)
-    success_msg = f"Successfully deleted snippet #{snippet_id} from {deleted_snippet['filepath']}"
-    console.print(Panel(Markdown(success_msg), title="🗑️ Snippet Deleted", border_style="green"))
-    return success_msg
+    return results

 def get_memory_value(key: str) -> str:
    """Get a value from global memory.
--- a/tests/ra_aid/tools/test_memory.py
+++ b/tests/ra_aid/tools/test_memory.py
@ -4,7 +4,9 @@ from ra_aid.tools.memory import (
    get_memory_value,
    emit_research_subtask,
    emit_key_facts,
-    delete_key_facts
+    delete_key_facts,
+    emit_key_snippets,
+    delete_key_snippets
 )

@pytest.fixture
@ -12,6 +14,8 @@ def reset_memory():
    """Reset global memory before each test"""
    _global_memory['key_facts'] = {}
    _global_memory['key_fact_id_counter'] = 0
+    _global_memory['key_snippets'] = {}
+    _global_memory['key_snippet_id_counter'] = 0
    _global_memory['research_notes'] = []
    _global_memory['plans'] = []
    _global_memory['tasks'] = []
@ -20,6 +24,8 @@ def reset_memory():
    # Clean up after test
    _global_memory['key_facts'] = {}
    _global_memory['key_fact_id_counter'] = 0
+    _global_memory['key_snippets'] = {}
+    _global_memory['key_snippet_id_counter'] = 0
    _global_memory['research_notes'] = []
    _global_memory['plans'] = []
    _global_memory['tasks'] = []
@ -122,6 +128,167 @@ def test_delete_key_facts(reset_memory):
    assert 2 in _global_memory['key_facts']  # ID 2 should remain
    assert _global_memory['key_facts'][2] == "Third fact"

+def test_emit_key_snippets(reset_memory):
+    """Test emitting multiple code snippets at once"""
+    # Test snippets with and without descriptions
+    snippets = [
+        {
+            "filepath": "test.py",
+            "line_number": 10,
+            "snippet": "def test():\n    pass",
+            "description": "Test function"
+        },
+        {
+            "filepath": "main.py",
+            "line_number": 20,
+            "snippet": "print('hello')",
+            "description": None
+        }
+    ]
+    
+    # Emit snippets
+    results = emit_key_snippets.invoke({"snippets": snippets})
+    
+    # Verify return messages
+    assert results == ["Stored snippet #0", "Stored snippet #1"]
+    
+    # Verify snippets stored correctly
+    assert _global_memory['key_snippets'][0] == snippets[0]
+    assert _global_memory['key_snippets'][1] == snippets[1]
+    
+    # Verify counter incremented correctly
+    assert _global_memory['key_snippet_id_counter'] == 2
+
+def test_delete_key_snippets(reset_memory):
+    """Test deleting multiple code snippets"""
+    # Add test snippets
+    snippets = [
+        {
+            "filepath": "test1.py",
+            "line_number": 1,
+            "snippet": "code1",
+            "description": None
+        },
+        {
+            "filepath": "test2.py",
+            "line_number": 2,
+            "snippet": "code2",
+            "description": None
+        },
+        {
+            "filepath": "test3.py",
+            "line_number": 3,
+            "snippet": "code3",
+            "description": None
+        }
+    ]
+    emit_key_snippets.invoke({"snippets": snippets})
+    
+    # Test deleting mix of valid and invalid IDs
+    results = delete_key_snippets.invoke({"snippet_ids": [0, 1, 999]})
+    
+    # Verify success messages
+    assert results == [
+        "Successfully deleted snippet #0 from test1.py",
+        "Successfully deleted snippet #1 from test2.py"
+    ]
+    
+    # Verify correct snippets removed
+    assert 0 not in _global_memory['key_snippets']
+    assert 1 not in _global_memory['key_snippets']
+    assert 2 in _global_memory['key_snippets']
+    assert _global_memory['key_snippets'][2]['filepath'] == "test3.py"
+
+def test_delete_key_snippets_empty(reset_memory):
+    """Test deleting snippets with empty ID list"""
+    # Add a test snippet
+    snippet = {
+        "filepath": "test.py",
+        "line_number": 1,
+        "snippet": "code",
+        "description": None
+    }
+    emit_key_snippets.invoke({"snippets": [snippet]})
+    
+    # Test with empty list
+    results = delete_key_snippets.invoke({"snippet_ids": []})
+    assert results == []
+    
+    # Verify snippet still exists
+    assert 0 in _global_memory['key_snippets']
+
+def test_key_snippets_integration(reset_memory):
+    """Integration test for key snippets functionality"""
+    # Initial snippets to add
+    snippets = [
+        {
+            "filepath": "file1.py",
+            "line_number": 10,
+            "snippet": "def func1():\n    pass",
+            "description": "First function"
+        },
+        {
+            "filepath": "file2.py",
+            "line_number": 20,
+            "snippet": "def func2():\n    return True",
+            "description": "Second function"
+        },
+        {
+            "filepath": "file3.py",
+            "line_number": 30,
+            "snippet": "class TestClass:\n    pass",
+            "description": "Test class"
+        }
+    ]
+    
+    # Add all snippets
+    results = emit_key_snippets.invoke({"snippets": snippets})
+    assert results == ["Stored snippet #0", "Stored snippet #1", "Stored snippet #2"]
+    assert _global_memory['key_snippet_id_counter'] == 3
+    
+    # Verify all snippets were stored correctly
+    assert len(_global_memory['key_snippets']) == 3
+    assert _global_memory['key_snippets'][0] == snippets[0]
+    assert _global_memory['key_snippets'][1] == snippets[1]
+    assert _global_memory['key_snippets'][2] == snippets[2]
+    
+    # Delete some but not all snippets (0 and 2)
+    results = delete_key_snippets.invoke({"snippet_ids": [0, 2]})
+    assert len(results) == 2
+    assert "Successfully deleted snippet #0 from file1.py" in results
+    assert "Successfully deleted snippet #2 from file3.py" in results
+    
+    # Verify remaining snippet is intact
+    assert len(_global_memory['key_snippets']) == 1
+    assert 1 in _global_memory['key_snippets']
+    assert _global_memory['key_snippets'][1] == snippets[1]
+    
+    # Counter should remain unchanged after deletions
+    assert _global_memory['key_snippet_id_counter'] == 3
+    
+    # Add new snippet to verify counter continues correctly
+    new_snippet = {
+        "filepath": "file4.py",
+        "line_number": 40,
+        "snippet": "def func4():\n    return False",
+        "description": "Fourth function"
+    }
+    results = emit_key_snippets.invoke({"snippets": [new_snippet]})
+    assert results == ["Stored snippet #3"]
+    assert _global_memory['key_snippet_id_counter'] == 4
+    
+    # Delete remaining snippets
+    results = delete_key_snippets.invoke({"snippet_ids": [1, 3]})
+    assert len(results) == 2
+    assert "Successfully deleted snippet #1 from file2.py" in results
+    assert "Successfully deleted snippet #3 from file4.py" in results
+    
+    # Verify all snippets are gone
+    assert len(_global_memory['key_snippets']) == 0
+    
+    # Counter should still maintain its value
+    assert _global_memory['key_snippet_id_counter'] == 4
+
 def test_emit_research_subtask(reset_memory):
    """Test emitting research subtasks"""
    # Test adding a research subtask