my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	from unittest.mock import Mock, patch
#2
#3	import pytest
#4
#5	from mem0.configs.llms.base import BaseLlmConfig
#6	from mem0.llms.together import TogetherLLM
#7
#8
#9	@pytest.fixture
#10	def mock_together_client():
#11	with patch("mem0.llms.together.Together") as mock_together:
#12	mock_client = Mock()
#13	mock_together.return_value = mock_client
#14	yield mock_client
#15
#16
#17	def test_generate_response_without_tools(mock_together_client):
#18	config = BaseLlmConfig(model="mistralai/Mixtral-8x7B-Instruct-v0.1", temperature=0.7, max_tokens=100, top_p=1.0)
#19	llm = TogetherLLM(config)
#20	messages = [
#21	{"role": "system", "content": "You are a helpful assistant."},
#22	{"role": "user", "content": "Hello, how are you?"},
#23	]
#24
#25	mock_response = Mock()
#26	mock_response.choices = [Mock(message=Mock(content="I'm doing well, thank you for asking!"))]
#27	mock_together_client.chat.completions.create.return_value = mock_response
#28
#29	response = llm.generate_response(messages)
#30
#31	mock_together_client.chat.completions.create.assert_called_once_with(
#32	model="mistralai/Mixtral-8x7B-Instruct-v0.1", messages=messages, temperature=0.7, max_tokens=100, top_p=1.0
#33	)
#34	assert response == "I'm doing well, thank you for asking!"
#35
#36
#37	def test_generate_response_with_tools(mock_together_client):
#38	config = BaseLlmConfig(model="mistralai/Mixtral-8x7B-Instruct-v0.1", temperature=0.7, max_tokens=100, top_p=1.0)
#39	llm = TogetherLLM(config)
#40	messages = [
#41	{"role": "system", "content": "You are a helpful assistant."},
#42	{"role": "user", "content": "Add a new memory: Today is a sunny day."},
#43	]
#44	tools = [
#45	{
#46	"type": "function",
#47	"function": {
#48	"name": "add_memory",
#49	"description": "Add a memory",
#50	"parameters": {
#51	"type": "object",
#52	"properties": {"data": {"type": "string", "description": "Data to add to memory"}},
#53	"required": ["data"],
#54	},
#55	},
#56	}
#57	]
#58
#59	mock_response = Mock()
#60	mock_message = Mock()
#61	mock_message.content = "I've added the memory for you."
#62
#63	mock_tool_call = Mock()
#64	mock_tool_call.function.name = "add_memory"
#65	mock_tool_call.function.arguments = '{"data": "Today is a sunny day."}'
#66
#67	mock_message.tool_calls = [mock_tool_call]
#68	mock_response.choices = [Mock(message=mock_message)]
#69	mock_together_client.chat.completions.create.return_value = mock_response
#70
#71	response = llm.generate_response(messages, tools=tools)
#72
#73	mock_together_client.chat.completions.create.assert_called_once_with(
#74	model="mistralai/Mixtral-8x7B-Instruct-v0.1",
#75	messages=messages,
#76	temperature=0.7,
#77	max_tokens=100,
#78	top_p=1.0,
#79	tools=tools,
#80	tool_choice="auto",
#81	)
#82
#83	assert response["content"] == "I've added the memory for you."
#84	assert len(response["tool_calls"]) == 1
#85	assert response["tool_calls"][0]["name"] == "add_memory"
#86	assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."}
#87

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public