519 lines
17 KiB
Python
519 lines
17 KiB
Python
# Copyright 2025 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import asyncio
|
|
import os
|
|
from typing import Any, Optional, Union
|
|
|
|
import google.auth
|
|
import pydantic
|
|
|
|
from ._api_client import BaseApiClient
|
|
from ._base_url import get_base_url
|
|
from ._replay_api_client import ReplayApiClient
|
|
from .batches import AsyncBatches, Batches
|
|
from .caches import AsyncCaches, Caches
|
|
from .chats import AsyncChats, Chats
|
|
from .file_search_stores import AsyncFileSearchStores, FileSearchStores
|
|
from .files import AsyncFiles, Files
|
|
from .live import AsyncLive
|
|
from .models import AsyncModels, Models
|
|
from .operations import AsyncOperations, Operations
|
|
from .tokens import AsyncTokens, Tokens
|
|
from .tunings import AsyncTunings, Tunings
|
|
from .types import HttpOptions, HttpOptionsDict, HttpRetryOptions
|
|
|
|
import warnings
|
|
|
|
from . import _common
|
|
|
|
from ._gaos.google_genai import (
|
|
AsyncGeminiNextGenAgents,
|
|
AsyncGeminiNextGenInteractions,
|
|
AsyncGeminiNextGenWebhooks,
|
|
GeminiNextGenAgents,
|
|
GeminiNextGenInteractions,
|
|
GeminiNextGenWebhooks,
|
|
build_google_genai_async_client,
|
|
build_google_genai_client,
|
|
)
|
|
from ._gaos.sdk import AsyncGenAI as AsyncGeminiNextGenAPI
|
|
from ._gaos.sdk import GenAI as GeminiNextGenAPI
|
|
|
|
_agent_experimental_warned = False
|
|
|
|
|
|
class AsyncClient:
|
|
"""Client for making asynchronous (non-blocking) requests."""
|
|
|
|
def __init__(self, api_client: BaseApiClient):
|
|
|
|
self._api_client = api_client
|
|
self._models = AsyncModels(self._api_client)
|
|
self._tunings = AsyncTunings(self._api_client)
|
|
self._caches = AsyncCaches(self._api_client)
|
|
self._batches = AsyncBatches(self._api_client)
|
|
self._files = AsyncFiles(self._api_client)
|
|
self._file_search_stores = AsyncFileSearchStores(self._api_client)
|
|
self._live = AsyncLive(self._api_client)
|
|
self._tokens = AsyncTokens(self._api_client)
|
|
self._operations = AsyncOperations(self._api_client)
|
|
self._nextgen_client_instance: Optional[AsyncGeminiNextGenAPI] = None
|
|
self._agents: Optional[AsyncGeminiNextGenAgents] = None
|
|
self._interactions: Optional[AsyncGeminiNextGenInteractions] = None
|
|
self._webhooks: Optional[AsyncGeminiNextGenWebhooks] = None
|
|
|
|
@property
|
|
def _nextgen_client(self) -> AsyncGeminiNextGenAPI:
|
|
if self._nextgen_client_instance is None:
|
|
self._nextgen_client_instance = build_google_genai_async_client(
|
|
self._api_client
|
|
)
|
|
return self._nextgen_client_instance
|
|
|
|
@property
|
|
def interactions(self) -> AsyncGeminiNextGenInteractions:
|
|
if self._interactions is None:
|
|
self._interactions = AsyncGeminiNextGenInteractions(self._api_client)
|
|
return self._interactions
|
|
|
|
@property
|
|
def webhooks(self) -> AsyncGeminiNextGenWebhooks:
|
|
if self._webhooks is None:
|
|
self._webhooks = AsyncGeminiNextGenWebhooks(self._api_client)
|
|
return self._webhooks
|
|
|
|
@property
|
|
def agents(self) -> AsyncGeminiNextGenAgents:
|
|
global _agent_experimental_warned
|
|
if not _agent_experimental_warned:
|
|
_agent_experimental_warned = True
|
|
warnings.warn(
|
|
'Agents usage is experimental and may change in future versions.',
|
|
category=UserWarning,
|
|
stacklevel=1,
|
|
)
|
|
if self._agents is None:
|
|
self._agents = AsyncGeminiNextGenAgents(self._api_client)
|
|
return self._agents
|
|
|
|
@property
|
|
def models(self) -> AsyncModels:
|
|
return self._models
|
|
|
|
@property
|
|
def tunings(self) -> AsyncTunings:
|
|
return self._tunings
|
|
|
|
@property
|
|
def caches(self) -> AsyncCaches:
|
|
return self._caches
|
|
|
|
@property
|
|
def file_search_stores(self) -> AsyncFileSearchStores:
|
|
return self._file_search_stores
|
|
|
|
@property
|
|
def batches(self) -> AsyncBatches:
|
|
return self._batches
|
|
|
|
@property
|
|
def chats(self) -> AsyncChats:
|
|
return AsyncChats(modules=self.models)
|
|
|
|
@property
|
|
def files(self) -> AsyncFiles:
|
|
return self._files
|
|
|
|
@property
|
|
def live(self) -> AsyncLive:
|
|
return self._live
|
|
|
|
@property
|
|
def auth_tokens(self) -> AsyncTokens:
|
|
return self._tokens
|
|
|
|
@property
|
|
def operations(self) -> AsyncOperations:
|
|
return self._operations
|
|
|
|
async def aclose(self) -> None:
|
|
"""Closes the async client explicitly.
|
|
|
|
However, it doesn't close the sync client, which can be closed using the
|
|
Client.close() method or using the context manager.
|
|
|
|
Usage:
|
|
.. code-block:: python
|
|
|
|
from google.genai import Client
|
|
|
|
async_client = Client(
|
|
vertexai=True, project='my-project-id', location='us-central1'
|
|
).aio
|
|
response_1 = await async_client.models.generate_content(
|
|
model='gemini-2.0-flash',
|
|
contents='Hello World',
|
|
)
|
|
response_2 = await async_client.models.generate_content(
|
|
model='gemini-2.0-flash',
|
|
contents='Hello World',
|
|
)
|
|
# Close the client to release resources.
|
|
await async_client.aclose()
|
|
"""
|
|
await self._api_client.aclose()
|
|
|
|
async def __aenter__(self) -> 'AsyncClient':
|
|
return self
|
|
|
|
async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
|
|
del args, kwargs
|
|
await self.aclose()
|
|
|
|
def __del__(self) -> None:
|
|
try:
|
|
asyncio.get_running_loop().create_task(self.aclose())
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
class DebugConfig(pydantic.BaseModel):
|
|
"""Configuration options that change client network behavior when testing."""
|
|
|
|
client_mode: Optional[str] = pydantic.Field(
|
|
default_factory=lambda: os.getenv('GOOGLE_GENAI_CLIENT_MODE', None)
|
|
)
|
|
|
|
replays_directory: Optional[str] = pydantic.Field(
|
|
default_factory=lambda: os.getenv('GOOGLE_GENAI_REPLAYS_DIRECTORY', None)
|
|
)
|
|
|
|
replay_id: Optional[str] = pydantic.Field(
|
|
default_factory=lambda: os.getenv('GOOGLE_GENAI_REPLAY_ID', None)
|
|
)
|
|
|
|
|
|
class Client:
|
|
"""Client for making synchronous requests.
|
|
|
|
Use this client to make a request to the Gemini Developer API or Gemini
|
|
Enterprise Agent Platform (previously Vertex AI API) and then wait for the
|
|
response.
|
|
|
|
To initialize the client, provide the required arguments either directly
|
|
or by using environment variables. Gemini API users and Vertex AI users in
|
|
`api_key="your-api-key"` or by defining `GOOGLE_API_KEY="your-api-key"` as an
|
|
environment variable
|
|
|
|
Gemini Enterprise Agent Platform API users can provide inputs argument as
|
|
`enterprise=True,
|
|
project="your-project-id", location="us-central1"` or by defining
|
|
`GOOGLE_GENAI_USE_ENTERPRISE=true`, `GOOGLE_CLOUD_PROJECT` and
|
|
`GOOGLE_CLOUD_LOCATION` environment variables.
|
|
|
|
Attributes:
|
|
api_key: The `API key <https://ai.google.dev/gemini-api/docs/api-key>`_ to
|
|
use for authentication. Applies to the Gemini Developer API only.
|
|
enterprise (bool): Indicates whether the client should use the Gemini
|
|
Enterprise Agent Platform endpoints (previously Vertex AI API).
|
|
Defaults to False (uses Gemini Developer API endpoints). When
|
|
`enterprise` and `vertexai` are both set, and they have conflicting
|
|
values, a `ValueError` will be raised.
|
|
vertexai (bool): Legacy flag for `enterprise`.
|
|
credentials: The credentials to use for authentication when calling the
|
|
Gemini Enterprise Agent Platform APIs. Credentials can be obtained from
|
|
environment variables and default credentials. For more information, see
|
|
`Set up Application Default Credentials
|
|
<https://cloud.google.com/docs/authentication/provide-credentials-adc>`_.
|
|
Applies to the Vertex AI API only.
|
|
project: The `Google Cloud project ID
|
|
<https://cloud.google.com/vertex-ai/docs/start/cloud-environment>`_ to use
|
|
for quota. Can be obtained from environment variables (for example,
|
|
``GOOGLE_CLOUD_PROJECT``). Applies to the Vertex AI API only.
|
|
Find your `Google Cloud project ID
|
|
<https://cloud.google.com/resource-manager/docs/creating-managing-projects#identifying_projects>`_.
|
|
location: The `location
|
|
<https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations>`_
|
|
to send API requests to (for example, ``us-central1``). Can be obtained
|
|
from environment variables. Applies to the Vertex AI API only.
|
|
debug_config: Config settings that control network behavior of the client.
|
|
This is typically used when running test code.
|
|
http_options: Http options to use for the client. These options will be
|
|
applied to all requests made by the client. Example usage: `client =
|
|
genai.Client(http_options=types.HttpOptions(api_version='v1'))`.
|
|
|
|
Usage for the Gemini Developer API:
|
|
|
|
.. code-block:: python
|
|
|
|
from google import genai
|
|
|
|
client = genai.Client(api_key='my-api-key')
|
|
|
|
Usage for the Gemini Enterprise Agent Platform API:
|
|
|
|
.. code-block:: python
|
|
|
|
from google import genai
|
|
|
|
client = genai.Client(
|
|
enterprise=True, project='my-project-id', location='us-central1'
|
|
)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
enterprise: Optional[bool] = None,
|
|
vertexai: Optional[bool] = None,
|
|
api_key: Optional[str] = None,
|
|
credentials: Optional[google.auth.credentials.Credentials] = None,
|
|
project: Optional[str] = None,
|
|
location: Optional[str] = None,
|
|
debug_config: Optional[DebugConfig] = None,
|
|
http_options: Optional[Union[HttpOptions, HttpOptionsDict]] = None,
|
|
):
|
|
"""Initializes the client.
|
|
|
|
Args:
|
|
enterprise (bool): Indicates whether the client should use the Gemini
|
|
Enterprise Agent Platform endpoints (previously Vertex AI API).
|
|
Defaults to False (uses Gemini Developer API endpoints). When
|
|
`enterprise` and `vertexai` are both set, and they have conflicting
|
|
values, a `ValueError` will be raised.
|
|
vertexai (bool): Legacy flag for `enterprise`.
|
|
api_key (str): The `API key
|
|
<https://ai.google.dev/gemini-api/docs/api-key>`_ to use for
|
|
authentication. Applies to the Gemini Developer API only.
|
|
credentials (google.auth.credentials.Credentials): The credentials to use
|
|
for authentication when calling the Vertex AI APIs. Credentials can be
|
|
obtained from environment variables and default credentials. For more
|
|
information, see `Set up Application Default Credentials
|
|
<https://cloud.google.com/docs/authentication/provide-credentials-adc>`_.
|
|
Applies to the Vertex AI API only.
|
|
project (str): The `Google Cloud project ID
|
|
<https://cloud.google.com/vertex-ai/docs/start/cloud-environment>`_ to
|
|
use for quota. Can be obtained from environment variables (for example,
|
|
``GOOGLE_CLOUD_PROJECT``). Applies to the Vertex AI API only.
|
|
location (str): The `location
|
|
<https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations>`_
|
|
to send API requests to (for example, ``us-central1``). Can be obtained
|
|
from environment variables. Applies to the Vertex AI API only.
|
|
debug_config (DebugConfig): Config settings that control network behavior
|
|
of the client. This is typically used when running test code.
|
|
http_options (Union[HttpOptions, HttpOptionsDict]): Http options to use
|
|
for the client.
|
|
"""
|
|
|
|
self._debug_config = debug_config or DebugConfig()
|
|
|
|
if enterprise is not None and vertexai is not None and enterprise != vertexai:
|
|
raise ValueError(
|
|
'enterprise and vertexai flags have conflicting values, please set'
|
|
' enterprise value only.'
|
|
)
|
|
|
|
resolved_vertexai = enterprise if enterprise is not None else vertexai
|
|
|
|
if isinstance(http_options, dict):
|
|
http_options = HttpOptions(**http_options)
|
|
|
|
base_url = get_base_url(resolved_vertexai or False, http_options)
|
|
if base_url:
|
|
if http_options:
|
|
http_options.base_url = base_url
|
|
else:
|
|
http_options = HttpOptions(base_url=base_url)
|
|
|
|
self._api_client = self._get_api_client(
|
|
vertexai=resolved_vertexai,
|
|
api_key=api_key,
|
|
credentials=credentials,
|
|
project=project,
|
|
location=location,
|
|
debug_config=self._debug_config,
|
|
http_options=http_options,
|
|
)
|
|
|
|
self._aio = AsyncClient(self._api_client)
|
|
self._models = Models(self._api_client)
|
|
self._tunings = Tunings(self._api_client)
|
|
self._caches = Caches(self._api_client)
|
|
self._file_search_stores = FileSearchStores(self._api_client)
|
|
self._batches = Batches(self._api_client)
|
|
self._files = Files(self._api_client)
|
|
self._tokens = Tokens(self._api_client)
|
|
self._operations = Operations(self._api_client)
|
|
self._nextgen_client_instance: Optional[GeminiNextGenAPI] = None
|
|
self._agents: Optional[GeminiNextGenAgents] = None
|
|
self._interactions: Optional[GeminiNextGenInteractions] = None
|
|
self._webhooks: Optional[GeminiNextGenWebhooks] = None
|
|
|
|
@staticmethod
|
|
def _get_api_client(
|
|
vertexai: Optional[bool] = None,
|
|
api_key: Optional[str] = None,
|
|
credentials: Optional[google.auth.credentials.Credentials] = None,
|
|
project: Optional[str] = None,
|
|
location: Optional[str] = None,
|
|
debug_config: Optional[DebugConfig] = None,
|
|
http_options: Optional[HttpOptions] = None,
|
|
) -> BaseApiClient:
|
|
if debug_config and debug_config.client_mode in [
|
|
'record',
|
|
'replay',
|
|
'auto',
|
|
]:
|
|
return ReplayApiClient(
|
|
mode=debug_config.client_mode, # type: ignore[arg-type]
|
|
replay_id=debug_config.replay_id, # type: ignore[arg-type]
|
|
replays_directory=debug_config.replays_directory,
|
|
vertexai=vertexai, # type: ignore[arg-type]
|
|
api_key=api_key,
|
|
credentials=credentials,
|
|
project=project,
|
|
location=location,
|
|
http_options=http_options,
|
|
)
|
|
|
|
return BaseApiClient(
|
|
vertexai=vertexai,
|
|
api_key=api_key,
|
|
credentials=credentials,
|
|
project=project,
|
|
location=location,
|
|
http_options=http_options,
|
|
)
|
|
|
|
@property
|
|
def _nextgen_client(self) -> GeminiNextGenAPI:
|
|
if self._nextgen_client_instance is None:
|
|
self._nextgen_client_instance = build_google_genai_client(
|
|
self._api_client
|
|
)
|
|
return self._nextgen_client_instance
|
|
|
|
@property
|
|
def interactions(self) -> GeminiNextGenInteractions:
|
|
if self._interactions is None:
|
|
self._interactions = GeminiNextGenInteractions(self._api_client)
|
|
return self._interactions
|
|
|
|
@property
|
|
def webhooks(self) -> GeminiNextGenWebhooks:
|
|
if self._webhooks is None:
|
|
self._webhooks = GeminiNextGenWebhooks(self._api_client)
|
|
return self._webhooks
|
|
|
|
@property
|
|
def agents(self) -> GeminiNextGenAgents:
|
|
global _agent_experimental_warned
|
|
if not _agent_experimental_warned:
|
|
_agent_experimental_warned = True
|
|
warnings.warn(
|
|
'Agents usage is experimental and may change in future versions.',
|
|
category=UserWarning,
|
|
stacklevel=2,
|
|
)
|
|
if self._agents is None:
|
|
self._agents = GeminiNextGenAgents(self._api_client)
|
|
return self._agents
|
|
|
|
@property
|
|
def chats(self) -> Chats:
|
|
return Chats(modules=self.models)
|
|
|
|
@property
|
|
def aio(self) -> AsyncClient:
|
|
return self._aio
|
|
|
|
@property
|
|
def models(self) -> Models:
|
|
return self._models
|
|
|
|
@property
|
|
def tunings(self) -> Tunings:
|
|
return self._tunings
|
|
|
|
@property
|
|
def caches(self) -> Caches:
|
|
return self._caches
|
|
|
|
@property
|
|
def file_search_stores(self) -> FileSearchStores:
|
|
return self._file_search_stores
|
|
|
|
@property
|
|
def batches(self) -> Batches:
|
|
return self._batches
|
|
|
|
@property
|
|
def files(self) -> Files:
|
|
return self._files
|
|
|
|
@property
|
|
def auth_tokens(self) -> Tokens:
|
|
return self._tokens
|
|
|
|
@property
|
|
def operations(self) -> Operations:
|
|
return self._operations
|
|
|
|
@property
|
|
def vertexai(self) -> bool:
|
|
"""Returns whether the client is using the Vertex AI API."""
|
|
return self._api_client.vertexai or False
|
|
|
|
def close(self) -> None:
|
|
"""Closes the synchronous client explicitly.
|
|
|
|
However, it doesn't close the async client, which can be closed using the
|
|
Client.aio.aclose() method or using the async context manager.
|
|
|
|
Usage:
|
|
.. code-block:: python
|
|
|
|
from google.genai import Client
|
|
|
|
client = Client(
|
|
vertexai=True, project='my-project-id', location='us-central1'
|
|
)
|
|
response_1 = client.models.generate_content(
|
|
model='gemini-2.0-flash',
|
|
contents='Hello World',
|
|
)
|
|
response_2 = client.models.generate_content(
|
|
model='gemini-2.0-flash',
|
|
contents='Hello World',
|
|
)
|
|
# Close the client to release resources.
|
|
client.close()
|
|
"""
|
|
self._api_client.close()
|
|
|
|
def __enter__(self) -> 'Client':
|
|
return self
|
|
|
|
def __exit__(self, *args: Any, **kwargs: Any) -> None:
|
|
del args, kwargs
|
|
self.close()
|
|
|
|
def __del__(self) -> None:
|
|
try:
|
|
self.close()
|
|
except Exception:
|
|
pass
|