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

google/genai/_api_client.py

@@ -996,14 +996,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 +1038,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