tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-09 08:23:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Feat] Add Support for calling Gemini/Vertex models in their native format (#12046)

Browse Source 
* 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
This commit is contained in:
Ishaan Jaff
2025-06-25 18:37:03 -07:00
committed by GitHub
parent 51c1c7bd36
commit 35e46784d3
15 changed files with 1637 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/google_genai/Readme.md
+123
View File
@@ -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.
litellm/google_genai/__init__.py
+19
View File
@@ -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",
]
litellm/google_genai/main.py
+439
View File
@@ -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,
)
litellm/google_genai/streaming_iterator.py
+118
View File
@@ -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
litellm/llms/base_llm/google_genai/transformation.py
+201
View File
@@ -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,
)
litellm/llms/custom_httpx/llm_http_handler.py
+225 -1
View File
@@ -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,
)
litellm/llms/gemini/google_genai/transformation.py
+295
View File
@@ -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)
litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py
+3 -21
View File
@@ -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,
litellm/llms/vertex_ai/vertex_llm_base.py
+32 -5
View File
@@ -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")
)
litellm/types/google_genai/__init__.py
+13
View File
@@ -0,0 +1,13 @@
from .main import (
ContentListUnion,
ContentListUnionDict,
GenerateContentConfigOrDict,
GenerateContentResponse,
)
__all__ = [
"ContentListUnion",
"ContentListUnionDict",
"GenerateContentConfigOrDict",
"GenerateContentResponse",
]
litellm/types/google_genai/main.py
+18
View File
@@ -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]
litellm/utils.py
+18
View File
@@ -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(
requirements.txt
+1 -1
View File
@@ -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
tests/unified_google_tests/base_google_test.py
+122
View File
@@ -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
tests/unified_google_tests/test_google_ai_studio.py
+10
View File
@@ -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",
}