diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 126eb09a51..3367f567a7 100644 --- a/litellm/batches/main.py +++ b/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: diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 12ce8b48aa..36f5e65e7a 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/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, diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 96e0963a92..3004f39b97 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/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 diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index af8942f89e..1987a7a545 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/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"] diff --git a/tests/local_testing/test_completion.py b/tests/local_testing/test_completion.py index 8d815829d4..de10034ca9 100644 --- a/tests/local_testing/test_completion.py +++ b/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}") diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index fdba86af4a..803584b561 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/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": {}, }