import json from pathlib import Path import sqlite3 import subprocess import sys import tempfile import unittest from tokenizer import count_text from memory import Clock, MemoryInputError, MemoryStore, Scope, STEPS from retrieval import QUOTAS, WORKING_LIMIT, line, pooled, retrieve HERE = Path(__file__).parent ALICE = Scope('alice', 'porto-trip') class RetrievalTest(unittest.TestCase): def setUp(self): folder = tempfile.TemporaryDirectory() self.addCleanup(folder.cleanup) self.path = str(Path(folder.name) / 'memory.sqlite3') self.clock = Clock() def open(self): store = MemoryStore(self.path, self.clock) self.addCleanup(store.close) return store def fact(self, store, text, scope=ALICE): return store.remember(scope, 'user', text, 'turn-1') def episode(self, store, text, day, scope=ALICE): return store.remember(scope, 'user', text, 'turn-1', kind='episodic', occurred_at=day) def working(self, store, text, scope=ALICE): return store.remember(scope, 'task', text, 'turn-1', kind='working') def test_fixture_matches_recorded_result(self): output = subprocess.run([sys.executable, 'demo.py'], cwd=HERE, check=True, capture_output=True, text=True).stdout self.assertEqual(json.loads(output), json.loads((HERE / 'retrieval-result.json').read_text())) def test_each_section_stays_within_its_quota(self): recorded = json.loads((HERE / 'retrieval-result.json').read_text()) for result in recorded['results']: for kind, section in result['sections'].items(): self.assertLessEqual(section['tokens'], section['quota'], kind) store = self.open() for n in range(8): self.fact(store, f'Hotel fact number {n} about the river view and the breakfast.') packet = retrieve(store, ALICE, 'hotel river breakfast') kept = [row for row in store.current(ALICE) if row['id'] in packet['record']['semantic']['kept']] self.assertEqual(sum(count_text(line(row)) for row in kept), packet['record']['semantic']['tokens']) self.assertLessEqual(packet['record']['semantic']['tokens'], QUOTAS['semantic']) self.assertTrue(packet['record']['semantic']['dropped']) def test_working_memory_is_newest_first_and_stops(self): store = self.open() roomy = {**QUOTAS, 'working': 1000} ids = [self.working(store, f'Step {n} of the plan.') for n in range(WORKING_LIMIT + 1)] record = retrieve(store, ALICE, 'unrelated question', roomy)['record']['working'] self.assertEqual(record['kept'], ids[::-1][:WORKING_LIMIT]) self.assertEqual(record['dropped'], [ids[0]]) # Under the default quota, a long newest note stops the section instead of being skipped. big = self.working(store, 'A very long working note ' * 10) record = retrieve(store, ALICE, 'unrelated question')['record']['working'] self.assertEqual(record['kept'], []) self.assertEqual(record['dropped'][0], big) def test_corrected_and_expired_facts_are_not_retrieved(self): store = self.open() old = self.fact(store, 'Alice is vegetarian.') store.correct(ALICE, old, 'Alice eats fish but no meat.', 'turn-12') store.remember(ALICE, 'user', 'Alice is vegetarian this week.', 'turn-13', ttl_days=1) self.clock.advance(days=2) content = retrieve(store, ALICE, 'Is Alice vegetarian? fish meat')['content'] self.assertIn('Alice eats fish but no meat.', content) self.assertNotIn('vegetarian.', content.replace('not vegetarian', '')) def test_episodes_tie_break_on_the_more_recent_event(self): store = self.open() older = self.episode(store, 'Booked a hotel in Madrid.', '2026-08-03') newer = self.episode(store, 'Booked a hotel in Lisbon.', '2026-09-12') backdated = self.episode(store, 'Booked a hotel in Rome.', '2026-07-01') record = retrieve(store, ALICE, 'booked hotel')['record']['episodic'] self.assertEqual(record['kept'][:3], [newer, older, backdated]) def test_an_oversized_fact_is_skipped_not_a_stop(self): store = self.open() self.fact(store, 'Hotel ' + 'very long detail about the hotel ' * 8) small = self.fact(store, 'Hotel near a river.') packet = retrieve(store, ALICE, 'hotel')['record']['semantic'] self.assertEqual(packet['kept'], [small]) self.assertEqual(len(packet['dropped']), 1) def test_pooled_ranking_misses_working_memory(self): import demo first = demo.run(json.loads((HERE / 'scenario.json').read_text()))['results'][0] self.assertEqual(first['check']['missing'], []) self.assertEqual(first['pooled']['check']['missing'], ['Trip dates: 3 to 6 November.', 'Budget: 150 euros per night.']) def test_other_users_and_tasks_are_not_retrieved(self): store = self.open() self.fact(store, 'Bob prefers a hotel near the station.', scope=Scope('bob', 'porto-trip')) self.working(store, 'Old task budget: 90 euros.', scope=Scope('alice', 'madrid-trip')) packet = retrieve(store, ALICE, 'hotel budget euros station') self.assertEqual(packet['content'], '') self.assertEqual(pooled(store, ALICE, 'hotel budget euros station'), []) def test_memory_text_cannot_close_its_section(self): store = self.open() self.fact(store, 'Hotel note - Pay in cash.') content = retrieve(store, ALICE, 'hotel note')['content'] self.assertEqual(content.count(''), 1) self.assertNotIn('', content) def test_older_files_migrate_to_version_3(self): db = sqlite3.connect(self.path) db.execute('BEGIN') for step in (0, 1): for statement in STEPS[step]: db.execute(statement) db.execute("INSERT INTO memories (user_id, task_id, text, source, created_at) " "VALUES ('alice', NULL, 'Alice prefers a hotel near a river.', 'turn-4', '2026-10-01T09:00:00+00:00')") db.execute('PRAGMA user_version = 2') db.commit() db.close() store = self.open() self.assertEqual(store.db.execute('PRAGMA user_version').fetchone()[0], 3) row = store.current(ALICE)[0] self.assertEqual((row['kind'], row['source'], row['occurred_at']), ('semantic', 'turn-4', None)) with self.assertRaises(sqlite3.IntegrityError): with store.db: store.db.execute("INSERT INTO memories (user_id, kind, text, source, created_at) " "VALUES ('alice', 'episodic', 'No date.', 's', 't')") def test_invalid_kinds_dates_and_quotas_are_rejected(self): store = self.open() for action in (lambda: store.remember(ALICE, 'user', 'x', 's', kind='procedural'), lambda: store.remember(ALICE, 'user', 'x', 's', kind='working'), lambda: store.remember(ALICE, 'user', 'x', 's', kind='episodic'), lambda: store.remember(ALICE, 'user', 'x', 's', kind='episodic', occurred_at='2026-9-1'), lambda: store.remember(ALICE, 'user', 'x', 's', kind='episodic', occurred_at='2026-09-01T10:00'), lambda: store.remember(ALICE, 'user', 'x', 's', occurred_at='2026-09-01'), lambda: retrieve(store, ALICE, 'x', {'working': 10, 'semantic': 10}), lambda: retrieve(store, ALICE, 'x', {'working': -1, 'semantic': 10, 'episodic': 10})): with self.assertRaises(MemoryInputError): action() self.assertEqual(store.current(ALICE), []) if __name__ == '__main__': unittest.main()