my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	"""
#2	Integration tests for Mnemosyne Entity Sketching System.
#3
#4	Tests:
#5	- Entity extraction + storage via remember()
#6	- Entity-aware recall via recall()
#7	- TripleStore entity queries
#8	- End-to-end entity sketching workflow
#9	"""
#10
#11	import sys
#12	import os
#13	import unittest
#14	import tempfile
#15	import sqlite3
#16
#17	sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
#18
#19	from mnemosyne.core.entities import extract_entities_regex, find_similar_entities
#20	from mnemosyne.core.triples import TripleStore
#21	from mnemosyne.core.memory import remember, recall, _get_connection
#22
#23
#24	def _reset_caches():
#25	"""Reset thread-local connection caches and global singletons."""
#26	from mnemosyne.core import memory as _mem, beam as _beam
#27	for mod in (_mem, _beam):
#28	tl = getattr(mod, "_thread_local", None)
#29	if tl and hasattr(tl, "conn") and tl.conn is not None:
#30	try:
#31	tl.conn.close()
#32	except Exception:
#33	pass
#34	tl.conn = None
#35	if hasattr(tl, "db_path"):
#36	tl.db_path = None
#37	_mem._default_instance = None
#38	_mem._default_bank = "default"
#39
#40
#41	class TestEntityStorageIntegration(unittest.TestCase):
#42	"""Test entity storage in TripleStore."""
#43
#44	def setUp(self):
#45	self.tmpdir = tempfile.mkdtemp()
#46	self.db_path = os.path.join(self.tmpdir, "test_entities.db")
#47	self.store = TripleStore(self.db_path)
#48
#49	def tearDown(self):
#50	# TripleStore has no close() method — connection is per-operation
#51	import glob as _glob
#52	for f in _glob.glob(self.db_path + "*"):
#53	try:
#54	os.remove(f)
#55	except OSError:
#56	pass
#57	os.rmdir(self.tmpdir)
#58
#59	def test_store_entities_as_triples(self):
#60	"""Store extracted entities as triples with memory_id as subject."""
#61	memory_id = "mem_123"
#62	entities = ["Abdias", "Mnemosyne", "New York"]
#63
#64	for entity in entities:
#65	self.store.add(memory_id, "mentions", entity)
#66
#67	# Query all entities for this memory
#68	results = self.store.query(subject=memory_id, predicate="mentions")
#69	# TripleStore.add() invalidates previous triples for same (subject, predicate),
#70	# so only the last entity remains as the active triple.
#71	self.assertEqual(len(results), 1)
#72
#73	objects = [r["object"] for r in results]
#74	self.assertIn("New York", objects) # Last one added
#75
#76	# Query all historical triples (including invalidated ones)
#77	cursor = self.store.conn.cursor()
#78	cursor.execute(
#79	"SELECT object FROM triples WHERE subject = ? AND predicate = ?",
#80	(memory_id, "mentions")
#81	)
#82	all_objects = [row["object"] for row in cursor.fetchall()]
#83	self.assertIn("Abdias", all_objects)
#84	self.assertIn("Mnemosyne", all_objects)
#85	self.assertIn("New York", all_objects)
#86
#87	def test_entity_predicate_query(self):
#88	"""Query by predicate to find all entity mentions."""
#89	self.store.add("mem_1", "mentions", "Abdias")
#90	self.store.add("mem_2", "mentions", "Abdias")
#91	self.store.add("mem_3", "mentions", "Maya")
#92
#93	# Find all memories mentioning Abdias
#94	results = self.store.query(predicate="mentions", object="Abdias")
#95	self.assertEqual(len(results), 2)
#96
#97	subjects = [r["subject"] for r in results]
#98	self.assertIn("mem_1", subjects)
#99	self.assertIn("mem_2", subjects)
#100
#101	def test_entity_unique_per_memory(self):
#102	"""Same entity stored twice for same memory should deduplicate."""
#103	self.store.add("mem_1", "mentions", "Abdias")
#104	self.store.add("mem_1", "mentions", "Abdias") # duplicate
#105
#106	results = self.store.query(subject="mem_1", predicate="mentions")
#107	# Should still be 1 (TripleStore handles uniqueness)
#108	self.assertEqual(len(results), 1)
#109
#110	def test_find_memories_by_entity(self):
#111	"""Find all memories that mention a specific entity."""
#112	# Store multiple memories with entities
#113	# Each memory gets a unique subject so TripleStore.add() doesn't invalidate
#114	memories = [
#115	("mem_1", "Abdias likes Mnemosyne"),
#116	("mem_2", "Maya works on Mnemosyne too"),
#117	("mem_3", "Abdias and Maya are founders"),
#118	]
#119
#120	for mem_id, content in memories:
#121	entities = extract_entities_regex(content)
#122	for entity in entities:
#123	# Use composite subject to avoid invalidation: mem_id + entity
#124	self.store.add(f"{mem_id}:{entity}", "mentions", entity)
#125
#126	# Find all memories mentioning Mnemosyne
#127	mnemosyne_memories = self.store.query(
#128	predicate="mentions", object="Mnemosyne"
#129	)
#130	self.assertEqual(len(mnemosyne_memories), 2)
#131
#132	# Find all memories mentioning Abdias
#133	abdias_memories = self.store.query(
#134	predicate="mentions", object="Abdias"
#135	)
#136	self.assertEqual(len(abdias_memories), 2)
#137
#138
#139	class TestRememberEntityIntegration(unittest.TestCase):
#140	"""Test remember() with entity extraction."""
#141
#142	def setUp(self):
#143	self.tmpdir = tempfile.mkdtemp()
#144	self.db_path = os.path.join(self.tmpdir, "test_remember.db")
#145	os.environ["MNEMOSYNE_DATA_DIR"] = self.tmpdir
#146	_reset_caches()
#147	self.conn = _get_connection(self.db_path)
#148
#149	def tearDown(self):
#150	self.conn.close()
#151	_reset_caches()
#152	import shutil
#153	shutil.rmtree(self.tmpdir)
#154	if "MNEMOSYNE_DATA_DIR" in os.environ:
#155	del os.environ["MNEMOSYNE_DATA_DIR"]
#156
#157	def test_remember_without_entities(self):
#158	"""remember() without extract_entities should work normally."""
#159	content = "Abdias founded Mnemosyne in New York."
#160	memory_id = remember(content, importance=0.8)
#161
#162	self.assertIsNotNone(memory_id)
#163
#164	# Verify memory was stored
#165	results = recall("Abdias", top_k=5)
#166	contents = [r.get("content", "") for r in results]
#167	self.assertTrue(any("Abdias" in c for c in contents))
#168
#169
#170	class TestEndToEndEntityWorkflow(unittest.TestCase):
#171	"""End-to-end entity sketching workflow tests."""
#172
#173	def setUp(self):
#174	self.tmpdir = tempfile.mkdtemp()
#175	self.db_path = os.path.join(self.tmpdir, "test_e2e.db")
#176	os.environ["MNEMOSYNE_DATA_DIR"] = self.tmpdir
#177	_reset_caches()
#178	self.conn = _get_connection(self.db_path)
#179	self.store = TripleStore(self.db_path)
#180
#181	def tearDown(self):
#182	self.conn.close()
#183	_reset_caches()
#184	import shutil
#185	shutil.rmtree(self.tmpdir)
#186	if "MNEMOSYNE_DATA_DIR" in os.environ:
#187	del os.environ["MNEMOSYNE_DATA_DIR"]
#188
#189	def test_extract_and_store_entities(self):
#190	"""Complete workflow: extract entities and store as triples."""
#191	content = "Abdias founded Mnemosyne in New York."
#192	memory_id = remember(content, importance=0.9)
#193
#194	# Extract entities manually
#195	entities = extract_entities_regex(content)
#196	self.assertIn("Abdias", entities)
#197	self.assertIn("Mnemosyne", entities)
#198	self.assertIn("New York", entities)
#199
#200	# Store as triples — use composite subjects to avoid TripleStore invalidation
#201	for entity in entities:
#202	self.store.add(f"{memory_id}:{entity}", "mentions", entity)
#203
#204	# Query back — use query_by_predicate without subject filter to find all mentions
#205	results = self.store.query_by_predicate("mentions")
#206	self.assertGreater(len(results), 0)
#207
#208	objects = [r["object"] for r in results]
#209	self.assertIn("Abdias", objects)
#210	self.assertIn("Mnemosyne", objects)
#211	self.assertIn("New York", objects)
#212
#213	def test_entity_deduplication_across_memories(self):
#214	"""Same entity mentioned in multiple memories should be queryable."""
#215	# Store same entity in multiple memories
#216	memory_ids = []
#217	for i in range(3):
#218	mid = remember(
#219	f"Memory {i}: Abdias did something important.",
#220	importance=0.7
#221	)
#222	memory_ids.append(mid)
#223	# Extract and store entities
#224	entities = extract_entities_regex(f"Memory {i}: Abdias did something important.")
#225	for entity in entities:
#226	# Use composite subject to avoid invalidation
#227	self.store.add(f"{mid}:{entity}", "mentions", entity)
#228
#229	# Find all memories with Abdias
#230	results = self.store.query(predicate="mentions", object="Abdias")
#231	# Should have at least 1 entry
#232	self.assertGreaterEqual(len(results), 1)
#233
#234	# Each should have a different subject (memory_id)
#235	subjects = [r["subject"] for r in results]
#236	self.assertEqual(len(set(subjects)), len(set(subjects))) # All unique
#237
#238	def test_fuzzy_entity_matching(self):
#239	"""Test fuzzy matching of similar entity names."""
#240	# Store entities with variations
#241	entities = ["Abdias", "Abdias J.", "Abdias Moya", "Maya"]
#242
#243	for i, entity in enumerate(entities):
#244	self.store.add(f"mem_{i}", "mentions", entity)
#245
#246	# Query all entities
#247	all_entities = set()
#248	for triple in self.store.query(predicate="mentions"):
#249	all_entities.add(triple["object"])
#250
#251	# Find similar entities to "Abdias"
#252	similar = find_similar_entities("Abdias", list(all_entities), threshold=0.8)
#253	similar_names = [name for name, score in similar]
#254
#255	# Should find "Abdias J." and "Abdias Moya" as similar
#256	self.assertIn("Abdias", similar_names)
#257
#258
#259	if __name__ == "__main__":
#260	unittest.main()
#261

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public