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	from google.genai import types
#5
#6	from mem0.configs.llms.base import BaseLlmConfig
#7	from mem0.llms.gemini import GeminiLLM
#8
#9
#10	@pytest.fixture
#11	def mock_gemini_client():
#12	with patch("mem0.llms.gemini.genai.Client") as mock_client_class:
#13	mock_client = Mock()
#14	mock_client_class.return_value = mock_client
#15	yield mock_client
#16
#17
#18	def test_generate_response_without_tools(mock_gemini_client: Mock):
#19	config = BaseLlmConfig(model="gemini-2.0-flash-latest", temperature=0.7, max_tokens=100, top_p=1.0)
#20	llm = GeminiLLM(config)
#21	messages = [
#22	{"role": "system", "content": "You are a helpful assistant."},
#23	{"role": "user", "content": "Hello, how are you?"},
#24	]
#25
#26	mock_part = Mock(text="I'm doing well, thank you for asking!")
#27	mock_content = Mock(parts=[mock_part])
#28	mock_candidate = Mock(content=mock_content)
#29	mock_response = Mock(candidates=[mock_candidate])
#30
#31	mock_gemini_client.models.generate_content.return_value = mock_response
#32
#33	response = llm.generate_response(messages)
#34
#35	# Check the actual call - system instruction is now in config
#36	mock_gemini_client.models.generate_content.assert_called_once()
#37	call_args = mock_gemini_client.models.generate_content.call_args
#38
#39	# Verify model and contents
#40	assert call_args.kwargs["model"] == "gemini-2.0-flash-latest"
#41	assert len(call_args.kwargs["contents"]) == 1 # Only user message
#42
#43	# Verify config has system instruction
#44	config_arg = call_args.kwargs["config"]
#45	assert config_arg.system_instruction == "You are a helpful assistant."
#46	assert config_arg.temperature == 0.7
#47	assert config_arg.max_output_tokens == 100
#48	assert config_arg.top_p == 1.0
#49
#50	assert response == "I'm doing well, thank you for asking!"
#51
#52
#53	def test_generate_response_with_tools(mock_gemini_client: Mock):
#54	config = BaseLlmConfig(model="gemini-1.5-flash-latest", temperature=0.7, max_tokens=100, top_p=1.0)
#55	llm = GeminiLLM(config)
#56	messages = [
#57	{"role": "system", "content": "You are a helpful assistant."},
#58	{"role": "user", "content": "Add a new memory: Today is a sunny day."},
#59	]
#60	tools = [
#61	{
#62	"type": "function",
#63	"function": {
#64	"name": "add_memory",
#65	"description": "Add a memory",
#66	"parameters": {
#67	"type": "object",
#68	"properties": {"data": {"type": "string", "description": "Data to add to memory"}},
#69	"required": ["data"],
#70	},
#71	},
#72	}
#73	]
#74
#75	mock_tool_call = Mock()
#76	mock_tool_call.name = "add_memory"
#77	mock_tool_call.args = {"data": "Today is a sunny day."}
#78
#79	# Create mock parts with both text and function_call
#80	mock_text_part = Mock()
#81	mock_text_part.text = "I've added the memory for you."
#82	mock_text_part.function_call = None
#83
#84	mock_func_part = Mock()
#85	mock_func_part.text = None
#86	mock_func_part.function_call = mock_tool_call
#87
#88	mock_content = Mock()
#89	mock_content.parts = [mock_text_part, mock_func_part]
#90
#91	mock_candidate = Mock()
#92	mock_candidate.content = mock_content
#93
#94	mock_response = Mock(candidates=[mock_candidate])
#95	mock_gemini_client.models.generate_content.return_value = mock_response
#96
#97	response = llm.generate_response(messages, tools=tools)
#98
#99	# Check the actual call
#100	mock_gemini_client.models.generate_content.assert_called_once()
#101	call_args = mock_gemini_client.models.generate_content.call_args
#102
#103	# Verify model and contents
#104	assert call_args.kwargs["model"] == "gemini-1.5-flash-latest"
#105	assert len(call_args.kwargs["contents"]) == 1 # Only user message
#106
#107	# Verify config has system instruction and tools
#108	config_arg = call_args.kwargs["config"]
#109	assert config_arg.system_instruction == "You are a helpful assistant."
#110	assert config_arg.temperature == 0.7
#111	assert config_arg.max_output_tokens == 100
#112	assert config_arg.top_p == 1.0
#113	assert len(config_arg.tools) == 1
#114	assert config_arg.tool_config.function_calling_config.mode == types.FunctionCallingConfigMode.AUTO
#115
#116	assert response["content"] == "I've added the memory for you."
#117	assert len(response["tool_calls"]) == 1
#118	assert response["tool_calls"][0]["name"] == "add_memory"
#119	assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."}
#120

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public