my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2	from unittest.mock import Mock, patch
#3
#4	import pytest
#5
#6	from mem0.configs.base import MemoryConfig
#7	from mem0.memory.main import Memory
#8
#9
#10	@pytest.fixture(autouse=True)
#11	def mock_openai():
#12	os.environ["OPENAI_API_KEY"] = "123"
#13	with patch("openai.OpenAI") as mock:
#14	mock.return_value = Mock()
#15	yield mock
#16
#17
#18	@pytest.fixture
#19	def memory_instance():
#20	with (
#21	patch("mem0.utils.factory.EmbedderFactory") as mock_embedder,
#22	patch("mem0.memory.main.VectorStoreFactory") as mock_vector_store,
#23	patch("mem0.utils.factory.LlmFactory") as mock_llm,
#24	patch("mem0.memory.telemetry.capture_event"),
#25	patch("mem0.memory.graph_memory.MemoryGraph"),
#26	patch("mem0.memory.main.GraphStoreFactory") as mock_graph_store,
#27	):
#28	mock_embedder.create.return_value = Mock()
#29	mock_vector_store.create.return_value = Mock()
#30	mock_vector_store.create.return_value.search.return_value = []
#31	mock_llm.create.return_value = Mock()
#32
#33	# Create a mock instance that won't try to access config attributes
#34	mock_graph_instance = Mock()
#35	mock_graph_store.create.return_value = mock_graph_instance
#36
#37	config = MemoryConfig(version="v1.1")
#38	config.graph_store.config = {"some_config": "value"}
#39	return Memory(config)
#40
#41
#42	@pytest.fixture
#43	def memory_custom_instance():
#44	with (
#45	patch("mem0.utils.factory.EmbedderFactory") as mock_embedder,
#46	patch("mem0.memory.main.VectorStoreFactory") as mock_vector_store,
#47	patch("mem0.utils.factory.LlmFactory") as mock_llm,
#48	patch("mem0.memory.telemetry.capture_event"),
#49	patch("mem0.memory.graph_memory.MemoryGraph"),
#50	patch("mem0.memory.main.GraphStoreFactory") as mock_graph_store,
#51	):
#52	mock_embedder.create.return_value = Mock()
#53	mock_vector_store.create.return_value = Mock()
#54	mock_vector_store.create.return_value.search.return_value = []
#55	mock_llm.create.return_value = Mock()
#56
#57	# Create a mock instance that won't try to access config attributes
#58	mock_graph_instance = Mock()
#59	mock_graph_store.create.return_value = mock_graph_instance
#60
#61	config = MemoryConfig(
#62	version="v1.1",
#63	custom_fact_extraction_prompt="custom prompt extracting memory",
#64	custom_update_memory_prompt="custom prompt determining memory update",
#65	)
#66	config.graph_store.config = {"some_config": "value"}
#67	return Memory(config)
#68
#69
#70	@pytest.mark.parametrize("version, enable_graph", [("v1.0", False), ("v1.1", True)])
#71	def test_add(memory_instance, version, enable_graph):
#72	memory_instance.config.version = version
#73	memory_instance.enable_graph = enable_graph
#74	memory_instance._add_to_vector_store = Mock(return_value=[{"memory": "Test memory", "event": "ADD"}])
#75	memory_instance._add_to_graph = Mock(return_value=[])
#76
#77	result = memory_instance.add(messages=[{"role": "user", "content": "Test message"}], user_id="test_user")
#78
#79	if enable_graph:
#80	assert "results" in result
#81	assert result["results"] == [{"memory": "Test memory", "event": "ADD"}]
#82	assert "relations" in result
#83	assert result["relations"] == []
#84	else:
#85	assert "results" in result
#86	assert result["results"] == [{"memory": "Test memory", "event": "ADD"}]
#87
#88	memory_instance._add_to_vector_store.assert_called_once_with(
#89	[{"role": "user", "content": "Test message"}], {"user_id": "test_user"}, {"user_id": "test_user"}, True
#90	)
#91
#92	# Remove the conditional assertion for _add_to_graph
#93	memory_instance._add_to_graph.assert_called_once_with(
#94	[{"role": "user", "content": "Test message"}], {"user_id": "test_user"}
#95	)
#96
#97
#98	def test_get(memory_instance):
#99	mock_memory = Mock(
#100	id="test_id",
#101	payload={
#102	"data": "Test memory",
#103	"user_id": "test_user",
#104	"hash": "test_hash",
#105	"created_at": "2023-01-01T00:00:00",
#106	"updated_at": "2023-01-02T00:00:00",
#107	"extra_field": "extra_value",
#108	},
#109	)
#110	memory_instance.vector_store.get = Mock(return_value=mock_memory)
#111
#112	result = memory_instance.get("test_id")
#113
#114	assert result["id"] == "test_id"
#115	assert result["memory"] == "Test memory"
#116	assert result["user_id"] == "test_user"
#117	assert result["hash"] == "test_hash"
#118	assert result["created_at"] == "2023-01-01T00:00:00"
#119	assert result["updated_at"] == "2023-01-02T00:00:00"
#120	assert result["metadata"] == {"extra_field": "extra_value"}
#121
#122
#123	@pytest.mark.parametrize("version, enable_graph", [("v1.0", False), ("v1.1", True)])
#124	def test_search(memory_instance, version, enable_graph):
#125	memory_instance.config.version = version
#126	memory_instance.enable_graph = enable_graph
#127	mock_memories = [
#128	Mock(id="1", payload={"data": "Memory 1", "user_id": "test_user"}, score=0.9),
#129	Mock(id="2", payload={"data": "Memory 2", "user_id": "test_user"}, score=0.8),
#130	]
#131	memory_instance.vector_store.search = Mock(return_value=mock_memories)
#132	memory_instance.embedding_model.embed = Mock(return_value=[0.1, 0.2, 0.3])
#133	memory_instance.graph.search = Mock(return_value=[{"relation": "test_relation"}])
#134
#135	result = memory_instance.search("test query", user_id="test_user")
#136
#137	if version == "v1.1":
#138	assert "results" in result
#139	assert len(result["results"]) == 2
#140	assert result["results"][0]["id"] == "1"
#141	assert result["results"][0]["memory"] == "Memory 1"
#142	assert result["results"][0]["user_id"] == "test_user"
#143	assert result["results"][0]["score"] == 0.9
#144	if enable_graph:
#145	assert "relations" in result
#146	assert result["relations"] == [{"relation": "test_relation"}]
#147	else:
#148	assert "relations" not in result
#149	else:
#150	assert isinstance(result, dict)
#151	assert "results" in result
#152	assert len(result["results"]) == 2
#153	assert result["results"][0]["id"] == "1"
#154	assert result["results"][0]["memory"] == "Memory 1"
#155	assert result["results"][0]["user_id"] == "test_user"
#156	assert result["results"][0]["score"] == 0.9
#157
#158	memory_instance.vector_store.search.assert_called_once_with(
#159	query="test query", vectors=[0.1, 0.2, 0.3], limit=100, filters={"user_id": "test_user"}
#160	)
#161	memory_instance.embedding_model.embed.assert_called_once_with("test query", "search")
#162
#163	if enable_graph:
#164	memory_instance.graph.search.assert_called_once_with("test query", {"user_id": "test_user"}, 100)
#165	else:
#166	memory_instance.graph.search.assert_not_called()
#167
#168
#169	def test_update(memory_instance):
#170	memory_instance.embedding_model = Mock()
#171	memory_instance.embedding_model.embed = Mock(return_value=[0.1, 0.2, 0.3])
#172
#173	memory_instance._update_memory = Mock()
#174
#175	result = memory_instance.update("test_id", "Updated memory")
#176
#177	memory_instance._update_memory.assert_called_once_with(
#178	"test_id", "Updated memory", {"Updated memory": [0.1, 0.2, 0.3]}
#179	)
#180
#181	assert result["message"] == "Memory updated successfully!"
#182
#183
#184	def test_delete(memory_instance):
#185	memory_instance._delete_memory = Mock()
#186
#187	result = memory_instance.delete("test_id")
#188
#189	memory_instance._delete_memory.assert_called_once_with("test_id")
#190	assert result["message"] == "Memory deleted successfully!"
#191
#192
#193	@pytest.mark.parametrize("version, enable_graph", [("v1.0", False), ("v1.1", True)])
#194	def test_delete_all(memory_instance, version, enable_graph):
#195	memory_instance.config.version = version
#196	memory_instance.enable_graph = enable_graph
#197	mock_memories = [Mock(id="1"), Mock(id="2")]
#198	memory_instance.vector_store.list = Mock(return_value=(mock_memories, None))
#199	memory_instance._delete_memory = Mock()
#200	memory_instance.graph.delete_all = Mock()
#201
#202	result = memory_instance.delete_all(user_id="test_user")
#203
#204	assert memory_instance._delete_memory.call_count == 2
#205
#206	if enable_graph:
#207	memory_instance.graph.delete_all.assert_called_once_with({"user_id": "test_user"})
#208	else:
#209	memory_instance.graph.delete_all.assert_not_called()
#210
#211	assert result["message"] == "Memories deleted successfully!"
#212
#213
#214	@pytest.mark.parametrize(
#215	"version, enable_graph, expected_result",
#216	[
#217	("v1.0", False, {"results": [{"id": "1", "memory": "Memory 1", "user_id": "test_user"}]}),
#218	("v1.1", False, {"results": [{"id": "1", "memory": "Memory 1", "user_id": "test_user"}]}),
#219	(
#220	"v1.1",
#221	True,
#222	{
#223	"results": [{"id": "1", "memory": "Memory 1", "user_id": "test_user"}],
#224	"relations": [{"source": "entity1", "relationship": "rel", "target": "entity2"}],
#225	},
#226	),
#227	],
#228	)
#229	def test_get_all(memory_instance, version, enable_graph, expected_result):
#230	memory_instance.config.version = version
#231	memory_instance.enable_graph = enable_graph
#232	mock_memories = [Mock(id="1", payload={"data": "Memory 1", "user_id": "test_user"})]
#233	memory_instance.vector_store.list = Mock(return_value=(mock_memories, None))
#234	memory_instance.graph.get_all = Mock(
#235	return_value=[{"source": "entity1", "relationship": "rel", "target": "entity2"}]
#236	)
#237
#238	result = memory_instance.get_all(user_id="test_user")
#239
#240	assert isinstance(result, dict)
#241	assert "results" in result
#242	assert len(result["results"]) == len(expected_result["results"])
#243	for expected_item, result_item in zip(expected_result["results"], result["results"]):
#244	assert all(key in result_item for key in expected_item)
#245	assert result_item["id"] == expected_item["id"]
#246	assert result_item["memory"] == expected_item["memory"]
#247	assert result_item["user_id"] == expected_item["user_id"]
#248
#249	if enable_graph:
#250	assert "relations" in result
#251	assert result["relations"] == expected_result["relations"]
#252	else:
#253	assert "relations" not in result
#254
#255	memory_instance.vector_store.list.assert_called_once_with(filters={"user_id": "test_user"}, limit=100)
#256
#257	if enable_graph:
#258	memory_instance.graph.get_all.assert_called_once_with({"user_id": "test_user"}, 100)
#259	else:
#260	memory_instance.graph.get_all.assert_not_called()
#261
#262
#263	def test_custom_prompts(memory_custom_instance):
#264	messages = [{"role": "user", "content": "Test message"}]
#265	from mem0.embeddings.mock import MockEmbeddings
#266
#267	memory_custom_instance.llm.generate_response = Mock()
#268	memory_custom_instance.llm.generate_response.return_value = '{"facts": ["fact1", "fact2"]}'
#269	memory_custom_instance.embedding_model = MockEmbeddings()
#270
#271	with patch("mem0.memory.main.parse_messages", return_value="Test message") as mock_parse_messages:
#272	with patch(
#273	"mem0.memory.main.get_update_memory_messages", return_value="custom update memory prompt"
#274	) as mock_get_update_memory_messages:
#275	memory_custom_instance.add(messages=messages, user_id="test_user")
#276
#277	## custom prompt
#278	##
#279	mock_parse_messages.assert_called_once_with(messages)
#280
#281	memory_custom_instance.llm.generate_response.assert_any_call(
#282	messages=[
#283	{"role": "system", "content": memory_custom_instance.config.custom_fact_extraction_prompt},
#284	{"role": "user", "content": f"Input:\n{mock_parse_messages.return_value}"},
#285	],
#286	response_format={"type": "json_object"},
#287	)
#288
#289	## custom update memory prompt
#290	##
#291	mock_get_update_memory_messages.assert_called_once_with(
#292	[], ["fact1", "fact2"], memory_custom_instance.config.custom_update_memory_prompt
#293	)
#294
#295	memory_custom_instance.llm.generate_response.assert_any_call(
#296	messages=[{"role": "user", "content": mock_get_update_memory_messages.return_value}],
#297	response_format={"type": "json_object"},
#298	)
#299

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public