Code formatting (#1986)

tests/llms/test_gemini_llm.py

@@ -33,19 +33,19 @@ def test_generate_response_without_tools(mock_gemini_client: Mock):
     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"}
+        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"},
         ],
-        generation_config = GenerationConfig(temperature=0.7, max_output_tokens=100, top_p=1.0),
-        tools = None,
-        tool_config = content_types.to_tool_config(
-            {"function_calling_config":
-                {"mode": 'auto', "allowed_function_names": None}
-            })
+        generation_config=GenerationConfig(temperature=0.7, max_output_tokens=100, top_p=1.0),
+        tools=None,
+        tool_config=content_types.to_tool_config(
+            {"function_calling_config": {"mode": "auto", "allowed_function_names": None}}
+        ),
     )
     assert response == "I'm doing well, thank you for asking!"
-    
+
+
 def test_generate_response_with_tools(mock_gemini_client: Mock):
     config = BaseLlmConfig(model="gemini-1.5-flash-latest", temperature=0.7, max_tokens=100, top_p=1.0)
     llm = GeminiLLM(config)
@@ -74,13 +74,13 @@ def test_generate_response_with_tools(mock_gemini_client: Mock):
 
     mock_part = Mock()
     mock_part.function_call = mock_tool_call
-    mock_part.text="I've added the memory for you."
+    mock_part.text = "I've added the memory for you."
 
     mock_content = Mock()
-    mock_content.parts=[mock_part]
+    mock_content.parts = [mock_part]
 
     mock_message = Mock()
-    mock_message.content=mock_content
+    mock_message.content = mock_content
 
     mock_response = Mock(candidates=[mock_message])
     mock_gemini_client.generate_content.return_value = mock_response
@@ -88,28 +88,29 @@ def test_generate_response_with_tools(mock_gemini_client: Mock):
     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"}
+        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"},
         ],
-        generation_config = GenerationConfig(temperature=0.7, max_output_tokens=100, top_p=1.0),
-        tools = [
+        generation_config=GenerationConfig(temperature=0.7, max_output_tokens=100, top_p=1.0),
+        tools=[
             {
-                "function_declarations": [{
-                    "name": "add_memory", 
-                    "description": "Add a memory",
-                    "parameters": {
-                        "type": "object", 
-                        "properties": {"data": {"type": "string", "description": "Data to add to memory"}},
-                        "required": ["data"]
+                "function_declarations": [
+                    {
+                        "name": "add_memory",
+                        "description": "Add a memory",
+                        "parameters": {
+                            "type": "object",
+                            "properties": {"data": {"type": "string", "description": "Data to add to memory"}},
+                            "required": ["data"],
+                        },
                     }
-                }]
+                ]
             }
         ],
-        tool_config = content_types.to_tool_config(
-            {"function_calling_config":
-                {"mode": 'auto', "allowed_function_names": None}
-            })
+        tool_config=content_types.to_tool_config(
+            {"function_calling_config": {"mode": "auto", "allowed_function_names": None}}
+        ),
     )
 
     assert response["content"] == "I've added the memory for you."

tests/llms/test_openai.py

@@ -31,8 +31,9 @@ def test_openai_llm_base_url():
 
     # case3: with config.openai_base_url
     config_base_url = "https://api.config.com/v1"
-    config = BaseLlmConfig(model="gpt-4o", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key",
-                           openai_base_url=config_base_url)
+    config = BaseLlmConfig(
+        model="gpt-4o", temperature=0.7, max_tokens=100, top_p=1.0, api_key="api_key", openai_base_url=config_base_url
+    )
     llm = OpenAILLM(config)
     # Note: openai client will parse the raw base_url into a URL object, which will have a trailing slash
     assert str(llm.client.base_url) == config_base_url + "/"