From 35e46784d3b01bdf80274e1ea212db8b5c9d2247 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 25 Jun 2025 18:37:03 -0700 Subject: [PATCH] [Feat] Add Support for calling Gemini/Vertex models in their native format (#12046) * init litellm google gen ai methods * feat init structure of functions for generate content * add init * add BaseGoogleGenAIGenerateContentConfig * add generate_content_handler * add get_provider_google_genai_generate_content_config * fixes for generate content * add get_vertex_ai_project etc to base * use VertexBase * fixes for BaseGoogleGenAIGenerateContentConfig * working validate env for google gemini * feat - add transform google response * fixes for transform_generate_content_request * fix get_supported_generate_content_optional_params * add BaseGoogleGenAITest * working e2e test * fixes init config * use correct types * fix test for google gen ai * fix types * add sync_get_auth_token_and_url * fixes for transform * add llm http handler for google * working non-streaming google endpoints * add BaseGoogleGenAIGenerateContentStreamingIterator * add GoogleGenAIGenerateContentStreamingIterator * fix working sync stream * fixes for litellm logging obj * working async streaming * add google gen ai types * fix - required imports * fix readme * fix deps * fix deps * fix ruff code QA checks * fix linting * fixes TYPE_CHECKING * fixes for typing --- litellm/google_genai/Readme.md | 123 +++++ litellm/google_genai/__init__.py | 19 + litellm/google_genai/main.py | 439 ++++++++++++++++++ litellm/google_genai/streaming_iterator.py | 118 +++++ .../base_llm/google_genai/transformation.py | 201 ++++++++ litellm/llms/custom_httpx/llm_http_handler.py | 226 ++++++++- .../gemini/google_genai/transformation.py | 295 ++++++++++++ .../transformation.py | 24 +- litellm/llms/vertex_ai/vertex_llm_base.py | 37 +- litellm/types/google_genai/__init__.py | 13 + litellm/types/google_genai/main.py | 18 + litellm/utils.py | 18 + requirements.txt | 2 +- .../unified_google_tests/base_google_test.py | 122 +++++ .../test_google_ai_studio.py | 10 + 15 files changed, 1637 insertions(+), 28 deletions(-) create mode 100644 litellm/google_genai/Readme.md create mode 100644 litellm/google_genai/__init__.py create mode 100644 litellm/google_genai/main.py create mode 100644 litellm/google_genai/streaming_iterator.py create mode 100644 litellm/llms/base_llm/google_genai/transformation.py create mode 100644 litellm/llms/gemini/google_genai/transformation.py create mode 100644 litellm/types/google_genai/__init__.py create mode 100644 litellm/types/google_genai/main.py create mode 100644 tests/unified_google_tests/base_google_test.py create mode 100644 tests/unified_google_tests/test_google_ai_studio.py diff --git a/litellm/google_genai/Readme.md b/litellm/google_genai/Readme.md new file mode 100644 index 0000000000..2c18292652 --- /dev/null +++ b/litellm/google_genai/Readme.md @@ -0,0 +1,123 @@ +# LiteLLM Google GenAI Interface + +Interface to interact with Google GenAI Functions in the native Google interface format. + +## Overview + +This module provides a native interface to Google's Generative AI API, allowing you to use Google's content generation capabilities with both streaming and non-streaming modes, in both synchronous and asynchronous contexts. + +## Available Functions + +### Non-Streaming Functions + +- `generate_content()` - Synchronous content generation +- `agenerate_content()` - Asynchronous content generation + +### Streaming Functions + +- `generate_content_stream()` - Synchronous streaming content generation +- `agenerate_content_stream()` - Asynchronous streaming content generation + +## Usage Examples + +### Basic Non-Streaming Usage + +```python +from litellm.google_genai import generate_content, agenerate_content +from google.genai.types import ContentDict, PartDict + +# Synchronous usage +contents = ContentDict( + parts=[ + PartDict(text="Hello, can you tell me a short joke?") + ], +) + +response = generate_content( + contents=contents, + model="gemini-pro", # or your preferred model + # Add other model-specific parameters as needed +) + +print(response) +``` + +### Async Non-Streaming Usage + +```python +import asyncio +from litellm.google_genai import agenerate_content +from google.genai.types import ContentDict, PartDict + +async def main(): + contents = ContentDict( + parts=[ + PartDict(text="Hello, can you tell me a short joke?") + ], + ) + + response = await agenerate_content( + contents=contents, + model="gemini-pro", + # Add other model-specific parameters as needed + ) + + print(response) + +# Run the async function +asyncio.run(main()) +``` + +### Streaming Usage + +```python +from litellm.google_genai import generate_content_stream +from google.genai.types import ContentDict, PartDict + +# Synchronous streaming +contents = ContentDict( + parts=[ + PartDict(text="Tell me a story about space exploration") + ], +) + +for chunk in generate_content_stream( + contents=contents, + model="gemini-pro", +): + print(f"Chunk: {chunk}") +``` + +### Async Streaming Usage + +```python +import asyncio +from litellm.google_genai import agenerate_content_stream +from google.genai.types import ContentDict, PartDict + +async def main(): + contents = ContentDict( + parts=[ + PartDict(text="Tell me a story about space exploration") + ], + ) + + async for chunk in agenerate_content_stream( + contents=contents, + model="gemini-pro", + ): + print(f"Async chunk: {chunk}") + +asyncio.run(main()) +``` + + +## Testing + +This module includes comprehensive tests covering: +- Sync and async non-streaming requests +- Sync and async streaming requests +- Response validation +- Error handling scenarios + +See `tests/unified_google_tests/base_google_test.py` for test implementation examples. \ No newline at end of file diff --git a/litellm/google_genai/__init__.py b/litellm/google_genai/__init__.py new file mode 100644 index 0000000000..faeb1f227d --- /dev/null +++ b/litellm/google_genai/__init__.py @@ -0,0 +1,19 @@ +""" +This allows using Google GenAI model in their native interface. + +This module provides generate_content functionality for Google GenAI models. +""" + +from .main import ( + agenerate_content, + agenerate_content_stream, + generate_content, + generate_content_stream, +) + +__all__ = [ + "generate_content", + "agenerate_content", + "generate_content_stream", + "agenerate_content_stream", +] \ No newline at end of file diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py new file mode 100644 index 0000000000..6c27ecdadf --- /dev/null +++ b/litellm/google_genai/main.py @@ -0,0 +1,439 @@ +import asyncio +import contextvars +from functools import partial +from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, Iterator, Optional, Union + +import httpx +from pydantic import BaseModel + +import litellm +from litellm.constants import request_timeout +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.google_genai.transformation import ( + BaseGoogleGenAIGenerateContentConfig, +) +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.types.router import GenericLiteLLMParams +from litellm.utils import ProviderConfigManager, client + +if TYPE_CHECKING: + from litellm.types.google_genai.main import ( + GenerateContentConfigDict, + GenerateContentContentListUnionDict, + GenerateContentResponse, + ) +else: + GenerateContentConfigDict = Any + GenerateContentContentListUnionDict = Any + GenerateContentResponse = Any + +####### ENVIRONMENT VARIABLES ################### +# Initialize any necessary instances or variables here +base_llm_http_handler = BaseLLMHTTPHandler() +################################################# + + +class GenerateContentSetupResult(BaseModel): + """Internal Type - Result of setting up a generate content call""" + model: str + request_body: Dict[str, Any] + custom_llm_provider: str + generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig + generate_content_config_dict: Dict[str, Any] + litellm_params: GenericLiteLLMParams + litellm_logging_obj: LiteLLMLoggingObj + litellm_call_id: Optional[str] + + class Config: + arbitrary_types_allowed = True + + +class GenerateContentHelper: + """Helper class for Google GenAI generate content operations""" + + @staticmethod + def mock_generate_content_response( + mock_response: str = "This is a mock response from Google GenAI generate_content.", + ) -> Dict[str, Any]: + """Mock response for generate_content for testing purposes""" + return { + "text": mock_response, + "candidates": [ + { + "content": { + "parts": [{"text": mock_response}], + "role": "model" + }, + "finishReason": "STOP", + "index": 0, + "safetyRatings": [] + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 20, + "totalTokenCount": 30 + } + } + + @staticmethod + def setup_generate_content_call( + model: str, + contents: GenerateContentContentListUnionDict, + config: Optional[GenerateContentConfigDict] = None, + custom_llm_provider: Optional[str] = None, + stream: bool = False, + **kwargs + ) -> GenerateContentSetupResult: + """ + Common setup logic for generate_content calls + + Args: + model: The model name + contents: The content to generate from + config: Optional configuration + custom_llm_provider: Optional custom LLM provider + stream: Whether this is a streaming call + local_vars: Local variables from the calling function + **kwargs: Additional keyword arguments + + Returns: + GenerateContentSetupResult containing all setup information + """ + litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") + litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + + # get llm provider logic + litellm_params = GenericLiteLLMParams(**kwargs) + + ## MOCK RESPONSE LOGIC (only for non-streaming) + if not stream and litellm_params.mock_response and isinstance(litellm_params.mock_response, str): + raise ValueError("Mock response should be handled by caller") + + ( + model, + custom_llm_provider, + dynamic_api_key, + dynamic_api_base, + ) = litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=litellm_params.api_base, + api_key=litellm_params.api_key, + ) + + # get provider config + generate_content_provider_config: Optional[BaseGoogleGenAIGenerateContentConfig] = ( + ProviderConfigManager.get_provider_google_genai_generate_content_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), + ) + ) + + if generate_content_provider_config is None: + operation = "streaming" if stream else "" + raise ValueError( + f"Generate content {operation} is not supported for {custom_llm_provider}".strip() + ) + + + ######################################################################################### + # Construct request body + ######################################################################################### + # Create Google Optional Params Config + generate_content_config_dict = generate_content_provider_config.map_generate_content_optional_params( + generate_content_config_dict=config or {}, + model=model, + ) + request_body = generate_content_provider_config.transform_generate_content_request( + model=model, + contents=contents, + generate_content_config_dict=generate_content_config_dict, + ) + + # Pre Call logging + if litellm_logging_obj is None: + raise ValueError("litellm_logging_obj is required, but got None") + + litellm_logging_obj.update_environment_variables( + model=model, + optional_params=dict(generate_content_config_dict), + litellm_params={ + "litellm_call_id": litellm_call_id, + }, + custom_llm_provider=custom_llm_provider, + ) + + return GenerateContentSetupResult( + model=model, + custom_llm_provider=custom_llm_provider, + request_body=request_body, + generate_content_provider_config=generate_content_provider_config, + generate_content_config_dict=generate_content_config_dict, + litellm_params=litellm_params, + litellm_logging_obj=litellm_logging_obj, + litellm_call_id=litellm_call_id + ) + + +@client +async def agenerate_content( + model: str, + contents: GenerateContentContentListUnionDict, + config: Optional[GenerateContentConfigDict] = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + # LiteLLM specific params, + custom_llm_provider: Optional[str] = None, + **kwargs +) -> Any: + """ + Async: Generate content using Google GenAI + """ + local_vars = locals() + try: + loop = asyncio.get_event_loop() + kwargs["agenerate_content"] = True + + # get custom llm provider so we can use this for mapping exceptions + if custom_llm_provider is None: + _, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + ) + + func = partial( + generate_content, + model=model, + contents=contents, + config=config, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + + ctx = contextvars.copy_context() + func_with_context = partial(ctx.run, func) + init_response = await loop.run_in_executor(None, func_with_context) + + if asyncio.iscoroutine(init_response): + response = await init_response + else: + response = init_response + + return response + except Exception as e: + raise litellm.exception_type( + model=model, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +def generate_content( + model: str, + contents: GenerateContentContentListUnionDict, + config: Optional[GenerateContentConfigDict] = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + # LiteLLM specific params, + custom_llm_provider: Optional[str] = None, + **kwargs, +) -> Any: + """ + Generate content using Google GenAI + """ + local_vars = locals() + try: + _is_async = kwargs.pop("agenerate_content", False) is True + + # Check for mock response first + litellm_params = GenericLiteLLMParams(**kwargs) + if litellm_params.mock_response and isinstance(litellm_params.mock_response, str): + return GenerateContentHelper.mock_generate_content_response( + mock_response=litellm_params.mock_response + ) + + # Setup the call + setup_result = GenerateContentHelper.setup_generate_content_call( + model=model, + contents=contents, + config=config, + custom_llm_provider=custom_llm_provider, + stream=False, + **kwargs + ) + + # Call the handler + response = base_llm_http_handler.generate_content_handler( + model=setup_result.model, + contents=contents, + generate_content_provider_config=setup_result.generate_content_provider_config, + generate_content_config_dict=setup_result.generate_content_config_dict, + custom_llm_provider=setup_result.custom_llm_provider, + litellm_params=setup_result.litellm_params, + logging_obj=setup_result.litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or request_timeout, + _is_async=_is_async, + client=kwargs.get("client"), + stream=False, + litellm_metadata=kwargs.get("litellm_metadata", {}), + ) + + return response + except Exception as e: + raise litellm.exception_type( + model=model, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +async def agenerate_content_stream( + model: str, + contents: GenerateContentContentListUnionDict, + config: Optional[GenerateContentConfigDict] = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + # LiteLLM specific params, + custom_llm_provider: Optional[str] = None, + **kwargs +) -> AsyncIterator[Any]: + """ + Async: Generate content using Google GenAI with streaming response + """ + local_vars = locals() + try: + kwargs["agenerate_content_stream"] = True + + # get custom llm provider so we can use this for mapping exceptions + if custom_llm_provider is None: + _, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=model, api_base=local_vars.get("base_url", None) + ) + + # Setup the call + setup_result = GenerateContentHelper.setup_generate_content_call( + model=model, + contents=contents, + config=config, + custom_llm_provider=custom_llm_provider, + stream=True, + **kwargs + ) + + # Call the handler with async enabled and streaming + async_generator = await base_llm_http_handler.generate_content_handler( + model=setup_result.model, + contents=contents, + generate_content_provider_config=setup_result.generate_content_provider_config, + generate_content_config_dict=setup_result.generate_content_config_dict, + custom_llm_provider=setup_result.custom_llm_provider, + litellm_params=setup_result.litellm_params, + logging_obj=setup_result.litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or request_timeout, + _is_async=True, + client=kwargs.get("client"), + stream=True, + litellm_metadata=kwargs.get("litellm_metadata", {}), + ) + + # Iterate over the async generator + async for chunk in async_generator: + yield chunk + + except Exception as e: + raise litellm.exception_type( + model=model, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + + +@client +def generate_content_stream( + model: str, + contents: GenerateContentContentListUnionDict, + config: Optional[GenerateContentConfigDict] = None, + # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. + # The extra values given here take precedence over values defined on the client or passed to this method. + extra_headers: Optional[Dict[str, Any]] = None, + extra_query: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + # LiteLLM specific params, + custom_llm_provider: Optional[str] = None, + **kwargs, +) -> Iterator[Any]: + """ + Generate content using Google GenAI with streaming response + """ + local_vars = locals() + try: + # Remove any async-related flags since this is the sync function + kwargs.pop("agenerate_content_stream", None) + + # Setup the call + setup_result = GenerateContentHelper.setup_generate_content_call( + model=model, + contents=contents, + config=config, + custom_llm_provider=custom_llm_provider, + stream=True, + **kwargs + ) + + # Call the handler with streaming enabled (sync version) + return base_llm_http_handler.generate_content_handler( + model=setup_result.model, + contents=contents, + generate_content_provider_config=setup_result.generate_content_provider_config, + generate_content_config_dict=setup_result.generate_content_config_dict, + custom_llm_provider=setup_result.custom_llm_provider, + litellm_params=setup_result.litellm_params, + logging_obj=setup_result.litellm_logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout or request_timeout, + _is_async=False, + client=kwargs.get("client"), + stream=True, + litellm_metadata=kwargs.get("litellm_metadata", {}), + ) + + except Exception as e: + raise litellm.exception_type( + model=model, + custom_llm_provider=custom_llm_provider, + original_exception=e, + completion_kwargs=local_vars, + extra_kwargs=kwargs, + ) + diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py new file mode 100644 index 0000000000..9e88e9d772 --- /dev/null +++ b/litellm/google_genai/streaming_iterator.py @@ -0,0 +1,118 @@ +from datetime import datetime +from typing import TYPE_CHECKING, Any, List, Optional + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +if TYPE_CHECKING: + from litellm.llms.base_llm.google_genai.transformation import ( + BaseGoogleGenAIGenerateContentConfig, + ) +else: + BaseGoogleGenAIGenerateContentConfig = Any + +class BaseGoogleGenAIGenerateContentStreamingIterator: + """ + Base class for Google GenAI Generate Content streaming iterators that provides common logic + for streaming response handling and logging. + """ + + def __init__( + self, + litellm_logging_obj: LiteLLMLoggingObj, + request_body: dict, + ): + self.litellm_logging_obj = litellm_logging_obj + self.request_body = request_body + self.start_time = datetime.now() + + async def _handle_streaming_logging(self, collected_chunks: List[bytes]): + """Handle the logging after all chunks have been collected.""" + pass + + +class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContentStreamingIterator): + """ + Streaming iterator specifically for Google GenAI generate content API. + """ + + def __init__( + self, + response, + model: str, + logging_obj: LiteLLMLoggingObj, + generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, + litellm_metadata: dict, + custom_llm_provider: str, + request_body: Optional[dict] = None, + ): + super().__init__( + litellm_logging_obj=logging_obj, + request_body=request_body or {}, + ) + self.response = response + self.model = model + self.generate_content_provider_config = generate_content_provider_config + self.litellm_metadata = litellm_metadata + self.custom_llm_provider = custom_llm_provider + # Store the iterator once to avoid multiple stream consumption + self.stream_iterator = response.iter_lines() + + def __iter__(self): + return self + + def __next__(self): + try: + # Get the next chunk from the stored iterator + chunk = next(self.stream_iterator) + # Just yield raw bytes + return chunk + except StopIteration: + raise StopIteration + + def __aiter__(self): + return self + + async def __anext__(self): + # This should not be used for sync responses + # If you need async iteration, use AsyncGoogleGenAIGenerateContentStreamingIterator + raise NotImplementedError("Use AsyncGoogleGenAIGenerateContentStreamingIterator for async iteration") + + +class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContentStreamingIterator): + """ + Async streaming iterator specifically for Google GenAI generate content API. + """ + + def __init__( + self, + response, + model: str, + logging_obj: LiteLLMLoggingObj, + generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, + litellm_metadata: dict, + custom_llm_provider: str, + request_body: Optional[dict] = None, + ): + super().__init__( + litellm_logging_obj=logging_obj, + request_body=request_body or {}, + ) + self.response = response + self.model = model + self.generate_content_provider_config = generate_content_provider_config + self.litellm_metadata = litellm_metadata + self.custom_llm_provider = custom_llm_provider + # Store the async iterator once to avoid multiple stream consumption + self.stream_iterator = response.aiter_lines() + + def __aiter__(self): + return self + + async def __anext__(self): + try: + # Get the next chunk from the stored async iterator + chunk = await self.stream_iterator.__anext__() + # Just yield raw bytes + return chunk + except StopAsyncIteration: + raise StopAsyncIteration \ No newline at end of file diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py new file mode 100644 index 0000000000..0b0df98800 --- /dev/null +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -0,0 +1,201 @@ +import types +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +import httpx + +if TYPE_CHECKING: + from litellm.types.google_genai.main import ( + GenerateContentConfigDict, + GenerateContentContentListUnionDict, + GenerateContentResponse, + ) +else: + GenerateContentConfigDict = Any + GenerateContentContentListUnionDict = Any + GenerateContentResponse = Any + +from litellm.types.router import GenericLiteLLMParams + + +class BaseGoogleGenAIGenerateContentConfig(ABC): + """Base configuration class for Google GenAI generate_content functionality""" + + def __init__(self): + pass + + @classmethod + def get_config(cls): + return { + k: v + for k, v in cls.__dict__.items() + if not k.startswith("__") + and not k.startswith("_abc") + and not isinstance( + v, + ( + types.FunctionType, + types.BuiltinFunctionType, + classmethod, + staticmethod, + ), + ) + and v is not None + } + + @abstractmethod + def get_supported_generate_content_optional_params(self, model: str) -> List[str]: + """ + Get the list of supported Google GenAI parameters for the model. + + Args: + model: The model name + + Returns: + List of supported parameter names + """ + raise NotImplementedError("get_supported_generate_content_optional_params is not implemented") + + + @abstractmethod + def map_generate_content_optional_params( + self, + generate_content_config_dict: GenerateContentConfigDict, + model: str, + ) -> Dict[str, Any]: + """ + Map Google GenAI parameters to provider-specific format. + + Args: + generate_content_optional_params: Optional parameters for generate content + model: The model name + + Returns: + Mapped parameters for the provider + """ + raise NotImplementedError("map_generate_content_optional_params is not implemented") + + @abstractmethod + def validate_environment( + self, + api_key: Optional[str], + headers: Optional[dict], + model: str, + litellm_params: Optional[Union[GenericLiteLLMParams, dict]] + ) -> dict: + """ + Validate the environment and return headers for the request. + + Args: + api_key: API key + headers: Existing headers + model: The model name + litellm_params: LiteLLM parameters + + Returns: + Updated headers + """ + raise NotImplementedError("validate_environment is not implemented") + + def sync_get_auth_token_and_url( + self, + api_base: Optional[str], + model: str, + litellm_params: dict, + stream: bool, + ) -> Tuple[dict, str]: + """ + Sync version of get_auth_token_and_url. + + Args: + api_base: Base API URL + model: The model name + litellm_params: LiteLLM parameters + stream: Whether this is a streaming call + + Returns: + Tuple of headers and API base + """ + raise NotImplementedError("sync_get_auth_token_and_url is not implemented") + + async def get_auth_token_and_url( + self, + api_base: Optional[str], + model: str, + litellm_params: dict, + stream: bool, + ) -> Tuple[dict, str]: + """ + Get the complete URL for the request. + + Args: + api_base: Base API URL + model: The model name + litellm_params: LiteLLM parameters + + Returns: + Tuple of headers and API base + """ + raise NotImplementedError("get_auth_token_and_url is not implemented") + + @abstractmethod + def transform_generate_content_request( + self, + model: str, + contents: GenerateContentContentListUnionDict, + generate_content_config_dict: Dict, + ) -> dict: + """ + Transform the request parameters for the generate content API. + + Args: + model: The model name + contents: Input contents + generate_content_request_params: Request parameters + litellm_params: LiteLLM parameters + headers: Request headers + + Returns: + Transformed request data + """ + pass + + @abstractmethod + def transform_generate_content_response( + self, + model: str, + raw_response: httpx.Response, + ) -> GenerateContentResponse: + """ + Transform the raw response from the generate content API. + + Args: + model: The model name + raw_response: Raw HTTP response + + Returns: + Transformed response data + """ + pass + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] + ) -> Exception: + """ + Get the appropriate exception class for the error. + + Args: + error_message: Error message + status_code: HTTP status code + headers: Response headers + + Returns: + Exception instance + """ + from litellm.llms.base_llm.chat.transformation import BaseLLMException + + return BaseLLMException( + status_code=status_code, + message=error_message, + headers=headers, + ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 92f9d6d958..769cfe75b6 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -31,6 +31,9 @@ from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.base_llm.files.transformation import BaseFilesConfig +from litellm.llms.base_llm.google_genai.transformation import ( + BaseGoogleGenAIGenerateContentConfig, +) from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig @@ -2349,7 +2352,7 @@ class BaseLLMHTTPHandler: self, e: Exception, provider_config: Union[ - BaseConfig, BaseRerankConfig, BaseResponsesAPIConfig, BaseImageEditConfig, BaseVectorStoreConfig + BaseConfig, BaseRerankConfig, BaseResponsesAPIConfig, BaseImageEditConfig, BaseVectorStoreConfig, BaseGoogleGenAIGenerateContentConfig ], ): status_code = getattr(e, "status_code", 500) @@ -2887,4 +2890,225 @@ class BaseLLMHTTPHandler: return vector_store_provider_config.transform_create_vector_store_response( response=response, ) + + ##################################################################### + ################ Google GenAI GENERATE CONTENT HANDLER ########################### + ##################################################################### + def generate_content_handler( + self, + model: str, + contents: Any, + generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, + generate_content_config_dict: Dict, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + _is_async: bool = False, + client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, + stream: bool = False, + litellm_metadata: Optional[Dict[str, Any]] = None, + ) -> Any: + """ + Handles Google GenAI generate content requests. + When _is_async=True, returns a coroutine instead of making the call directly. + """ + from litellm.google_genai.streaming_iterator import ( + GoogleGenAIGenerateContentStreamingIterator, + ) + + if _is_async: + return self.async_generate_content_handler( + model=model, + contents=contents, + generate_content_provider_config=generate_content_provider_config, + generate_content_config_dict=generate_content_config_dict, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + logging_obj=logging_obj, + extra_headers=extra_headers, + extra_body=extra_body, + timeout=timeout, + client=client if isinstance(client, AsyncHTTPHandler) else None, + stream=stream, + litellm_metadata=litellm_metadata, + ) + if client is None or not isinstance(client, HTTPHandler): + sync_httpx_client = _get_httpx_client( + params={"ssl_verify": litellm_params.get("ssl_verify", None)} + ) + else: + sync_httpx_client = client + + # Get headers and URL from the provider config + headers, api_base = generate_content_provider_config.sync_get_auth_token_and_url( + api_base=litellm_params.api_base, + model=model, + litellm_params=dict(litellm_params), + stream=stream, + ) + + if extra_headers: + headers.update(extra_headers) + + # Get the request body from the provider config + data = generate_content_provider_config.transform_generate_content_request( + model=model, + contents=contents, + generate_content_config_dict=generate_content_config_dict, + ) + + if extra_body: + data.update(extra_body) + + ## LOGGING + logging_obj.pre_call( + input=contents, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": headers, + }, + ) + + try: + if stream: + response = sync_httpx_client.post( + url=api_base, + headers=headers, + json=data, + timeout=timeout, + stream=True, + ) + # Return streaming iterator + return GoogleGenAIGenerateContentStreamingIterator( + response=response, + model=model, + logging_obj=logging_obj, + generate_content_provider_config=generate_content_provider_config, + litellm_metadata=litellm_metadata or {}, + custom_llm_provider=custom_llm_provider, + request_body=data, + ) + else: + response = sync_httpx_client.post( + url=api_base, + headers=headers, + json=data, + timeout=timeout, + ) + except Exception as e: + raise self._handle_error( + e=e, + provider_config=generate_content_provider_config, + ) + + return generate_content_provider_config.transform_generate_content_response( + model=model, + raw_response=response, + ) + + async def async_generate_content_handler( + self, + model: str, + contents: Any, + generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, + generate_content_config_dict: Dict, + custom_llm_provider: str, + litellm_params: GenericLiteLLMParams, + logging_obj: LiteLLMLoggingObj, + extra_headers: Optional[Dict[str, Any]] = None, + extra_body: Optional[Dict[str, Any]] = None, + timeout: Optional[Union[float, httpx.Timeout]] = None, + client: Optional[AsyncHTTPHandler] = None, + stream: bool = False, + litellm_metadata: Optional[Dict[str, Any]] = None, + ) -> Any: + """ + Async version of the generate content handler. + Uses async HTTP client to make requests. + """ + from litellm.google_genai.streaming_iterator import ( + AsyncGoogleGenAIGenerateContentStreamingIterator, + ) + + if client is None or not isinstance(client, AsyncHTTPHandler): + async_httpx_client = get_async_httpx_client( + llm_provider=litellm.LlmProviders(custom_llm_provider), + params={"ssl_verify": litellm_params.get("ssl_verify", None)}, + ) + else: + async_httpx_client = client + + # Get headers and URL from the provider config + headers, api_base = await generate_content_provider_config.get_auth_token_and_url( + model=model, + litellm_params=dict(litellm_params), + stream=stream, + api_base=litellm_params.api_base, + ) + + if extra_headers: + headers.update(extra_headers) + + # Get the request body from the provider config + data = generate_content_provider_config.transform_generate_content_request( + model=model, + contents=contents, + generate_content_config_dict=generate_content_config_dict, + ) + + if extra_body: + data.update(extra_body) + + ## LOGGING + logging_obj.pre_call( + input=contents, + api_key="", + additional_args={ + "complete_input_dict": data, + "api_base": api_base, + "headers": headers, + }, + ) + + try: + if stream: + response = await async_httpx_client.post( + url=api_base, + headers=headers, + json=data, + timeout=timeout, + stream=True, + ) + # Return async streaming iterator + return AsyncGoogleGenAIGenerateContentStreamingIterator( + response=response, + model=model, + logging_obj=logging_obj, + generate_content_provider_config=generate_content_provider_config, + litellm_metadata=litellm_metadata or {}, + custom_llm_provider=custom_llm_provider, + request_body=data, + ) + else: + response = await async_httpx_client.post( + url=api_base, + headers=headers, + json=data, + timeout=timeout, + ) + except Exception as e: + raise self._handle_error( + e=e, + provider_config=generate_content_provider_config, + ) + + return generate_content_provider_config.transform_generate_content_response( + model=model, + raw_response=response, + ) diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py new file mode 100644 index 0000000000..5751d28265 --- /dev/null +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -0,0 +1,295 @@ +""" +Transformation for Calling Google models in their native format. +""" +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast + +import httpx + +import litellm +from litellm.llms.base_llm.google_genai.transformation import ( + BaseGoogleGenAIGenerateContentConfig, +) +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams + +if TYPE_CHECKING: + from litellm.types.google_genai.main import ( + GenerateContentConfigDict, + GenerateContentContentListUnionDict, + GenerateContentResponse, + ) +else: + GenerateContentConfigDict = Any + GenerateContentContentListUnionDict = Any + GenerateContentResponse = Any + + +class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): + """ + Configuration for calling Google models in their native format. + """ + @property + def custom_llm_provider(self) -> Literal["gemini"]: + return "gemini" + + def __init__(self): + super().__init__() + VertexLLM.__init__(self) + + def get_supported_generate_content_optional_params(self, model: str) -> List[str]: + """ + Get the list of supported Google GenAI parameters for the model. + + Args: + model: The model name + + Returns: + List of supported parameter names + """ + return [ + "http_options", + "system_instruction", + "temperature", + "top_p", + "top_k", + "candidate_count", + "max_output_tokens", + "stop_sequences", + "response_logprobs", + "logprobs", + "presence_penalty", + "frequency_penalty", + "seed", + "response_mime_type", + "response_schema", + "routing_config", + "model_selection_config", + "safety_settings", + "tools", + "tool_config", + "labels", + "cached_content", + "response_modalities", + "media_resolution", + "speech_config", + "audio_timestamp", + "automatic_function_calling", + "thinking_config" + ] + + + def map_generate_content_optional_params( + self, + generate_content_config_dict: GenerateContentConfigDict, + model: str, + ) -> Dict[str, Any]: + """ + Map Google GenAI parameters to provider-specific format. + + Args: + generate_content_optional_params: Optional parameters for generate content + model: The model name + + Returns: + Mapped parameters for the provider + """ + from litellm.types.google_genai.main import GenerateContentConfigDict + _generate_content_config_dict = GenerateContentConfigDict() + supported_google_genai_params = self.get_supported_generate_content_optional_params(model) + for param, value in generate_content_config_dict.items(): + if param in supported_google_genai_params: + _generate_content_config_dict[param] = value + return dict(_generate_content_config_dict) + + def validate_environment( + self, + api_key: Optional[str], + headers: Optional[dict], + model: str, + litellm_params: Optional[Union[GenericLiteLLMParams, dict]] + ) -> dict: + default_headers = { + "Content-Type": "application/json", + } + if api_key is not None: + default_headers["Authorization"] = f"Bearer {api_key}" + if headers is not None: + default_headers.update(headers) + + return default_headers + + def _get_google_ai_studio_api_key(self, litellm_params: dict) -> Optional[str]: + return ( + litellm_params.pop("api_key", None) + or litellm_params.pop("gemini_api_key", None) + or get_secret_str("GEMINI_API_KEY") + or litellm.api_key + ) + + def _get_common_auth_components( + self, + litellm_params: dict, + ) -> Tuple[Any, Optional[str], Optional[str]]: + """ + Get common authentication components used by both sync and async methods. + + Returns: + Tuple of (vertex_credentials, vertex_project, vertex_location) + """ + vertex_credentials = self.get_vertex_ai_credentials(litellm_params) + vertex_project = self.get_vertex_ai_project(litellm_params) + vertex_location = self.get_vertex_ai_location(litellm_params) + return vertex_credentials, vertex_project, vertex_location + + def _build_final_headers_and_url( + self, + model: str, + auth_header: Optional[str], + vertex_project: Optional[str], + vertex_location: Optional[str], + vertex_credentials: Any, + stream: bool, + api_base: Optional[str], + litellm_params: dict, + ) -> Tuple[dict, str]: + """ + Build final headers and API URL from auth components. + """ + gemini_api_key = self._get_google_ai_studio_api_key(litellm_params) + + auth_header, api_base = self._get_token_and_url( + model=model, + gemini_api_key=gemini_api_key, + auth_header=auth_header, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_credentials=vertex_credentials, + stream=stream, + custom_llm_provider=self.custom_llm_provider, + api_base=api_base, + should_use_v1beta1_features=True, + ) + + headers = self.validate_environment( + api_key=auth_header, + headers=None, + model=model, + litellm_params=litellm_params, + ) + + return headers, api_base + + def sync_get_auth_token_and_url( + self, + api_base: Optional[str], + model: str, + litellm_params: dict, + stream: bool, + ) -> Tuple[dict, str]: + """ + Sync version of get_auth_token_and_url. + """ + vertex_credentials, vertex_project, vertex_location = self._get_common_auth_components(litellm_params) + + _auth_header, vertex_project = self._ensure_access_token( + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider=self.custom_llm_provider, + ) + + return self._build_final_headers_and_url( + model=model, + auth_header=_auth_header, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_credentials=vertex_credentials, + stream=stream, + api_base=api_base, + litellm_params=litellm_params, + ) + + async def get_auth_token_and_url( + self, + api_base: Optional[str], + model: str, + litellm_params: dict, + stream: bool, + ) -> Tuple[dict, str]: + """ + Get the complete URL for the request. + + Args: + api_base: Base API URL + model: The model name + litellm_params: LiteLLM parameters + + Returns: + Tuple of headers and API base + """ + vertex_credentials, vertex_project, vertex_location = self._get_common_auth_components(litellm_params) + + _auth_header, vertex_project = await self._ensure_access_token_async( + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider=self.custom_llm_provider, + ) + + return self._build_final_headers_and_url( + model=model, + auth_header=_auth_header, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_credentials=vertex_credentials, + stream=stream, + api_base=api_base, + litellm_params=litellm_params, + ) + + + def transform_generate_content_request( + self, + model: str, + contents: GenerateContentContentListUnionDict, + generate_content_config_dict: Dict, + ) -> dict: + from litellm.types.google_genai.main import ( + GenerateContentConfigDict, + GenerateContentRequestDict, + ) + typed_generate_content_request = GenerateContentRequestDict( + model=model, + contents=contents, + generationConfig=GenerateContentConfigDict(**generate_content_config_dict), + ) + + request_dict = cast(dict, typed_generate_content_request) + + return request_dict + + def transform_generate_content_response( + self, + model: str, + raw_response: httpx.Response, + ) -> GenerateContentResponse: + """ + Transform the raw response from the generate content API. + + Args: + model: The model name + raw_response: Raw HTTP response + + Returns: + Transformed response data + """ + from litellm.types.google_genai.main import GenerateContentResponse + try: + response = raw_response.json() + except Exception as e: + raise self.get_error_class( + error_message=f"Error transforming generate content response: {e}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + return GenerateContentResponse(**response) \ No newline at end of file diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index cb2a96e3bf..2133cac2c5 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -1,10 +1,8 @@ from typing import Any, Dict, List, Optional, Tuple -import litellm from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( AnthropicMessagesConfig, ) -from litellm.secret_managers.main import get_secret_str from litellm.types.llms.vertex_ai import VertexPartnerProvider from litellm.types.router import GenericLiteLLMParams @@ -28,25 +26,9 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert Validate the environment for the request """ if "Authorization" not in headers: - vertex_ai_project = ( - litellm_params.pop("vertex_project", None) - or litellm_params.pop("vertex_ai_project", None) - or litellm.vertex_project - or get_secret_str("VERTEXAI_PROJECT") - ) - vertex_credentials = ( - litellm_params.pop("vertex_credentials", None) - or litellm_params.pop("vertex_ai_credentials", None) - or get_secret_str("VERTEXAI_CREDENTIALS") - ) - - vertex_ai_location = ( - litellm_params.pop("vertex_location", None) - or litellm_params.pop("vertex_ai_location", None) - or litellm.vertex_location - or get_secret_str("VERTEXAI_LOCATION") - or get_secret_str("VERTEX_LOCATION") - ) + vertex_ai_project = VertexBase.get_vertex_ai_project(litellm_params) + vertex_credentials = VertexBase.get_vertex_ai_credentials(litellm_params) + vertex_ai_location = VertexBase.get_vertex_ai_location(litellm_params) access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index d66c498d4c..f45549368a 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -8,9 +8,11 @@ import json import os from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple +import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.secret_managers.main import get_secret_str from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES, VertexPartnerProvider from .common_utils import ( @@ -80,11 +82,9 @@ class VertexBase: # Check if the JSON object contains Workload Identity Federation configuration if "type" in json_obj and json_obj["type"] == "external_account": # If environment_id key contains "aws" value it corresponds to an AWS config file - if ( - "credential_source" in json_obj - and "environment_id" in json_obj["credential_source"] - and "aws" in json_obj["credential_source"]["environment_id"] - ): + credential_source = json_obj.get("credential_source", {}) + environment_id = credential_source.get("environment_id", "") if isinstance(credential_source, dict) else "" + if isinstance(environment_id, str) and "aws" in environment_id: creds = self._credentials_from_identity_pool_with_aws(json_obj) else: creds = self._credentials_from_identity_pool(json_obj) @@ -490,3 +490,30 @@ class VertexBase: headers.update(extra_headers) return headers + + @staticmethod + def get_vertex_ai_project(litellm_params: dict) -> Optional[str]: + return ( + litellm_params.pop("vertex_project", None) + or litellm_params.pop("vertex_ai_project", None) + or litellm.vertex_project + or get_secret_str("VERTEXAI_PROJECT") + ) + + @staticmethod + def get_vertex_ai_credentials(litellm_params: dict) -> Optional[str]: + return ( + litellm_params.pop("vertex_credentials", None) + or litellm_params.pop("vertex_ai_credentials", None) + or get_secret_str("VERTEXAI_CREDENTIALS") + ) + + @staticmethod + def get_vertex_ai_location(litellm_params: dict) -> Optional[str]: + return ( + litellm_params.pop("vertex_location", None) + or litellm_params.pop("vertex_ai_location", None) + or litellm.vertex_location + or get_secret_str("VERTEXAI_LOCATION") + or get_secret_str("VERTEX_LOCATION") + ) diff --git a/litellm/types/google_genai/__init__.py b/litellm/types/google_genai/__init__.py new file mode 100644 index 0000000000..f510f3cdbe --- /dev/null +++ b/litellm/types/google_genai/__init__.py @@ -0,0 +1,13 @@ +from .main import ( + ContentListUnion, + ContentListUnionDict, + GenerateContentConfigOrDict, + GenerateContentResponse, +) + +__all__ = [ + "ContentListUnion", + "ContentListUnionDict", + "GenerateContentConfigOrDict", + "GenerateContentResponse", +] \ No newline at end of file diff --git a/litellm/types/google_genai/main.py b/litellm/types/google_genai/main.py new file mode 100644 index 0000000000..1d69012547 --- /dev/null +++ b/litellm/types/google_genai/main.py @@ -0,0 +1,18 @@ +# Import types from the Google GenAI SDK +from typing import TYPE_CHECKING, Any, Optional, TypeAlias, TypedDict + +# During static type-checking we can rely on the real google-genai types. +from google.genai import types as _genai_types # type: ignore +from pydantic import BaseModel + +ContentListUnion = _genai_types.ContentListUnion +ContentListUnionDict = _genai_types.ContentListUnionDict +GenerateContentConfigOrDict = _genai_types.GenerateContentConfigOrDict +GenerateContentResponse = _genai_types.GenerateContentResponse + +GenerateContentContentListUnionDict = _genai_types.ContentListUnionDict +GenerateContentConfigDict = _genai_types.GenerateContentConfigDict +GenerateContentRequestParametersDict = _genai_types._GenerateContentParametersDict + +class GenerateContentRequestDict(GenerateContentRequestParametersDict): # type: ignore[misc] + generationConfig: Optional[Any] \ No newline at end of file diff --git a/litellm/utils.py b/litellm/utils.py index 05328c3750..6fb9b3e8e3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -129,6 +129,9 @@ from litellm.litellm_core_utils.redact_messages import ( from litellm.litellm_core_utils.rules import Rules from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.litellm_core_utils.token_counter import get_modified_max_tokens +from litellm.llms.base_llm.google_genai.transformation import ( + BaseGoogleGenAIGenerateContentConfig, +) from litellm.llms.bedrock.common_utils import BedrockModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.router_utils.get_retry_from_policy import ( @@ -6981,6 +6984,21 @@ class ProviderConfigManager: return AzureImageEditConfig() return None + + @staticmethod + def get_provider_google_genai_generate_content_config( + model: str, + provider: LlmProviders, + ) -> Optional[BaseGoogleGenAIGenerateContentConfig]: + if litellm.LlmProviders.GEMINI == provider: + from litellm.llms.gemini.google_genai.transformation import ( + GoogleGenAIConfig, + ) + + return GoogleGenAIConfig() + elif litellm.LlmProviders.VERTEX_AI == provider: + pass + return None def get_end_user_id_for_cost_tracking( diff --git a/requirements.txt b/requirements.txt index 905a41fd4a..41a61ced81 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ # LITELLM PROXY DEPENDENCIES # -anyio==4.5.0 # openai + http req. +anyio==4.8.0 # openai + http req. httpx==0.27.0 # Pin Httpx dependency openai==1.81.0 # openai req. fastapi==0.115.5 # server dep diff --git a/tests/unified_google_tests/base_google_test.py b/tests/unified_google_tests/base_google_test.py new file mode 100644 index 0000000000..266ce28582 --- /dev/null +++ b/tests/unified_google_tests/base_google_test.py @@ -0,0 +1,122 @@ +import asyncio +import json +import sys +import os +from typing import Any, AsyncIterator, Dict, List, Optional, Union +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.google_genai import ( + generate_content, + agenerate_content, + generate_content_stream, + agenerate_content_stream, +) +from google.genai.types import ContentDict, PartDict, GenerateContentResponse + + +class BaseGoogleGenAITest: + """Base class for Google GenAI generate content tests to reduce code duplication""" + + @property + def model_config(self) -> Dict[str, Any]: + """Override in subclasses to provide model-specific configuration""" + raise NotImplementedError("Subclasses must implement model_config") + + + def _validate_non_streaming_response(self, response: Any): + """Validate non-streaming response structure""" + # Handle type checking - response should be a dict for non-streaming + if isinstance(response, AsyncIterator): + pytest.fail("Expected non-streaming response but got AsyncIterator") + + assert isinstance(response, GenerateContentResponse), f"Expected dict response, got {type(response)}" + print(f"Response: {response.model_dump_json(indent=4)}") + + # Basic validation - adjust based on actual Google GenAI response structure + # The exact structure may vary, so we'll be flexible here + assert response is not None, "Response should not be None" + + def _validate_streaming_response(self, chunks: List[Any]): + """Validate streaming response chunks""" + assert isinstance(chunks, list), f"Expected list of chunks, got {type(chunks)}" + assert len(chunks) >= 0, "Should have at least 0 chunks" + print(f"Total chunks received: {len(chunks)}") + + @pytest.mark.parametrize("is_async", [False, True]) + @pytest.mark.asyncio + async def test_non_streaming_base(self, is_async: bool): + """Base test for non-streaming requests (parametrized for sync/async)""" + request_params = self.model_config + contents = ContentDict( + parts=[ + PartDict( + text="Hello, can you tell me a short joke?" + ) + ], + ) + litellm._turn_on_debug() + + print(f"Testing {'async' if is_async else 'sync'} non-streaming with model config: {request_params}") + print(f"Contents: {contents}") + + if is_async: + print("\n--- Testing async agenerate_content ---") + response = await agenerate_content( + contents=contents, + **request_params + ) + else: + print("\n--- Testing sync generate_content ---") + response = generate_content( + contents=contents, + **request_params + ) + + print(f"{'Async' if is_async else 'Sync'} response: {json.dumps(response, indent=2, default=str)}") + self._validate_non_streaming_response(response) + + return response + + @pytest.mark.parametrize("is_async", [False, True]) + @pytest.mark.asyncio + async def test_streaming_base(self, is_async: bool): + """Base test for streaming requests (parametrized for sync/async)""" + request_params = self.model_config + contents = ContentDict( + parts=[ + PartDict( + text="Hello, can you tell me a short joke?" + ) + ], + ) + + print(f"Testing {'async' if is_async else 'sync'} streaming with model config: {request_params}") + print(f"Contents: {contents}") + + chunks = [] + + if is_async: + print("\n--- Testing async agenerate_content_stream ---") + async for chunk in agenerate_content_stream( + contents=contents, + **request_params + ): + print(f"Async chunk: {chunk}") + chunks.append(chunk) + else: + print("\n--- Testing sync generate_content_stream ---") + for chunk in generate_content_stream( + contents=contents, + **request_params + ): + print(f"Sync chunk: {chunk}") + chunks.append(chunk) + + self._validate_streaming_response(chunks) + + return chunks diff --git a/tests/unified_google_tests/test_google_ai_studio.py b/tests/unified_google_tests/test_google_ai_studio.py new file mode 100644 index 0000000000..d662c189cf --- /dev/null +++ b/tests/unified_google_tests/test_google_ai_studio.py @@ -0,0 +1,10 @@ +from base_google_test import BaseGoogleGenAITest + +class TestGoogleGenAIStudio(BaseGoogleGenAITest): + """Test Google GenAI Studio""" + + @property + def model_config(self): + return { + "model": "gemini/gemini-1.5-flash", + } \ No newline at end of file