my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	from unittest.mock import MagicMock, Mock, patch
#2
#3	import pytest
#4
#5	from mem0 import AsyncMemory, Memory
#6	from mem0.configs.llms.base import BaseLlmConfig
#7	from mem0.llms.vllm import VllmLLM
#8
#9
#10	@pytest.fixture
#11	def mock_vllm_client():
#12	with patch("mem0.llms.vllm.OpenAI") as mock_openai:
#13	mock_client = Mock()
#14	mock_openai.return_value = mock_client
#15	yield mock_client
#16
#17
#18	def test_generate_response_without_tools(mock_vllm_client):
#19	config = BaseLlmConfig(model="Qwen/Qwen2.5-32B-Instruct", temperature=0.7, max_tokens=100, top_p=1.0)
#20	llm = VllmLLM(config)
#21	messages = [
#22	{"role": "system", "content": "You are a helpful assistant."},
#23	{"role": "user", "content": "Hello, how are you?"},
#24	]
#25
#26	mock_response = Mock()
#27	mock_response.choices = [Mock(message=Mock(content="I'm doing well, thank you for asking!"))]
#28	mock_vllm_client.chat.completions.create.return_value = mock_response
#29
#30	response = llm.generate_response(messages)
#31
#32	mock_vllm_client.chat.completions.create.assert_called_once_with(
#33	model="Qwen/Qwen2.5-32B-Instruct", messages=messages, temperature=0.7, max_tokens=100, top_p=1.0
#34	)
#35	assert response == "I'm doing well, thank you for asking!"
#36
#37
#38	def test_generate_response_with_tools(mock_vllm_client):
#39	config = BaseLlmConfig(model="Qwen/Qwen2.5-32B-Instruct", temperature=0.7, max_tokens=100, top_p=1.0)
#40	llm = VllmLLM(config)
#41	messages = [
#42	{"role": "system", "content": "You are a helpful assistant."},
#43	{"role": "user", "content": "Add a new memory: Today is a sunny day."},
#44	]
#45	tools = [
#46	{
#47	"type": "function",
#48	"function": {
#49	"name": "add_memory",
#50	"description": "Add a memory",
#51	"parameters": {
#52	"type": "object",
#53	"properties": {"data": {"type": "string", "description": "Data to add to memory"}},
#54	"required": ["data"],
#55	},
#56	},
#57	}
#58	]
#59
#60	mock_response = Mock()
#61	mock_message = Mock()
#62	mock_message.content = "I've added the memory for you."
#63
#64	mock_tool_call = Mock()
#65	mock_tool_call.function.name = "add_memory"
#66	mock_tool_call.function.arguments = '{"data": "Today is a sunny day."}'
#67
#68	mock_message.tool_calls = [mock_tool_call]
#69	mock_response.choices = [Mock(message=mock_message)]
#70	mock_vllm_client.chat.completions.create.return_value = mock_response
#71
#72	response = llm.generate_response(messages, tools=tools)
#73
#74	mock_vllm_client.chat.completions.create.assert_called_once_with(
#75	model="Qwen/Qwen2.5-32B-Instruct",
#76	messages=messages,
#77	temperature=0.7,
#78	max_tokens=100,
#79	top_p=1.0,
#80	tools=tools,
#81	tool_choice="auto",
#82	)
#83
#84	assert response["content"] == "I've added the memory for you."
#85	assert len(response["tool_calls"]) == 1
#86	assert response["tool_calls"][0]["name"] == "add_memory"
#87	assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."}
#88
#89
#90
#91	def create_mocked_memory():
#92	"""Create a fully mocked Memory instance for testing."""
#93	with patch('mem0.utils.factory.LlmFactory.create') as mock_llm_factory, \
#94	patch('mem0.utils.factory.EmbedderFactory.create') as mock_embedder_factory, \
#95	patch('mem0.utils.factory.VectorStoreFactory.create') as mock_vector_factory, \
#96	patch('mem0.memory.storage.SQLiteManager') as mock_sqlite:
#97
#98	mock_llm = MagicMock()
#99	mock_llm_factory.return_value = mock_llm
#100
#101	mock_embedder = MagicMock()
#102	mock_embedder.embed.return_value = [0.1, 0.2, 0.3]
#103	mock_embedder_factory.return_value = mock_embedder
#104
#105	mock_vector_store = MagicMock()
#106	mock_vector_store.search.return_value = []
#107	mock_vector_store.add.return_value = None
#108	mock_vector_factory.return_value = mock_vector_store
#109
#110	mock_sqlite.return_value = MagicMock()
#111
#112	memory = Memory()
#113	memory.api_version = "v1.0"
#114	return memory, mock_llm, mock_vector_store
#115
#116
#117	def create_mocked_async_memory():
#118	"""Create a fully mocked AsyncMemory instance for testing."""
#119	with patch('mem0.utils.factory.LlmFactory.create') as mock_llm_factory, \
#120	patch('mem0.utils.factory.EmbedderFactory.create') as mock_embedder_factory, \
#121	patch('mem0.utils.factory.VectorStoreFactory.create') as mock_vector_factory, \
#122	patch('mem0.memory.storage.SQLiteManager') as mock_sqlite:
#123
#124	mock_llm = MagicMock()
#125	mock_llm_factory.return_value = mock_llm
#126
#127	mock_embedder = MagicMock()
#128	mock_embedder.embed.return_value = [0.1, 0.2, 0.3]
#129	mock_embedder_factory.return_value = mock_embedder
#130
#131	mock_vector_store = MagicMock()
#132	mock_vector_store.search.return_value = []
#133	mock_vector_store.add.return_value = None
#134	mock_vector_factory.return_value = mock_vector_store
#135
#136	mock_sqlite.return_value = MagicMock()
#137
#138	memory = AsyncMemory()
#139	memory.api_version = "v1.0"
#140	return memory, mock_llm, mock_vector_store
#141
#142
#143	def test_thinking_tags_sync():
#144	"""Test thinking tags handling in Memory._add_to_vector_store (sync)."""
#145	memory, mock_llm, mock_vector_store = create_mocked_memory()
#146
#147	# Mock LLM responses for both phases
#148	mock_llm.generate_response.side_effect = [
#149	' <think>Sync fact extraction</think> \n{"facts": ["User loves sci-fi"]}',
#150	' <think>Sync memory actions</think> \n{"memory": [{"text": "Loves sci-fi", "event": "ADD"}]}'
#151	]
#152
#153	mock_vector_store.search.return_value = []
#154
#155	result = memory._add_to_vector_store(
#156	messages=[{"role": "user", "content": "I love sci-fi movies"}],
#157	metadata={},
#158	filters={},
#159	infer=True
#160	)
#161
#162	assert len(result) == 1
#163	assert result[0]["memory"] == "Loves sci-fi"
#164	assert result[0]["event"] == "ADD"
#165
#166
#167
#168	@pytest.mark.asyncio
#169	async def test_async_thinking_tags_async():
#170	"""Test thinking tags handling in AsyncMemory._add_to_vector_store."""
#171	memory, mock_llm, mock_vector_store = create_mocked_async_memory()
#172
#173	# Directly mock llm.generate_response instead of via asyncio.to_thread
#174	mock_llm.generate_response.side_effect = [
#175	' <think>Async fact extraction</think> \n{"facts": ["User loves sci-fi"]}',
#176	' <think>Async memory actions</think> \n{"memory": [{"text": "Loves sci-fi", "event": "ADD"}]}'
#177	]
#178
#179	# Mock asyncio.to_thread to call the function directly (bypass threading)
#180	async def mock_to_thread(func, args, *kwargs):
#181	if func == mock_llm.generate_response:
#182	return func(args, *kwargs)
#183	elif hasattr(func, '__name__') and 'embed' in func.__name__:
#184	return [0.1, 0.2, 0.3]
#185	elif hasattr(func, '__name__') and 'search' in func.__name__:
#186	return []
#187	else:
#188	return func(args, *kwargs)
#189
#190	with patch('mem0.memory.main.asyncio.to_thread', side_effect=mock_to_thread):
#191	result = await memory._add_to_vector_store(
#192	messages=[{"role": "user", "content": "I love sci-fi movies"}],
#193	metadata={},
#194	effective_filters={},
#195	infer=True
#196	)
#197
#198	assert len(result) == 1
#199	assert result[0]["memory"] == "Loves sci-fi"
#200	assert result[0]["event"] == "ADD"

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public