my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import numpy as np
#2	import pytest
#3	from unittest.mock import Mock, patch
#4	from mem0.memory.kuzu_memory import MemoryGraph
#5
#6
#7	class TestKuzu:
#8	"""Test that Kuzu memory works correctly"""
#9
#10	# Create distinct embeddings that won't match with threshold=0.7
#11	# Each embedding is mostly zeros with ones in different positions to ensure low similarity
#12	alice_emb = np.zeros(384)
#13	alice_emb[0:96] = 1.0
#14
#15	bob_emb = np.zeros(384)
#16	bob_emb[96:192] = 1.0
#17
#18	charlie_emb = np.zeros(384)
#19	charlie_emb[192:288] = 1.0
#20
#21	dave_emb = np.zeros(384)
#22	dave_emb[288:384] = 1.0
#23
#24	embeddings = {
#25	"alice": alice_emb.tolist(),
#26	"bob": bob_emb.tolist(),
#27	"charlie": charlie_emb.tolist(),
#28	"dave": dave_emb.tolist(),
#29	}
#30
#31	@pytest.fixture
#32	def mock_config(self):
#33	"""Create a mock configuration for testing"""
#34	config = Mock()
#35
#36	# Mock embedder config
#37	config.embedder.provider = "mock_embedder"
#38	config.embedder.config = {"model": "mock_model"}
#39	config.vector_store.config = {"dimensions": 384}
#40
#41	# Mock graph store config
#42	config.graph_store.config.db = ":memory:"
#43	config.graph_store.threshold = 0.7
#44
#45	# Mock LLM config
#46	config.llm.provider = "mock_llm"
#47	config.llm.config = {"api_key": "test_key"}
#48
#49	return config
#50
#51	@pytest.fixture
#52	def mock_embedding_model(self):
#53	"""Create a mock embedding model"""
#54	mock_model = Mock()
#55	mock_model.config.embedding_dims = 384
#56
#57	def mock_embed(text):
#58	return self.embeddings[text]
#59
#60	mock_model.embed.side_effect = mock_embed
#61	return mock_model
#62
#63	@pytest.fixture
#64	def mock_llm(self):
#65	"""Create a mock LLM"""
#66	mock_llm = Mock()
#67	mock_llm.generate_response.return_value = {
#68	"tool_calls": [
#69	{
#70	"name": "extract_entities",
#71	"arguments": {"entities": [{"entity": "test_entity", "entity_type": "test_type"}]},
#72	}
#73	]
#74	}
#75	return mock_llm
#76
#77	@patch("mem0.memory.kuzu_memory.EmbedderFactory")
#78	@patch("mem0.memory.kuzu_memory.LlmFactory")
#79	def test_kuzu_memory_initialization(
#80	self, mock_llm_factory, mock_embedder_factory, mock_config, mock_embedding_model, mock_llm
#81	):
#82	"""Test that Kuzu memory initializes correctly"""
#83	# Setup mocks
#84	mock_embedder_factory.create.return_value = mock_embedding_model
#85	mock_llm_factory.create.return_value = mock_llm
#86
#87	# Create instance
#88	kuzu_memory = MemoryGraph(mock_config)
#89
#90	# Verify initialization
#91	assert kuzu_memory.config == mock_config
#92	assert kuzu_memory.embedding_model == mock_embedding_model
#93	assert kuzu_memory.embedding_dims == 384
#94	assert kuzu_memory.llm == mock_llm
#95	assert kuzu_memory.threshold == 0.7
#96
#97	@pytest.mark.parametrize(
#98	"embedding_dims",
#99	[None, 0, -1],
#100	)
#101	@patch("mem0.memory.kuzu_memory.EmbedderFactory")
#102	def test_kuzu_memory_initialization_invalid_embedding_dims(
#103	self, mock_embedder_factory, embedding_dims, mock_config
#104	):
#105	"""Test that Kuzu memory raises ValuError when initialized with invalid embedding_dims"""
#106	# Setup mocks
#107	mock_embedding_model = Mock()
#108	mock_embedding_model.config.embedding_dims = embedding_dims
#109	mock_embedder_factory.create.return_value = mock_embedding_model
#110
#111	with pytest.raises(ValueError, match="must be a positive"):
#112	MemoryGraph(mock_config)
#113
#114	@patch("mem0.memory.kuzu_memory.EmbedderFactory")
#115	@patch("mem0.memory.kuzu_memory.LlmFactory")
#116	def test_kuzu(self, mock_llm_factory, mock_embedder_factory, mock_config, mock_embedding_model, mock_llm):
#117	"""Test adding memory to the graph"""
#118	mock_embedder_factory.create.return_value = mock_embedding_model
#119	mock_llm_factory.create.return_value = mock_llm
#120
#121	kuzu_memory = MemoryGraph(mock_config)
#122
#123	filters = {"user_id": "test_user", "agent_id": "test_agent", "run_id": "test_run"}
#124	data1 = [
#125	{"source": "alice", "destination": "bob", "relationship": "knows"},
#126	{"source": "bob", "destination": "charlie", "relationship": "knows"},
#127	{"source": "charlie", "destination": "alice", "relationship": "knows"},
#128	]
#129	data2 = [
#130	{"source": "charlie", "destination": "alice", "relationship": "likes"},
#131	]
#132
#133	result = kuzu_memory._add_entities(data1, filters, {})
#134	assert result[0] == [{"source": "alice", "relationship": "knows", "target": "bob"}]
#135	assert result[1] == [{"source": "bob", "relationship": "knows", "target": "charlie"}]
#136	assert result[2] == [{"source": "charlie", "relationship": "knows", "target": "alice"}]
#137	assert get_node_count(kuzu_memory) == 3
#138	assert get_edge_count(kuzu_memory) == 3
#139
#140	result = kuzu_memory._add_entities(data2, filters, {})
#141	assert result[0] == [{"source": "charlie", "relationship": "likes", "target": "alice"}]
#142	assert get_node_count(kuzu_memory) == 3
#143	assert get_edge_count(kuzu_memory) == 4
#144
#145	data3 = [
#146	{"source": "dave", "destination": "alice", "relationship": "admires"}
#147	]
#148	result = kuzu_memory._add_entities(data3, filters, {})
#149	assert result[0] == [{"source": "dave", "relationship": "admires", "target": "alice"}]
#150	assert get_node_count(kuzu_memory) == 4 # dave is new
#151	assert get_edge_count(kuzu_memory) == 5
#152
#153	results = kuzu_memory.get_all(filters)
#154	assert set([f"{result['source']}_{result['relationship']}_{result['target']}" for result in results]) == set([
#155	"alice_knows_bob",
#156	"bob_knows_charlie",
#157	"charlie_likes_alice",
#158	"charlie_knows_alice",
#159	"dave_admires_alice"
#160	])
#161
#162	results = kuzu_memory._search_graph_db(["bob"], filters, threshold=0.8)
#163	assert set([f"{result['source']}_{result['relationship']}_{result['destination']}" for result in results]) == set([
#164	"alice_knows_bob",
#165	"bob_knows_charlie",
#166	])
#167
#168	result = kuzu_memory._delete_entities(data2, filters)
#169	assert result[0] == [{"source": "charlie", "relationship": "likes", "target": "alice"}]
#170	assert get_node_count(kuzu_memory) == 4
#171	assert get_edge_count(kuzu_memory) == 4
#172
#173	result = kuzu_memory._delete_entities(data1, filters)
#174	assert result[0] == [{"source": "alice", "relationship": "knows", "target": "bob"}]
#175	assert result[1] == [{"source": "bob", "relationship": "knows", "target": "charlie"}]
#176	assert result[2] == [{"source": "charlie", "relationship": "knows", "target": "alice"}]
#177	assert get_node_count(kuzu_memory) == 4
#178	assert get_edge_count(kuzu_memory) == 1
#179
#180	result = kuzu_memory.delete_all(filters)
#181	assert get_node_count(kuzu_memory) == 0
#182	assert get_edge_count(kuzu_memory) == 0
#183
#184	result = kuzu_memory._add_entities(data2, filters, {})
#185	assert result[0] == [{"source": "charlie", "relationship": "likes", "target": "alice"}]
#186	assert get_node_count(kuzu_memory) == 2
#187	assert get_edge_count(kuzu_memory) == 1
#188
#189	result = kuzu_memory.reset()
#190	assert get_node_count(kuzu_memory) == 0
#191	assert get_edge_count(kuzu_memory) == 0
#192
#193	def get_node_count(kuzu_memory):
#194	results = kuzu_memory.kuzu_execute(
#195	"""
#196	MATCH (n)
#197	RETURN COUNT(n) as count
#198	"""
#199	)
#200	return int(results[0]['count'])
#201
#202	def get_edge_count(kuzu_memory):
#203	results = kuzu_memory.kuzu_execute(
#204	"""
#205	MATCH (n)-[e]->(m)
#206	RETURN COUNT(e) as count
#207	"""
#208	)
#209	return int(results[0]['count'])
#210

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public