google-genai 1.21.1__py3-none-any.whl → 1.23.0__py3-none-any.whl

google/genai/_api_client.py

@@ -60,6 +60,11 @@ from .types import HttpOptionsOrDict
 from .types import HttpResponse as SdkHttpResponse
 from .types import HttpRetryOptions
 
+try:
+  from websockets.asyncio.client import connect as ws_connect
+except ModuleNotFoundError:
+  # This try/except is for TAP, mypy complains about it which is why we have the type: ignore
+  from websockets.client import connect as ws_connect  # type: ignore
 
 has_aiohttp = False
 try:
@@ -227,11 +232,13 @@ class HttpResponse:
       headers: Union[dict[str, str], httpx.Headers, 'CIMultiDictProxy[str]'],
       response_stream: Union[Any, str] = None,
       byte_stream: Union[Any, bytes] = None,
+      session: Optional['aiohttp.ClientSession'] = None,
   ):
     self.status_code: int = 200
     self.headers = headers
     self.response_stream = response_stream
     self.byte_stream = byte_stream
+    self._session = session
 
   # Async iterator for async streaming.
   def __aiter__(self) -> 'HttpResponse':
@@ -291,16 +298,23 @@ class HttpResponse:
               chunk = chunk[len('data: ') :]
             yield json.loads(chunk)
       elif hasattr(self.response_stream, 'content'):
-        async for chunk in self.response_stream.content.iter_any():
-          # This is aiohttp.ClientResponse.
-          if chunk:
+        # This is aiohttp.ClientResponse.
+        try:
+          while True:
+            chunk = await self.response_stream.content.readline()
+            if not chunk:
+              break
             # In async streaming mode, the chunk of JSON is prefixed with
             # "data:" which we must strip before parsing.
-            if not isinstance(chunk, str):
-              chunk = chunk.decode('utf-8')
+            chunk = chunk.decode('utf-8')
             if chunk.startswith('data: '):
               chunk = chunk[len('data: ') :]
-            yield json.loads(chunk)
+            chunk = chunk.strip()
+            if chunk:
+              yield json.loads(chunk)
+        finally:
+          if hasattr(self, '_session') and self._session:
+            await self._session.close()
       else:
         raise ValueError('Error parsing streaming response.')
 
@@ -538,6 +552,7 @@ class BaseApiClient:
     # Default options for both clients.
     self._http_options.headers = {'Content-Type': 'application/json'}
     if self.api_key:
+      self.api_key = self.api_key.strip()
       if self._http_options.headers is not None:
         self._http_options.headers['x-goog-api-key'] = self.api_key
     # Update the http options with the user provided http options.
@@ -558,7 +573,10 @@ class BaseApiClient:
       # Do it once at the genai.Client level. Share among all requests.
       self._async_client_session_request_args = self._ensure_aiohttp_ssl_ctx(
           self._http_options
-      )
+      ) 
+    self._websocket_ssl_ctx = self._ensure_websocket_ssl_ctx(
+        self._http_options
+    )
 
     retry_kwargs = _retry_args(self._http_options.retry_options)
     self._retry = tenacity.Retrying(**retry_kwargs, reraise=True)
@@ -688,6 +706,63 @@ class BaseApiClient:
 
     return _maybe_set(async_args, ctx)
 
+
+  @staticmethod
+  def _ensure_websocket_ssl_ctx(options: HttpOptions) -> dict[str, Any]:
+    """Ensures the SSL context is present in the async client args.
+
+    Creates a default SSL context if one is not provided.
+
+    Args:
+      options: The http options to check for SSL context.
+
+    Returns:
+      An async aiohttp ClientSession._request args.
+    """
+
+    verify = 'ssl'  # keep it consistent with httpx.
+    async_args = options.async_client_args
+    ctx = async_args.get(verify) if async_args else None
+
+    if not ctx:
+      # Initialize the SSL context for the httpx client.
+      # Unlike requests, the aiohttp package does not automatically pull in the
+      # environment variables SSL_CERT_FILE or SSL_CERT_DIR. They need to be
+      # enabled explicitly. Instead of 'verify' at client level in httpx,
+      # aiohttp uses 'ssl' at request level.
+      ctx = ssl.create_default_context(
+          cafile=os.environ.get('SSL_CERT_FILE', certifi.where()),
+          capath=os.environ.get('SSL_CERT_DIR'),
+      )
+
+    def _maybe_set(
+        args: Optional[dict[str, Any]],
+        ctx: ssl.SSLContext,
+    ) -> dict[str, Any]:
+      """Sets the SSL context in the client args if not set.
+
+      Does not override the SSL context if it is already set.
+
+      Args:
+        args: The client args to to check for SSL context.
+        ctx: The SSL context to set.
+
+      Returns:
+        The client args with the SSL context included.
+      """
+      if not args or not args.get(verify):
+        args = (args or {}).copy()
+        args[verify] = ctx
+      # Drop the args that isn't in the aiohttp RequestOptions.
+      copied_args = args.copy()
+      for key in copied_args.copy():
+        if key not in inspect.signature(ws_connect).parameters and key != 'ssl':
+          del copied_args[key]
+      return copied_args
+
+    return _maybe_set(async_args, ctx)
+
+
   def _websocket_base_url(self) -> str:
     url_parts = urlparse(self._http_options.base_url)
     return url_parts._replace(scheme='wss').geturl()  # type: ignore[arg-type, return-value]
