95 lines
3.8 KiB
Python
95 lines
3.8 KiB
Python
"""
|
|
Tests for layered memory retrieval with source attribution.
|
|
|
|
Run: python -m unittest brain.scripts.test_layered_retrieval -v
|
|
"""
|
|
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from brain.scripts.layered_memory_store import LayeredMemoryStore
|
|
from brain.scripts.memory_policy import MemoryPolicy
|
|
|
|
|
|
class TestLayeredRetrieval(unittest.TestCase):
|
|
def test_retrieve_all_returns_records_with_attribution(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
project_root = Path(tmpdir)
|
|
# Configure policy to read from multiple layers
|
|
policy = MemoryPolicy(
|
|
project_root=project_root,
|
|
read_layers=["project_agent", "project_global"],
|
|
write_layers=["project_agent", "project_global"],
|
|
)
|
|
store = LayeredMemoryStore(policy=policy, agent_id="agent-1")
|
|
|
|
# Write to project_global
|
|
store.append_entry(
|
|
layer="project_global",
|
|
record={
|
|
"id": "global-1",
|
|
"created_at": "2026-02-11T00:00:00Z",
|
|
"scope": "project_global",
|
|
"entry_type": "fact",
|
|
"content": "Global Fact",
|
|
"project_id": "rlm-mem",
|
|
},
|
|
)
|
|
|
|
# Write to project_agent
|
|
store.append_entry(
|
|
layer="project_agent",
|
|
record={
|
|
"id": "agent-1",
|
|
"created_at": "2026-02-11T00:00:01Z",
|
|
"scope": "project_agent",
|
|
"agent_id": "agent-1",
|
|
"entry_type": "note",
|
|
"content": "Agent Note",
|
|
"project_id": "rlm-mem",
|
|
},
|
|
)
|
|
|
|
all_records = store.get_all_records()
|
|
|
|
# Precedence should be project_agent then project_global
|
|
self.assertEqual(len(all_records), 2)
|
|
|
|
# First record should be from project_agent
|
|
self.assertEqual(all_records[0]["id"], "agent-1")
|
|
self.assertEqual(all_records[0]["source_layer"], "project_agent")
|
|
self.assertIn("agents", all_records[0]["source_path"])
|
|
self.assertIn("agent-1", all_records[0]["source_path"])
|
|
|
|
# Second record should be from project_global
|
|
self.assertEqual(all_records[1]["id"], "global-1")
|
|
self.assertEqual(all_records[1]["source_layer"], "project_global")
|
|
self.assertIn("global", all_records[1]["source_path"])
|
|
|
|
def test_retrieval_respects_policy_layer_order(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
project_root = Path(tmpdir)
|
|
# Reverse order for testing
|
|
policy = MemoryPolicy(
|
|
project_root=project_root,
|
|
read_layers=["project_global", "project_agent"],
|
|
write_layers=["project_agent", "project_global"],
|
|
)
|
|
store = LayeredMemoryStore(policy=policy, agent_id="agent-1")
|
|
|
|
store.append_entry(layer="project_global", record={
|
|
"id": "g1", "created_at": "2026-02-11T00:00:00Z", "scope": "project_global",
|
|
"entry_type": "fact", "content": "G", "project_id": "rlm-mem"
|
|
})
|
|
store.append_entry(layer="project_agent", record={
|
|
"id": "a1", "created_at": "2026-02-11T00:00:00Z", "scope": "project_agent",
|
|
"agent_id": "agent-1", "entry_type": "fact", "content": "A", "project_id": "rlm-mem"
|
|
})
|
|
|
|
all_records = store.get_all_records()
|
|
self.assertEqual(all_records[0]["id"], "g1")
|
|
self.assertEqual(all_records[1]["id"], "a1")
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2)
|