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.llms.base import BaseLlmConfig
#7	from mem0.configs.llms.deepseek import DeepSeekConfig
#8	from mem0.llms.deepseek import DeepSeekLLM
#9
#10
#11	@pytest.fixture
#12	def mock_deepseek_client():
#13	with patch("mem0.llms.deepseek.OpenAI") as mock_openai:
#14	mock_client = Mock()
#15	mock_openai.return_value = mock_client
#16	yield mock_client
#17
#18
#19	def test_deepseek_llm_base_url():
#20	# case1: default config with deepseek official base url
#21	config = BaseLlmConfig(model="deepseek-chat", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key")
#22	llm = DeepSeekLLM(config)
#23	assert str(llm.client.base_url) == "https://api.deepseek.com"
#24
#25	# case2: with env variable DEEPSEEK_API_BASE
#26	provider_base_url = "https://api.provider.com/v1/"
#27	os.environ["DEEPSEEK_API_BASE"] = provider_base_url
#28	config = DeepSeekConfig(model="deepseek-chat", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key")
#29	llm = DeepSeekLLM(config)
#30	assert str(llm.client.base_url) == provider_base_url
#31
#32	# case3: with config.deepseek_base_url
#33	config_base_url = "https://api.config.com/v1/"
#34	config = DeepSeekConfig(
#35	model="deepseek-chat",
#36	temperature=0.7,
#37	max_tokens=100,
#38	top_p=1.0,
#39	api_key="api_key",
#40	deepseek_base_url=config_base_url,
#41	)
#42	llm = DeepSeekLLM(config)
#43	assert str(llm.client.base_url) == config_base_url
#44
#45
#46	def test_generate_response_without_tools(mock_deepseek_client):
#47	config = BaseLlmConfig(model="deepseek-chat", temperature=0.7, max_tokens=100, top_p=1.0)
#48	llm = DeepSeekLLM(config)
#49	messages = [
#50	{"role": "system", "content": "You are a helpful assistant."},
#51	{"role": "user", "content": "Hello, how are you?"},
#52	]
#53
#54	mock_response = Mock()
#55	mock_response.choices = [Mock(message=Mock(content="I'm doing well, thank you for asking!"))]
#56	mock_deepseek_client.chat.completions.create.return_value = mock_response
#57
#58	response = llm.generate_response(messages)
#59
#60	mock_deepseek_client.chat.completions.create.assert_called_once_with(
#61	model="deepseek-chat", messages=messages, temperature=0.7, max_tokens=100, top_p=1.0
#62	)
#63	assert response == "I'm doing well, thank you for asking!"
#64
#65
#66	def test_generate_response_with_tools(mock_deepseek_client):
#67	config = BaseLlmConfig(model="deepseek-chat", temperature=0.7, max_tokens=100, top_p=1.0)
#68	llm = DeepSeekLLM(config)
#69	messages = [
#70	{"role": "system", "content": "You are a helpful assistant."},
#71	{"role": "user", "content": "Add a new memory: Today is a sunny day."},
#72	]
#73	tools = [
#74	{
#75	"type": "function",
#76	"function": {
#77	"name": "add_memory",
#78	"description": "Add a memory",
#79	"parameters": {
#80	"type": "object",
#81	"properties": {"data": {"type": "string", "description": "Data to add to memory"}},
#82	"required": ["data"],
#83	},
#84	},
#85	}
#86	]
#87
#88	mock_response = Mock()
#89	mock_message = Mock()
#90	mock_message.content = "I've added the memory for you."
#91
#92	mock_tool_call = Mock()
#93	mock_tool_call.function.name = "add_memory"
#94	mock_tool_call.function.arguments = '{"data": "Today is a sunny day."}'
#95
#96	mock_message.tool_calls = [mock_tool_call]
#97	mock_response.choices = [Mock(message=mock_message)]
#98	mock_deepseek_client.chat.completions.create.return_value = mock_response
#99
#100	response = llm.generate_response(messages, tools=tools)
#101
#102	mock_deepseek_client.chat.completions.create.assert_called_once_with(
#103	model="deepseek-chat",
#104	messages=messages,
#105	temperature=0.7,
#106	max_tokens=100,
#107	top_p=1.0,
#108	tools=tools,
#109	tool_choice="auto",
#110	)
#111
#112	assert response["content"] == "I've added the memory for you."
#113	assert len(response["tool_calls"]) == 1
#114	assert response["tool_calls"][0]["name"] == "add_memory"
#115	assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."}
#116

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public