From c1d495f09eded4339e6afefb2e7b71ba3b907edd Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 30 Jun 2025 15:17:46 -0700 Subject: [PATCH] [Bug Fix] Allow passing litellm_params when using generateContent API endpoint (#12177) * add _add_generic_litellm_params_to_request * fix type * fix: setup_result.litellm_params * fix passing litellm params * test_api_base_and_api_key_passthrough * fixes for passing litellm params * fix linting errors --- litellm/google_genai/adapters/handler.py | 8 ++ .../google_genai/adapters/transformation.py | 48 +++++++-- litellm/google_genai/main.py | 3 + .../google_genai/test_google_genai_adapter.py | 100 ++++++++++++++++++ 4 files changed, 151 insertions(+), 8 deletions(-) diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index 651a6413ca..ee7ddd0f2c 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -1,6 +1,7 @@ from typing import Any, AsyncIterator, Coroutine, Dict, List, Optional, Union, cast import litellm +from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ModelResponse from .transformation import GoogleGenAIAdapter @@ -18,6 +19,7 @@ class GenerateContentToCompletionHandler: contents: Union[List[Dict[str, Any]], Dict[str, Any]], config: Optional[Dict[str, Any]] = None, stream: bool = False, + litellm_params: Optional[GenericLiteLLMParams] = None, extra_kwargs: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """Prepare kwargs for litellm.completion/acompletion""" @@ -27,6 +29,7 @@ class GenerateContentToCompletionHandler: model=model, contents=contents, config=config, + litellm_params=litellm_params, **(extra_kwargs or {}) ) @@ -41,6 +44,7 @@ class GenerateContentToCompletionHandler: async def async_generate_content_handler( model: str, contents: Union[List[Dict[str, Any]], Dict[str, Any]], + litellm_params: GenericLiteLLMParams, config: Optional[Dict[str, Any]] = None, stream: bool = False, **kwargs, @@ -52,6 +56,7 @@ class GenerateContentToCompletionHandler: contents=contents, config=config, stream=stream, + litellm_params=litellm_params, extra_kwargs=kwargs, ) @@ -82,6 +87,7 @@ class GenerateContentToCompletionHandler: def generate_content_handler( model: str, contents: Union[List[Dict[str, Any]], Dict[str, Any]], + litellm_params: GenericLiteLLMParams, config: Optional[Dict[str, Any]] = None, stream: bool = False, _is_async: bool = False, @@ -95,6 +101,7 @@ class GenerateContentToCompletionHandler: contents=contents, config=config, stream=stream, + litellm_params=litellm_params, **kwargs, ) @@ -103,6 +110,7 @@ class GenerateContentToCompletionHandler: contents=contents, config=config, stream=stream, + litellm_params=litellm_params, extra_kwargs=kwargs, ) diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index e80da47d5e..7d6fdaf6e0 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -13,6 +13,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParam, ChatCompletionUserMessage, ) +from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( AdapterCompletionStreamWrapper, Choices, @@ -107,8 +108,9 @@ class GoogleGenAIAdapter: model: str, contents: Union[List[Dict[str, Any]], Dict[str, Any]], config: Optional[Dict[str, Any]] = None, + litellm_params: Optional[GenericLiteLLMParams] = None, **kwargs, - ) -> ChatCompletionRequest: + ) -> Dict[str, Any]: """ Transform generate_content request to litellm completion format @@ -119,7 +121,7 @@ class GoogleGenAIAdapter: **kwargs: Additional parameters Returns: - ChatCompletionRequest in OpenAI format + Dict in OpenAI format """ # Normalize contents to list format @@ -131,11 +133,11 @@ class GoogleGenAIAdapter: # Transform contents to OpenAI messages format messages = self._transform_contents_to_messages(contents_list) - # Create base request - completion_request: ChatCompletionRequest = ChatCompletionRequest( - model=model, - messages=messages, - ) + # Create base request as dict (which is compatible with ChatCompletionRequest) + completion_request: ChatCompletionRequest = { + "model": model, + "messages": messages, + } ######################################################### # Supported OpenAI chat completion params @@ -182,8 +184,38 @@ class GoogleGenAIAdapter: ) if tool_choice: completion_request["tool_choice"] = tool_choice + + ######################################################### + # forward any litellm specific params + ######################################################### + completion_request_dict = dict(completion_request) + if litellm_params: + completion_request_dict = self._add_generic_litellm_params_to_request( + completion_request_dict=completion_request_dict, + litellm_params=litellm_params + ) - return completion_request + return completion_request_dict + + def _add_generic_litellm_params_to_request( + self, + completion_request_dict: Dict[str, Any], + litellm_params: Optional[GenericLiteLLMParams] = None + ) -> dict: + """Add generic litellm params to request. e.g add api_base, api_key, api_version, etc. + + Args: + completion_request_dict: Dict[str, Any] + litellm_params: GenericLiteLLMParams + + Returns: + Dict[str, Any] + """ + if litellm_params: + litellm_dict = litellm_params.model_dump(exclude_none=True) + for key, value in litellm_dict.items(): + completion_request_dict[key] = value + return completion_request_dict def translate_completion_output_params_streaming( self, completion_stream: Any diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index f1847057db..5bf85bfc55 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -300,6 +300,7 @@ def generate_content( config=setup_result.generate_content_config_dict, stream=False, _is_async=_is_async, + litellm_params=setup_result.litellm_params, **kwargs ) @@ -377,6 +378,7 @@ async def agenerate_content_stream( model=setup_result.model, contents=contents, # type: ignore config=setup_result.generate_content_config_dict, + litellm_params=setup_result.litellm_params, stream=True, **kwargs ) @@ -452,6 +454,7 @@ def generate_content_stream( config=setup_result.generate_content_config_dict, stream=True, _is_async=_is_async, + litellm_params=setup_result.litellm_params, **kwargs ) diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/test_litellm/google_genai/test_google_genai_adapter.py index 45fa937523..3e79434f29 100644 --- a/tests/test_litellm/google_genai/test_google_genai_adapter.py +++ b/tests/test_litellm/google_genai/test_google_genai_adapter.py @@ -18,6 +18,8 @@ import sys import pytest +import litellm + def test_adapter_import(): """Test that the adapter can be imported successfully""" @@ -869,6 +871,104 @@ def test_handler_parameter_exclusion(): assert "temperature" in completion_kwargs assert completion_kwargs["temperature"] == 0.7 +@pytest.mark.parametrize("function_name,is_async,is_stream", [ + ("generate_content", False, False), + ("agenerate_content", True, False), + ("generate_content_stream", False, True), + ("agenerate_content_stream", True, True), +]) +def test_api_base_and_api_key_passthrough(function_name, is_async, is_stream): + """Test that api_base and api_key parameters are passed through to litellm.completion/acompletion when using generate_content""" + import asyncio + import unittest.mock + + litellm._turn_on_debug() + + # Import the specific function being tested + if function_name == "generate_content": + from litellm.google_genai.main import generate_content as test_function + elif function_name == "agenerate_content": + from litellm.google_genai.main import agenerate_content as test_function + elif function_name == "generate_content_stream": + from litellm.google_genai.main import generate_content_stream as test_function + elif function_name == "agenerate_content_stream": + from litellm.google_genai.main import agenerate_content_stream as test_function + + # Test input parameters + model = "gpt-3.5-turbo" + test_api_base = "https://test-api.example.com" + test_api_key = "test-api-key-123" + + # Mock the appropriate litellm function (completion vs acompletion) + mock_target = 'litellm.acompletion' if is_async else 'litellm.completion' + + with unittest.mock.patch(mock_target) as mock_completion: + # Mock return value + mock_return = unittest.mock.MagicMock() + if is_async: + # For async functions, return a coroutine that resolves to the mock + async def mock_async_return(): + return mock_return + mock_completion.return_value = mock_async_return() + else: + mock_completion.return_value = mock_return + + # Define the test call + def make_test_call(): + return test_function( + model=model, + contents={ + "role": "user", + "parts": [{"text": "Hello, world!"}] + }, + config={ + "temperature": 0.7, + }, + api_base=test_api_base, + api_key=test_api_key + ) + + # Call the handler with api_base and api_key + try: + if is_async: + # Run the async function + async def run_async_test(): + return await make_test_call() + + asyncio.run(run_async_test()) + else: + make_test_call() + except Exception: + # Ignore any errors from the mock response processing + pass + + # Verify that the appropriate litellm function was called + mock_completion.assert_called_once() + + # Get the arguments passed to litellm.completion/acompletion + call_args, call_kwargs = mock_completion.call_args + + # Verify that api_base and api_key were passed through + assert "api_base" in call_kwargs, f"api_base not found in completion kwargs: {call_kwargs.keys()}" + assert call_kwargs["api_base"] == test_api_base, f"Expected api_base {test_api_base}, got {call_kwargs['api_base']}" + + assert "api_key" in call_kwargs, f"api_key not found in completion kwargs: {call_kwargs.keys()}" + assert call_kwargs["api_key"] == test_api_key, f"Expected api_key {test_api_key}, got {call_kwargs['api_key']}" + + # Verify other expected parameters + assert call_kwargs["model"] == model + assert len(call_kwargs["messages"]) == 1 + assert call_kwargs["messages"][0]["role"] == "user" + assert call_kwargs["messages"][0]["content"] == "Hello, world!" + assert call_kwargs["temperature"] == 0.7 + + # Verify stream parameter for streaming functions + if is_stream: + assert call_kwargs.get("stream") is True, f"Expected stream=True for {function_name}" + else: + # For non-streaming, stream should be False or not present + assert call_kwargs.get("stream") is not True, f"Expected stream not True for {function_name}" + def test_shared_schema_normalization_utilities(): """Test the shared schema normalization utility functions work correctly""" from litellm.litellm_core_utils.json_validation_rule import (