From cb8ee556e5b8ba50991df160a2300c92db850bfa Mon Sep 17 00:00:00 2001 From: AI Christianson Date: Wed, 11 Dec 2024 13:52:08 -0500 Subject: [PATCH] emit snippets in batches --- ra_aid/__main__.py | 8 +- ra_aid/prompts.py | 14 +-- ra_aid/tools/__init__.py | 8 +- ra_aid/tools/memory.py | 114 ++++++++++---------- tests/ra_aid/tools/test_memory.py | 169 +++++++++++++++++++++++++++++- 5 files changed, 243 insertions(+), 70 deletions(-) diff --git a/ra_aid/__main__.py b/ra_aid/__main__.py index 1a1e58e..7e0fe78 100644 --- 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) diff --git a/ra_aid/prompts.py b/ra_aid/prompts.py index 38e6d17..cd1285c 100644 --- 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. diff --git a/ra_aid/tools/__init__.py b/ra_aid/tools/__init__.py index 3818ea0..56d4267 100644 --- 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', - 'emit_expert_context', + 'delete_key_snippets', + 'emit_expert_context', 'emit_key_facts', - 'emit_key_snippet', + 'emit_key_snippets', 'emit_plan', 'emit_related_files', 'emit_research_notes', diff --git a/ra_aid/tools/memory.py b/ra_aid/tools/memory.py index 17a2efb..22e2bb2 100644 --- 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 - - # Store snippet info - snippet_info: SnippetInfo = { - 'filepath': filepath, - 'line_number': line_number, - 'snippet': snippet, - 'description': description - } - _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]) - - # 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 = [] + 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 + _global_memory['key_snippets'][snippet_id] = snippet_info + + # 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")) + + results.append(f"Stored snippet #{snippet_id}") + + return results -@tool("delete_key_snippet") -def delete_key_snippet(snippet_id: int) -> str: - """Delete a key snippet from global memory by its ID. +@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 - - # 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 + 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) + + return results def get_memory_value(key: str) -> str: """Get a value from global memory. diff --git a/tests/ra_aid/tools/test_memory.py b/tests/ra_aid/tools/test_memory.py index bc25a9d..bfcdf60 100644 --- 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