allow key facts to be added/removed in batches

This commit is contained in:
AI Christianson 2024-12-11 11:52:57 -05:00
parent 48991e5abd
commit 6fe7905a82
2 changed files with 96 additions and 2 deletions

View File

@ -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.

View File

@ -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