Fix CI issues related to missing dependency (#3096)

tests/llms/test_gemini.py

@@ -9,7 +9,7 @@ from mem0.llms.gemini import GeminiLLM
 
 @pytest.fixture
 def mock_gemini_client():
-    with patch("mem0.llms.gemini.genai") as mock_client_class:
+    with patch("mem0.llms.gemini.genai.Client") as mock_client_class:
         mock_client = Mock()
         mock_client_class.return_value = mock_client
         yield mock_client
@@ -24,43 +24,30 @@ def test_generate_response_without_tools(mock_gemini_client: Mock):
     ]
 
     mock_part = Mock(text="I'm doing well, thank you for asking!")
-    mock_embedding = Mock()
-    mock_embedding.values = [0.1, 0.2, 0.3]
-
-    mock_response = Mock()
-    mock_response.candidates = [Mock()]
-    mock_response.candidates[0].content.parts = [Mock()]
-    mock_response.candidates[0].content.parts[0].text = "I'm doing well, thank you for asking!"
+    mock_content = Mock(parts=[mock_part])
+    mock_candidate = Mock(content=mock_content)
+    mock_response = Mock(candidates=[mock_candidate])
 
     mock_gemini_client.models.generate_content.return_value = mock_response
-    mock_content = Mock(parts=[mock_part])
-    mock_message = Mock(content=mock_content)
-    mock_response = Mock(candidates=[mock_message])
-    mock_gemini_client.generate_content.return_value = mock_response
 
     response = llm.generate_response(messages)
 
-    mock_gemini_client.generate_content.assert_called_once_with(
-        contents=[
-            {"parts": "THIS IS A SYSTEM PROMPT. YOU MUST OBEY THIS: You are a helpful assistant.", "role": "user"},
-            {"parts": "Hello, how are you?", "role": "user"},
-        ],
-        config=types.GenerateContentConfig(
-              temperature=0.7, 
-              max_output_tokens=100, 
-              top_p=1.0,
-              tools=None,
-              tool_config=types.ToolConfig(
-                function_calling_config=types.FunctionCallingConfig(
-                    allowed_function_names=None,
-                    mode="auto"
-                
-                )
-            )
-    )   )
-    
-    assert response == "I'm doing well, thank you for asking!"
+    # Check the actual call - system instruction is now in config
+    mock_gemini_client.models.generate_content.assert_called_once()
+    call_args = mock_gemini_client.models.generate_content.call_args
 
+    # Verify model and contents
+    assert call_args.kwargs['model'] == "gemini-2.0-flash-latest"
+    assert len(call_args.kwargs['contents']) == 1  # Only user message
+
+    # Verify config has system instruction
+    config_arg = call_args.kwargs['config']
+    assert config_arg.system_instruction == "You are a helpful assistant."
+    assert config_arg.temperature == 0.7
+    assert config_arg.max_output_tokens == 100
+    assert config_arg.top_p == 1.0
+
+    assert response == "I'm doing well, thank you for asking!"
 
 
 def test_generate_response_with_tools(mock_gemini_client: Mock):
@@ -89,64 +76,42 @@ def test_generate_response_with_tools(mock_gemini_client: Mock):
     mock_tool_call.name = "add_memory"
     mock_tool_call.args = {"data": "Today is a sunny day."}
 
-    mock_part = Mock()
-    mock_part.function_call = mock_tool_call
-    mock_part.text = "I've added the memory for you."
+    # Create mock parts with both text and function_call
+    mock_text_part = Mock()
+    mock_text_part.text = "I've added the memory for you."
+    mock_text_part.function_call = None
+
+    mock_func_part = Mock()
+    mock_func_part.text = None
+    mock_func_part.function_call = mock_tool_call
 
     mock_content = Mock()
-    mock_content.parts = [mock_part]
+    mock_content.parts = [mock_text_part, mock_func_part]
 
-    mock_message = Mock()
-    mock_message.content = mock_content
+    mock_candidate = Mock()
+    mock_candidate.content = mock_content
 
-    mock_response = Mock(candidates=[mock_message])
-    mock_gemini_client.generate_content.return_value = mock_response
+    mock_response = Mock(candidates=[mock_candidate])
+    mock_gemini_client.models.generate_content.return_value = mock_response
 
     response = llm.generate_response(messages, tools=tools)
 
-    mock_gemini_client.generate_content.assert_called_once_with(
-        contents=[
-            {
-                "parts": "THIS IS A SYSTEM PROMPT. YOU MUST OBEY THIS: You are a helpful assistant.",
-                "role": "user"
-            },
-            {
-                "parts": "Add a new memory: Today is a sunny day.",
-                "role": "user"
-            },
-        ],
-        config=types.GenerateContentConfig(
-            temperature=0.7,
-            max_output_tokens=100,
-            top_p=1.0,
-            tools=[
-                types.Tool(
-                    function_declarations=[
-                        types.FunctionDeclaration(
-                            name="add_memory",
-                            description="Add a memory",
-                            parameters={
-                                "type": "object",
-                                "properties": {
-                                    "data": {
-                                        "type": "string",
-                                        "description": "Data to add to memory"
-                                    }
-                                },
-                                "required": ["data"]
-                            }
-                        )
-                    ]
-                )
-            ],
-            tool_config=types.ToolConfig(
-                function_calling_config=types.FunctionCallingConfig(
-                    allowed_function_names=None,
-                    mode="auto"
-                )
-            )
-        )
-    )
+    # Check the actual call
+    mock_gemini_client.models.generate_content.assert_called_once()
+    call_args = mock_gemini_client.models.generate_content.call_args
+
+    # Verify model and contents
+    assert call_args.kwargs['model'] == "gemini-1.5-flash-latest"
+    assert len(call_args.kwargs['contents']) == 1  # Only user message
+
+    # Verify config has system instruction and tools
+    config_arg = call_args.kwargs['config']
+    assert config_arg.system_instruction == "You are a helpful assistant."
+    assert config_arg.temperature == 0.7
+    assert config_arg.max_output_tokens == 100
+    assert config_arg.top_p == 1.0
+    assert len(config_arg.tools) == 1
+    assert config_arg.tool_config.function_calling_config.mode == types.FunctionCallingConfigMode.AUTO
 
     assert response["content"] == "I've added the memory for you."
     assert len(response["tool_calls"]) == 1

tests/llms/test_lm_studio.py

@@ -43,6 +43,7 @@ def test_generate_response_without_tools(mock_lm_studio_client):
 
     assert response == "I'm doing well, thank you for asking!"
 
+
 def test_generate_response_specifying_response_format(mock_lm_studio_client):
     config = BaseLlmConfig(
         model="lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
@@ -68,4 +69,4 @@ def test_generate_response_specifying_response_format(mock_lm_studio_client):
         response_format={"type": "json_schema"},
     )
 
-    assert response == "I'm doing well, thank you for asking!"
+    assert response == "I'm doing well, thank you for asking!"

tests/llms/test_vllm.py

@@ -71,7 +71,13 @@ def test_generate_response_with_tools(mock_vllm_client):
     response = llm.generate_response(messages, tools=tools)
 
     mock_vllm_client.chat.completions.create.assert_called_once_with(
-        model="Qwen/Qwen2.5-32B-Instruct", messages=messages, temperature=0.7, max_tokens=100, top_p=1.0, tools=tools, tool_choice="auto"
+        model="Qwen/Qwen2.5-32B-Instruct",
+        messages=messages,
+        temperature=0.7,
+        max_tokens=100,
+        top_p=1.0,
+        tools=tools,
+        tool_choice="auto",
     )
 
     assert response["content"] == "I've added the memory for you."