Merge pull request #19556 from BerriAI/litellm_fix_gemini_batch_jan22

Fix: generation config empty for batch

litellm/batches/main.py

@@ -404,6 +404,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
     _retrieve_batch_request: RetrieveBatchRequest,
     _is_async: bool,
     custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
+    logging_obj: Optional[Any] = None,
 ):
     api_base: Optional[str] = None
     if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
@@ -499,6 +500,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
             vertex_credentials=vertex_credentials,
             timeout=timeout,
             max_retries=optional_params.max_retries,
+            logging_obj=logging_obj,
         )
     elif custom_llm_provider == "anthropic":
         api_base = (
@@ -662,6 +664,7 @@ def retrieve_batch(
             _retrieve_batch_request=_retrieve_batch_request,
             _is_async=_is_async,
             timeout=timeout,
+            logging_obj=litellm_logging_obj,
         )
 
     except Exception as e:

litellm/llms/vertex_ai/batches/handler.py

@@ -142,6 +142,7 @@ class VertexAIBatchPrediction(VertexLLM):
         vertex_location: Optional[str],
         timeout: Union[float, httpx.Timeout],
         max_retries: Optional[int],
+        logging_obj: Optional[Any] = None,
     ) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]:
         sync_handler = _get_httpx_client()
 
@@ -187,8 +188,30 @@ class VertexAIBatchPrediction(VertexLLM):
             return self._async_retrieve_batch(
                 api_base=api_base,
                 headers=headers,
+                logging_obj=logging_obj,
             )
 
+        # Log the request using logging_obj if available
+        if logging_obj is not None:
+            from litellm.litellm_core_utils.litellm_logging import Logging
+            if isinstance(logging_obj, Logging):
+                logging_obj.pre_call(
+                    input="",
+                    api_key="",
+                    additional_args={
+                        "complete_input_dict": {},
+                        "api_base": api_base,
+                        "headers": headers,
+                        "request_str": (
+                            f"\nGET Request Sent from LiteLLM:\n"
+                            f"curl -X GET \\\n"
+                            f"{api_base} \\\n"
+                            f"-H 'Authorization: Bearer ***REDACTED***' \\\n"
+                            f"-H 'Content-Type: application/json; charset=utf-8'\n"
+                        ),
+                    },
+                )
+
         response = sync_handler.get(
             url=api_base,
             headers=headers,
@@ -207,10 +230,33 @@ class VertexAIBatchPrediction(VertexLLM):
         self,
         api_base: str,
         headers: Dict[str, str],
+        logging_obj: Optional[Any] = None,
     ) -> LiteLLMBatch:
         client = get_async_httpx_client(
             llm_provider=litellm.LlmProviders.VERTEX_AI,
         )
+        
+        # Log the request using logging_obj if available
+        if logging_obj is not None:
+            from litellm.litellm_core_utils.litellm_logging import Logging
+            if isinstance(logging_obj, Logging):
+                logging_obj.pre_call(
+                    input="",
+                    api_key="",
+                    additional_args={
+                        "complete_input_dict": {},
+                        "api_base": api_base,
+                        "headers": headers,
+                        "request_str": (
+                            f"\nGET Request Sent from LiteLLM:\n"
+                            f"curl -X GET \\\n"
+                            f"{api_base} \\\n"
+                            f"-H 'Authorization: Bearer ***REDACTED***' \\\n"
+                            f"-H 'Content-Type: application/json; charset=utf-8'\n"
+                        ),
+                    },
+                )
+    
         response = await client.get(
             url=api_base,
             headers=headers,

litellm/llms/vertex_ai/gemini/transformation.py

@@ -591,7 +591,7 @@ def _transform_request_body(
             data["toolConfig"] = tool_choice
         if safety_settings is not None:
             data["safetySettings"] = safety_settings
-        if generation_config is not None:
+        if generation_config is not None and len(generation_config) > 0:
             data["generationConfig"] = generation_config
         if cached_content is not None:
             data["cachedContent"] = cached_content

tests/local_testing/test_amazing_vertex_completion.py

@@ -2352,8 +2352,7 @@ async def test_completion_fine_tuned_model():
     expected_payload = {
         "contents": [
             {"role": "user", "parts": [{"text": "Write a short poem about the sky"}]}
-        ],
-        "generationConfig": {},
+        ]
     }
 
     with patch(
@@ -2833,7 +2832,6 @@ def test_gemini_function_call_parameter_in_messages():
                 }
             ],
             "toolConfig": {"functionCallingConfig": {"mode": "AUTO"}},
-            "generationConfig": {},
         } == mock_client.call_args.kwargs["json"]
 
 

tests/local_testing/test_completion.py

@@ -2330,10 +2330,6 @@ async def test_completion_functions_param():
                 "litellm_param_is_function_call"
                 not in mock_client.call_args.kwargs["json"]
             )
-            assert (
-                "litellm_param_is_function_call"
-                not in mock_client.call_args.kwargs["json"]["generationConfig"]
-            )
             assert response.choices[0].message.function_call is not None
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")

tests/test_litellm/llms/vertex_ai/test_vertex.py

@@ -410,7 +410,6 @@ def test_multiple_function_call():
                 },
                 {"role": "user", "parts": [{"text": "tell me the results."}]},
             ],
-            "generationConfig": {},
         }