diff --git a/ra_aid/tools/memory.py b/ra_aid/tools/memory.py index 505ec5f..3d2cd8b 100644 --- a/ra_aid/tools/memory.py +++ b/ra_aid/tools/memory.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Any, Union, TypedDict, Optional +from typing import Dict, List, Any, Union, TypedDict, Optional, Sequence from rich.console import Console from rich.markdown import Markdown from rich.panel import Panel @@ -118,6 +118,33 @@ def emit_key_fact(fact: str) -> str: # Return fact with ID return f"Stored fact #{fact_id}: {fact}" +@tool("emit_key_facts") +def emit_key_facts(facts: List[str]) -> List[str]: + """Store multiple key facts about the project or current task in global memory. + + Args: + facts: List of key facts to store + + Returns: + List of stored fact confirmation messages + """ + results = [] + for fact in facts: + # 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")) + + # Add result message + results.append(f"Stored fact #{fact_id}: {fact}") + + return results + @tool("delete_key_fact") def delete_key_fact(fact_id: int) -> str: """Delete a key fact from global memory by its ID. @@ -139,6 +166,28 @@ def delete_key_fact(fact_id: int) -> str: 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]: + """Delete multiple key facts from global memory by their IDs. + Silently skips any IDs that don't exist. + + Args: + fact_ids: List of fact IDs to delete + + Returns: + List of success messages for deleted facts + """ + results = [] + for fact_id in fact_ids: + if fact_id in _global_memory['key_facts']: + # 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")) + results.append(success_msg) + + return results + @tool("request_implementation") def request_implementation(reason: str) -> str: """Request that implementation proceed after research/planning. diff --git a/tests/ra_aid/tools/test_memory.py b/tests/ra_aid/tools/test_memory.py index b34bd07..db231e4 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 ( emit_key_fact, delete_key_fact, get_memory_value, - emit_research_subtask + emit_research_subtask, + emit_key_facts, + delete_key_facts ) @pytest.fixture @@ -90,6 +92,49 @@ def test_get_memory_value_other_types(reset_memory): # Test with non-existent key assert get_memory_value('nonexistent') == "" +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}) + + # Verify return messages + assert results == [ + "Stored fact #0: First fact", + "Stored fact #1: Second fact", + "Stored fact #2: Third fact" + ] + + # Verify facts stored in memory with correct IDs + assert _global_memory['key_facts'][0] == "First fact" + assert _global_memory['key_facts'][1] == "Second fact" + assert _global_memory['key_facts'][2] == "Third fact" + + # Verify counter incremented correctly + assert _global_memory['key_fact_id_counter'] == 3 + +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") + + # Test deleting mix of existing and non-existing IDs + results = delete_key_facts({'fact_ids': [0, 1, 999]}) + + # Verify only success messages for existing facts + assert results == [ + "Successfully deleted fact #0: First fact", + "Successfully deleted fact #1: Second fact" + ] + + # Verify correct facts removed from memory + assert 0 not in _global_memory['key_facts'] + assert 1 not in _global_memory['key_facts'] + assert 2 in _global_memory['key_facts'] # ID 2 should remain + assert _global_memory['key_facts'][2] == "Third fact" + def test_emit_research_subtask(reset_memory): """Test emitting research subtasks""" # Test adding a research subtask