allow key facts to be added/removed in batches
This commit is contained in:
parent
48991e5abd
commit
6fe7905a82
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue