From 3ac2c3c66d062e251224f0cd9522ff59a969d4eb Mon Sep 17 00:00:00 2001 From: AI Christianson Date: Wed, 11 Dec 2024 12:35:22 -0500 Subject: [PATCH] allow emitting facts in batches --- .gitignore | 1 + ra_aid/__main__.py | 11 +++--- ra_aid/prompts.py | 18 +++++---- ra_aid/tools/__init__.py | 10 ++--- ra_aid/tools/memory.py | 51 ------------------------ ra_aid/tools/programmer.py | 27 +------------ tests/ra_aid/tools/test_memory.py | 65 +++++++++++++------------------ 7 files changed, 51 insertions(+), 132 deletions(-) diff --git a/.gitignore b/.gitignore index d091c0a..6873fb1 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ __pycache__/ .aider* .env /work +/dist diff --git a/ra_aid/__main__.py b/ra_aid/__main__.py index 9bbad77..1d4c1dd 100644 --- a/ra_aid/__main__.py +++ b/ra_aid/__main__.py @@ -10,7 +10,7 @@ from langgraph.prebuilt import create_react_agent from ra_aid.tools import ( ask_expert, run_shell_command, run_programming_task, emit_research_notes, emit_plan, emit_related_file, emit_task, - emit_expert_context, get_memory_value, emit_key_fact, delete_key_fact, + emit_expert_context, get_memory_value, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, 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_file, emit_key_fact, delete_key_fact, 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_file, emit_key_fact, delete_key_fact, 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_file, emit_key_fact, delete_key_fact, 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_file, 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_file, 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_file, emit_key_facts, delete_key_facts, emit_key_snippet, delete_key_snippet, 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) @@ -285,7 +285,8 @@ Be very thorough in your research and emit lots of snippets, key facts. If you t research_notes=get_memory_value('research_notes'), key_facts=get_memory_value('key_facts'), key_snippets=get_memory_value('key_snippets'), - base_task=base_task + base_task=base_task, + related_files="\n".join(related_files) ) # Run planning agent diff --git a/ra_aid/prompts.py b/ra_aid/prompts.py index dc969ab..3659625 100644 --- a/ra_aid/prompts.py +++ b/ra_aid/prompts.py @@ -90,6 +90,9 @@ Research Notes: {research_notes} +Relevant Files: +{related_files} + Key Facts: {key_facts} @@ -98,8 +101,8 @@ Key Snippets: Fact Management: Each fact is identified with [Fact ID: X]. - Facts may be deleted if they become outdated, irrelevant, or duplicates. - Use delete_key_fact with the specific Fact ID to remove unnecessary facts. + Facts may be deleted if they become outdated, irrelevant, or duplicates. + Use delete_key_facts([id1, id2, ...]) with a list of numeric Fact IDs to remove unnecessary facts. Snippet Management: Each snippet is identified with [Snippet ID: X]. @@ -155,8 +158,8 @@ Key Snippets: Fact Management: Each fact is identified with [Fact ID: X]. - Facts may be deleted if they become outdated, irrelevant, or duplicates. - Use delete_key_fact with the specific Fact ID to remove unnecessary facts. + Facts may be deleted if they become outdated, irrelevant, or duplicates. + Use delete_key_facts([id1, id2, ...]) with a list of numeric Fact IDs to remove unnecessary facts. Snippet Management: Each snippet is identified with [Snippet ID: X]. @@ -192,7 +195,7 @@ Relevant Files: 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_fact to remove facts that become outdated, irrelevant, or duplicated. +- 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. @@ -202,9 +205,10 @@ Instructions: {task} 3. Work incrementally, validating as you go. -4. Update or remove any key facts that no longer apply. +4. Use delete_key_facts to remove any key facts that no longer apply. 5. Do not add features not explicitly required. 6. Only create or modify files directly related to this task. +7. For trivial changes, use sed and awk judiciously via the run_shell_command tool. Once the task is complete, ensure all updated files are emitted. -""" \ No newline at end of file +""" diff --git a/ra_aid/tools/__init__.py b/ra_aid/tools/__init__.py index 3df863a..f723fd0 100644 --- a/ra_aid/tools/__init__.py +++ b/ra_aid/tools/__init__.py @@ -6,17 +6,17 @@ from .fuzzy_find import fuzzy_find_project_files from .list_directory import list_directory_tree from .ripgrep import ripgrep_search from .memory import ( - emit_research_notes, emit_plan, emit_task, get_memory_value, emit_key_fact, - request_implementation, skip_implementation, delete_key_fact, emit_research_subtask, + 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 ) __all__ = [ 'ask_expert', - 'delete_key_fact', - 'delete_key_snippet', + 'delete_key_facts', + 'delete_key_snippet', 'emit_expert_context', - 'emit_key_fact', + 'emit_key_facts', 'emit_key_snippet', 'emit_plan', 'emit_related_file', diff --git a/ra_aid/tools/memory.py b/ra_aid/tools/memory.py index 3d2cd8b..17a2efb 100644 --- a/ra_aid/tools/memory.py +++ b/ra_aid/tools/memory.py @@ -86,37 +86,6 @@ def emit_research_subtask(subtask: str) -> str: console.print(Panel(Markdown(subtask), title="🔬 Research Subtask")) return f"Added research subtask: {subtask}" -@tool("emit_key_fact") -def emit_key_fact(fact: str) -> str: - """Store a key fact about the project or current task in global memory. - - Key facts are things like: - - Specific files/functions to look at and what they do - - Coding conventions - - Specific external interfaces related to the task - - Key facts should be objective and not restating things already specified in our top-level task. - - They are generally things that will not change throughout the duration of our top-level task. - - Args: - fact: The key fact to store - - Returns: - The stored fact - """ - # Get and increment fact ID - fact_id = _global_memory['key_fact_id_counter'] - _global_memory['key_fact_id_counter'] += 1 - - # Store fact with ID - _global_memory['key_facts'][fact_id] = fact - - # Display panel with ID - console.print(Panel(Markdown(fact), title=f"💡 Key Fact #{fact_id}", border_style="bright_cyan")) - - # Return fact with ID - return f"Stored fact #{fact_id}: {fact}" @tool("emit_key_facts") def emit_key_facts(facts: List[str]) -> List[str]: @@ -145,26 +114,6 @@ def emit_key_facts(facts: List[str]) -> List[str]: return results -@tool("delete_key_fact") -def delete_key_fact(fact_id: int) -> str: - """Delete a key fact from global memory by its ID. - - Args: - fact_id: The ID of the fact to delete - - Returns: - A message indicating success or failure - """ - if fact_id not in _global_memory['key_facts']: - error_msg = f"Error: No fact found with ID #{fact_id}" - console.print(Panel(Markdown(error_msg), title="❌ Delete Failed", border_style="red")) - return error_msg - - # Delete the fact - deleted_fact = _global_memory['key_facts'].pop(fact_id) - success_msg = f"Successfully deleted fact #{fact_id}: {deleted_fact}" - console.print(Panel(Markdown(success_msg), title="🗑️ Fact Deleted", border_style="green")) - return success_msg @tool("delete_key_facts") def delete_key_facts(fact_ids: List[int]) -> List[str]: diff --git a/ra_aid/tools/programmer.py b/ra_aid/tools/programmer.py index 30918a6..4d7948a 100644 --- a/ra_aid/tools/programmer.py +++ b/ra_aid/tools/programmer.py @@ -75,30 +75,7 @@ def run_programming_task(input: RunProgrammingTaskInput) -> Dict[str, Union[str, "-m" ] - # Inject key facts into instructions if they exist - key_facts = get_memory_value('key_facts') - enhanced_instructions = input.instructions - - # Get and format snippets if they exist - key_snippets = get_memory_value('key_snippets') - - # Combine all sections - enhanced_instructions = f"""Key Facts About This Project: - -{key_facts} - -Key Code Snippets: - -{key_snippets} - -Instructions: - -{input.instructions} - -Only implement the immediate instructions, do not expand scope. -""" - - command.append(enhanced_instructions) + command.append(input.instructions) # Use both input files and related files files_to_use = set(related_files) # Start with related files @@ -111,7 +88,7 @@ Only implement the immediate instructions, do not expand scope. # Create a pretty display of what we're doing task_display = [ "## Instructions\n", - f"{enhanced_instructions}\n" + f"{input.instructions}\n" ] if files_to_use: diff --git a/tests/ra_aid/tools/test_memory.py b/tests/ra_aid/tools/test_memory.py index db231e4..bc25a9d 100644 --- a/tests/ra_aid/tools/test_memory.py +++ b/tests/ra_aid/tools/test_memory.py @@ -1,8 +1,6 @@ import pytest from ra_aid.tools.memory import ( _global_memory, - emit_key_fact, - delete_key_fact, get_memory_value, emit_research_subtask, emit_key_facts, @@ -27,43 +25,35 @@ def reset_memory(): _global_memory['tasks'] = [] _global_memory['research_subtasks'] = [] -def test_emit_key_fact(reset_memory): - """Test emitting key facts with ID assignment""" - # First fact should get ID 0 - result = emit_key_fact("First fact") - assert result == "Stored fact #0: First fact" +def test_emit_key_facts_single_fact(reset_memory): + """Test emitting a single key fact using emit_key_facts""" + # Test with single fact + result = emit_key_facts.invoke({"facts": ["First fact"]}) + assert result[0] == "Stored fact #0: First fact" assert _global_memory['key_facts'][0] == "First fact" - - # Second fact should get ID 1 - result = emit_key_fact("Second fact") - assert result == "Stored fact #1: Second fact" - assert _global_memory['key_facts'][1] == "Second fact" - - # Counter should be at 2 - assert _global_memory['key_fact_id_counter'] == 2 + assert _global_memory['key_fact_id_counter'] == 1 -def test_delete_key_fact(reset_memory): - """Test deleting key facts""" - # Add some facts - emit_key_fact("First fact") - emit_key_fact("Second fact") +def test_delete_key_facts_single_fact(reset_memory): + """Test deleting a single key fact using delete_key_facts""" + # Add a fact + emit_key_facts.invoke({"facts": ["Test fact"]}) - # Delete fact #0 - result = delete_key_fact({'fact_id': 0}) - assert result == "Successfully deleted fact #0: First fact" + # Delete the fact + result = delete_key_facts.invoke({"fact_ids": [0]}) + assert result[0] == "Successfully deleted fact #0: Test fact" assert 0 not in _global_memory['key_facts'] - assert 1 in _global_memory['key_facts'] -def test_delete_invalid_fact(reset_memory): - """Test error handling when deleting non-existent facts""" - result = delete_key_fact({'fact_id': 999}) - assert result == "Error: No fact found with ID #999" +def test_delete_key_facts_invalid(reset_memory): + """Test deleting non-existent facts returns empty list""" + # Try to delete non-existent fact + result = delete_key_facts.invoke({"fact_ids": [999]}) + assert result == [] # Add and delete a fact, then try to delete it again - emit_key_fact("Test fact") - delete_key_fact({'fact_id': 0}) - result = delete_key_fact({'fact_id': 0}) - assert result == "Error: No fact found with ID #0" + emit_key_facts.invoke({"facts": ["Test fact"]}) + delete_key_facts.invoke({"fact_ids": [0]}) + result = delete_key_facts.invoke({"fact_ids": [0]}) + assert result == [] def test_get_memory_value_key_facts(reset_memory): """Test get_memory_value with key facts dictionary""" @@ -71,8 +61,7 @@ def test_get_memory_value_key_facts(reset_memory): assert get_memory_value('key_facts') == "" # Add some facts - emit_key_fact("First fact") - emit_key_fact("Second fact") + emit_key_facts.invoke({"facts": ["First fact", "Second fact"]}) # Should return markdown formatted list expected = "## 🔑 Key Fact #0\n\nFirst fact\n\n## 🔑 Key Fact #1\n\nSecond fact" @@ -96,7 +85,7 @@ def test_emit_key_facts(reset_memory): """Test emitting multiple key facts at once""" # Test emitting multiple facts facts = ["First fact", "Second fact", "Third fact"] - results = emit_key_facts({'facts': facts}) + results = emit_key_facts.invoke({"facts": facts}) # Verify return messages assert results == [ @@ -116,12 +105,10 @@ def test_emit_key_facts(reset_memory): def test_delete_key_facts(reset_memory): """Test deleting multiple key facts""" # Add some test facts - emit_key_fact("First fact") - emit_key_fact("Second fact") - emit_key_fact("Third fact") + emit_key_facts.invoke({"facts": ["First fact", "Second fact", "Third fact"]}) # Test deleting mix of existing and non-existing IDs - results = delete_key_facts({'fact_ids': [0, 1, 999]}) + results = delete_key_facts.invoke({"fact_ids": [0, 1, 999]}) # Verify only success messages for existing facts assert results == [