@@ -882,6 +957,7 @@ class BaseApiClient:
       self, http_request: HttpRequest, stream: bool = False
   ) -> HttpResponse:
     data: Optional[Union[str, bytes]] = None
+
     if self.vertexai and not self.api_key:
       http_request.headers['Authorization'] = (
           f'Bearer {await self._async_access_token()}'
@@ -912,8 +988,9 @@ class BaseApiClient:
             timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
             **self._async_client_session_request_args,
         )
+
         await errors.APIError.raise_for_async_response(response)
-        return HttpResponse(response.headers, response)
+        return HttpResponse(response.headers, response, session=session)
       else:
         # aiohttp is not available. Fall back to httpx.
         httpx_request = self._async_httpx_client.build_request(
@@ -996,14 +1073,14 @@ class BaseApiClient:
       path: str,
       request_dict: dict[str, object],
       http_options: Optional[HttpOptionsOrDict] = None,
-  ) -> Generator[Any, None, None]:
+  ) -> Generator[SdkHttpResponse, None, None]:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
 
     session_response = self._request(http_request, stream=True)
     for chunk in session_response.segments():
-      yield chunk
+      yield SdkHttpResponse(headers=session_response.headers, body=json.dumps(chunk))
 
   async def async_request(
       self,
@@ -1038,7 +1115,7 @@ class BaseApiClient:
 
     async def async_generator():  # type: ignore[no-untyped-def]
       async for chunk in response:
-        yield chunk
+        yield SdkHttpResponse(headers=response.headers, body=json.dumps(chunk))
 
     return async_generator()  # type: ignore[no-untyped-call]
 

google/genai/_common.py

@@ -16,12 +16,13 @@
 """Common utilities for the SDK."""
 
 import base64
+import collections.abc
 import datetime
 import enum
 import functools
 import logging
 import typing
-from typing import Any, Callable, Optional, Union, get_origin, get_args
+from typing import Any, Callable, Optional, FrozenSet, Union, get_args, get_origin
 import uuid
 import warnings
 
@@ -233,6 +234,179 @@ def _remove_extra_fields(
 T = typing.TypeVar('T', bound='BaseModel')
 
 
+def _pretty_repr(
+    obj: Any,
+    *,
+    indent_level: int = 0,
+    indent_delta: int = 2,
+    max_len: int = 100,
+    max_items: int = 5,
+    depth: int = 6,
+    visited: Optional[FrozenSet[int]] = None,
+) -> str:
+  """Returns a representation of the given object."""
+  if visited is None:
+    visited = frozenset()
+
+  obj_id = id(obj)
+  if obj_id in visited:
+    return '<... Circular reference ...>'
+
+  if depth < 0:
+    return '<... Max depth ...>'
+
+  visited = frozenset(list(visited) + [obj_id])
+
+  indent = ' ' * indent_level
+  next_indent_str = ' ' * (indent_level + indent_delta)
+
+  if isinstance(obj, pydantic.BaseModel):
+    cls_name = obj.__class__.__name__
+    items = []
+    # Sort fields for consistent output
+    fields = sorted(type(obj).model_fields)
+
+    for field_name in fields:
+      field_info = type(obj).model_fields[field_name]
+      if not field_info.repr:  # Respect Field(repr=False)
+        continue
+
+      try:
+        value = getattr(obj, field_name)
+      except AttributeError:
+        continue
+
+      if value is None:
+        continue
+
+      value_repr = _pretty_repr(
+          value,
+          indent_level=indent_level + indent_delta,
+          indent_delta=indent_delta,
+          max_len=max_len,
+          max_items=max_items,
+          depth=depth - 1,
+          visited=visited,
+      )
+      items.append(f'{next_indent_str}{field_name}={value_repr}')
+
+    if not items:
+      return f'{cls_name}()'
+    return f'{cls_name}(\n' + ',\n'.join(items) + f'\n{indent})'
+  elif isinstance(obj, str):
+    if '\n' in obj:
+      escaped = obj.replace('"""', '\\"\\"\\"')
+      # Indent the multi-line string block contents
+      return f'"""{escaped}"""'
+    return repr(obj)
+  elif isinstance(obj, bytes):
+    if len(obj) > max_len:
+      return f"{repr(obj[:max_len-3])[:-1]}...'"
+    return repr(obj)
+  elif isinstance(obj, collections.abc.Mapping):
+    if not obj:
+      return '{}'
+    if len(obj) > max_items:
+      return f'<dict len={len(obj)}>'
+    items = []
+    try:
+      sorted_keys = sorted(obj.keys(), key=str)
+    except TypeError:
+      sorted_keys = list(obj.keys())
+
+    for k in sorted_keys:
+      v = obj[k]
+      k_repr = _pretty_repr(
+          k,
+          indent_level=indent_level + indent_delta,
+          indent_delta=indent_delta,
+          max_len=max_len,
+          max_items=max_items,
+          depth=depth - 1,
+          visited=visited,
+      )
+      v_repr = _pretty_repr(
+          v,
+          indent_level=indent_level + indent_delta,
+          indent_delta=indent_delta,
+          max_len=max_len,
+          max_items=max_items,
+          depth=depth - 1,
+          visited=visited,
+      )
+      items.append(f'{next_indent_str}{k_repr}: {v_repr}')
+    return f'{{\n' + ',\n'.join(items) + f'\n{indent}}}'
+  elif isinstance(obj, (list, tuple, set)):
+    return _format_collection(
+        obj,
+        indent_level=indent_level,
+        indent_delta=indent_delta,
+        max_len=max_len,
+        max_items=max_items,
+        depth=depth,
+        visited=visited,
+    )
+  else:
+    # Fallback to standard repr, indenting subsequent lines only
+    raw_repr = repr(obj)
+    # Replace newlines with newline + indent
+    return raw_repr.replace('\n', f'\n{next_indent_str}')
+
+
+
+def _format_collection(
+    obj: Any,
+    *,
+    indent_level: int,
+    indent_delta: int,
+    max_len: int,
+    max_items: int,
+    depth: int,
+    visited: FrozenSet[int],
+) -> str:
+    """Formats a collection (list, tuple, set)."""
+    if isinstance(obj, list):
+        brackets = ('[', ']')
+    elif isinstance(obj, tuple):
+        brackets = ('(', ')')
+    elif isinstance(obj, set):
+        obj = list(obj)
+        if obj:
+          brackets = ('{', '}')
+        else:
+          brackets = ('set(', ')')
+    else:
+        raise ValueError(f"Unsupported collection type: {type(obj)}")
+
+    if not obj:
+        return brackets[0] + brackets[1]
+
+    indent = ' ' * indent_level
+    next_indent_str = ' ' * (indent_level + indent_delta)
+    elements = []
+    for i, elem in enumerate(obj):
+        if i >= max_items:
+            elements.append(
+                f'{next_indent_str}<... {len(obj) - max_items} more items ...>'
+            )
+            break
+        # Each element starts on a new line, fully indented
+        elements.append(
+            next_indent_str
+            + _pretty_repr(
+                elem,
+                indent_level=indent_level + indent_delta,
+                indent_delta=indent_delta,
+                max_len=max_len,
+                max_items=max_items,
+                depth=depth - 1,
+                visited=visited,
+            )
+        )
+
+    return f'{brackets[0]}\n' + ',\n'.join(elements) + "," + f'\n{indent}{brackets[1]}'
+
+
 class BaseModel(pydantic.BaseModel):
 
   model_config = pydantic.ConfigDict(
@@ -248,6 +422,12 @@ class BaseModel(pydantic.BaseModel):
       ignored_types=(typing.TypeVar,)
   )
 
+  def __repr__(self) -> str:
+    try:
+      return _pretty_repr(self)
+    except Exception:
+      return super().__repr__()
+
   @classmethod
   def _from_response(
       cls: typing.Type[T], *, response: dict[str, object], kwargs: dict[str, object]

google/genai/_extra_utils.py

@@ -25,6 +25,7 @@ import pydantic
 
 from . import _common
 from . import _mcp_utils
+from . import _transformers as t
 from . import errors
 from . import types
 from ._adapters import McpToGenAiToolAdapter
@@ -62,11 +63,37 @@ def _create_generate_content_config_model(
     return config
 
 
+def _get_gcs_uri(
+    src: Union[str, types.BatchJobSourceOrDict]
+) -> Optional[str]:
+  """Extracts the first GCS URI from the source, if available."""
+  if isinstance(src, str) and src.startswith('gs://'):
+    return src
+  elif isinstance(src, dict) and src.get('gcs_uri'):
+    return src['gcs_uri'][0] if src['gcs_uri'] else None
+  elif isinstance(src, types.BatchJobSource) and src.gcs_uri:
+    return src.gcs_uri[0] if src.gcs_uri else None
+  return None
+
+
+def _get_bigquery_uri(
+    src: Union[str, types.BatchJobSourceOrDict]
+) -> Optional[str]:
+  """Extracts the BigQuery URI from the source, if available."""
+  if isinstance(src, str) and src.startswith('bq://'):
+    return src
+  elif isinstance(src, dict) and src.get('bigquery_uri'):
+    return src['bigquery_uri']
+  elif isinstance(src, types.BatchJobSource) and src.bigquery_uri:
+    return src.bigquery_uri
+  return None
+
+
 def format_destination(
-    src: str,
+    src: Union[str, types.BatchJobSourceOrDict],
     config: Optional[types.CreateBatchJobConfigOrDict] = None,
 ) -> types.CreateBatchJobConfig:
-  """Formats the destination uri based on the source uri."""
+  """Formats the destination uri based on the source uri for Vertex AI."""
   config = (
       types._CreateBatchJobParameters(config=config).config
       or types.CreateBatchJobConfig()
@@ -78,15 +105,14 @@ def format_destination(
     config.display_name = f'genai_batch_job_{unique_name}'
 
   if not config.dest:
-    if src.startswith('gs://') and src.endswith('.jsonl'):
-      # If source uri is "gs://bucket/path/to/src.jsonl", then the destination
-      # uri prefix will be "gs://bucket/path/to/src/dest".
-      config.dest = f'{src[:-6]}/dest'
-    elif src.startswith('bq://'):
-      # If source uri is "bq://project.dataset.src", then the destination
-      # uri will be "bq://project.dataset.src_dest_TIMESTAMP_UUID".
+    gcs_source_uri = _get_gcs_uri(src)
+    bigquery_source_uri = _get_bigquery_uri(src)
+
+    if gcs_source_uri and gcs_source_uri.endswith('.jsonl'):
+      config.dest = f'{gcs_source_uri[:-6]}/dest'
+    elif bigquery_source_uri:
       unique_name = unique_name or _common.timestamped_unique_name()
-      config.dest = f'{src}_dest_{unique_name}'
+      config.dest = f'{bigquery_source_uri}_dest_{unique_name}'
     else:
       raise ValueError(f'Unsupported source: {src}')
   return config
@@ -506,3 +532,15 @@ async def parse_config_for_mcp_sessions(
         parsed_config_copy.tools.append(tool)
 
   return parsed_config_copy, mcp_to_genai_tool_adapters
+
+
+def append_chunk_contents(
+    contents: Union[types.ContentListUnion, types.ContentListUnionDict],
+    chunk: types.GenerateContentResponse,
+) -> None:
+  """Appends the contents of the chunk to the contents list."""
+  if chunk is not None and chunk.candidates is not None:
+    chunk_content = chunk.candidates[0].content
+    contents = t.t_contents(contents)  # type: ignore[assignment]
+    if isinstance(contents, list) and chunk_content is not None:
+      contents.append(chunk_content)  # type: ignore[arg-type]

google/genai/_transformers.py

@@ -26,7 +26,7 @@ import sys
 import time
 import types as builtin_types
 import typing
-from typing import Any, GenericAlias, Optional, Sequence, Union  # type: ignore[attr-defined]
+from typing import Any, GenericAlias, List, Optional, Sequence, Union  # type: ignore[attr-defined]
 from ._mcp_utils import mcp_to_gemini_tool
 
 if typing.TYPE_CHECKING:
@@ -787,15 +787,25 @@ def process_schema(
 def _process_enum(
     enum: EnumMeta, client: _api_client.BaseApiClient
 ) -> types.Schema:
+  is_integer_enum = False
+
   for member in enum:  # type: ignore
-    if not isinstance(member.value, str):
+    if isinstance(member.value, int):
+      is_integer_enum = True
+    elif not isinstance(member.value, str):
       raise TypeError(
-          f'Enum member {member.name} value must be a string, got'
+          f'Enum member {member.name} value must be a string or integer, got'
           f' {type(member.value)}'
       )
 
+  enum_to_process = enum
+  if is_integer_enum:
+    str_members = [str(member.value) for member in enum]  # type: ignore
+    str_enum = Enum(enum.__name__, str_members, type=str)  # type: ignore
+    enum_to_process = str_enum
+
   class Placeholder(pydantic.BaseModel):
-    placeholder: enum  # type: ignore[valid-type]
+    placeholder: enum_to_process  # type: ignore[valid-type]
 
   enum_schema = Placeholder.model_json_schema()
   process_schema(enum_schema, client)
@@ -944,19 +954,54 @@ def t_cached_content_name(client: _api_client.BaseApiClient, name: str) -> str:
   return _resource_name(client, name, collection_identifier='cachedContents')
 
 
-def t_batch_job_source(src: str) -> types.BatchJobSource:
-  if src.startswith('gs://'):
-    return types.BatchJobSource(
-        format='jsonl',
-        gcs_uri=[src],
-    )
-  elif src.startswith('bq://'):
-    return types.BatchJobSource(
-        format='bigquery',
-        bigquery_uri=src,
-    )
-  else:
-    raise ValueError(f'Unsupported source: {src}')
+def t_batch_job_source(
+    client: _api_client.BaseApiClient,
+    src: Union[
+        str, List[types.InlinedRequestOrDict], types.BatchJobSourceOrDict
+    ],
+) -> types.BatchJobSource:
+  if isinstance(src, dict):
+    src = types.BatchJobSource(**src)
+  if isinstance(src, types.BatchJobSource):
+    if client.vertexai:
+      if src.gcs_uri and src.bigquery_uri:
+        raise ValueError(
+            'Only one of `gcs_uri` or `bigquery_uri` can be set.'
+        )
+      elif not src.gcs_uri and not src.bigquery_uri:
+        raise ValueError(
+            'One of `gcs_uri` or `bigquery_uri` must be set.'
+        )
+    else:
+      if src.inlined_requests and src.file_name:
+        raise ValueError(
+            'Only one of `inlined_requests` or `file_name` can be set.'
+        )
+      elif not src.inlined_requests and not src.file_name:
+        raise ValueError(
+            'One of `inlined_requests` or `file_name` must be set.'
+        )
+    return src
+
+  elif isinstance(src, list):
+    return types.BatchJobSource(inlined_requests=src)
+  elif isinstance(src, str):
+    if src.startswith('gs://'):
+      return types.BatchJobSource(
+          format='jsonl',
+          gcs_uri=[src],
+      )
+    elif src.startswith('bq://'):
+      return types.BatchJobSource(
+          format='bigquery',
+          bigquery_uri=src,
+      )
+    elif src.startswith('files/'):
+      return types.BatchJobSource(
+          file_name=src,
+      )
+
+  raise ValueError(f'Unsupported source: {src}')
 
 
 def t_batch_job_destination(dest: str) -> types.BatchJobDestination:
@@ -976,10 +1021,15 @@ def t_batch_job_destination(dest: str) -> types.BatchJobDestination:
 
 def t_batch_job_name(client: _api_client.BaseApiClient, name: str) -> str:
   if not client.vertexai:
-    return name
+    mldev_pattern = r'batches/[^/]+$'
+    if re.match(mldev_pattern, name):
+      return name.split('/')[-1]
+    else:
+      raise ValueError(f'Invalid batch job name: {name}.')
+
+  vertex_pattern = r'^projects/[^/]+/locations/[^/]+/batchPredictionJobs/[^/]+$'
 
-  pattern = r'^projects/[^/]+/locations/[^/]+/batchPredictionJobs/[^/]+$'
-  if re.match(pattern, name):
+  if re.match(vertex_pattern, name):
     return name.split('/')[-1]
   elif name.isdigit():
     return name
@@ -987,6 +1037,21 @@ def t_batch_job_name(client: _api_client.BaseApiClient, name: str) -> str:
     raise ValueError(f'Invalid batch job name: {name}.')
 
 
+def t_job_state(state: str) -> str:
+  if state == 'BATCH_STATE_UNSPECIFIED':
+    return 'JOB_STATE_UNSPECIFIED'
+  elif state == 'BATCH_STATE_PENDING':
+    return 'JOB_STATE_PENDING'
+  elif state == 'BATCH_STATE_SUCCEEDED':
+    return 'JOB_STATE_SUCCEEDED'
+  elif state == 'BATCH_STATE_FAILED':
+    return 'JOB_STATE_FAILED'
+  elif state == 'BATCH_STATE_CANCELLED':
+    return 'JOB_STATE_CANCELLED'
+  else:
+    return state
+
+
 LRO_POLLING_INITIAL_DELAY_SECONDS = 1.0
 LRO_POLLING_MAXIMUM_DELAY_SECONDS = 20.0
 LRO_POLLING_TIMEOUT_SECONDS = 900.0