google-genai 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

google/genai/_replay_api_client.py

@@ -29,7 +29,7 @@ import google.auth
 from requests.exceptions import HTTPError
 
 from . import errors
-from ._api_client import ApiClient
+from ._api_client import BaseApiClient
 from ._api_client import HttpOptions
 from ._api_client import HttpRequest
 from ._api_client import HttpResponse
@@ -60,6 +60,10 @@ def _redact_request_headers(headers):
       redacted_headers[header_name] = _redact_language_label(
           _redact_version_numbers(header_value)
       )
+    elif header_name.lower() == 'x-goog-user-project':
+      continue
+    elif header_name.lower() == 'authorization':
+      continue
     else:
       redacted_headers[header_name] = header_value
   return redacted_headers
@@ -175,7 +179,7 @@ class ReplayFile(BaseModel):
   interactions: list[ReplayInteraction]
 
 
-class ReplayApiClient(ApiClient):
+class ReplayApiClient(BaseApiClient):
   """For integration testing, send recorded response or records a response."""
 
   def __init__(
@@ -409,6 +413,34 @@ class ReplayApiClient(ApiClient):
     else:
       return self._build_response_from_replay(http_request)
 
+  async def _async_request(
+      self,
+      http_request: HttpRequest,
+      stream: bool = False,
+  ) -> HttpResponse:
+    self._initialize_replay_session_if_not_loaded()
+    if self._should_call_api():
+      _debug_print('api mode request: %s' % http_request)
+      try:
+        result = await super()._async_request(http_request, stream)
+      except errors.APIError as e:
+        self._record_interaction(http_request, e)
+        raise e
+      if stream:
+        result_segments = []
+        async for segment in result.async_segments():
+          result_segments.append(json.dumps(segment))
+        result = HttpResponse(result.headers, result_segments)
+        self._record_interaction(http_request, result)
+        # Need to return a RecordedResponse that rebuilds the response
+        # segments since the stream has been consumed.
+      else:
+        self._record_interaction(http_request, result)
+      _debug_print('api mode result: %s' % result.json)
+      return result
+    else:
+      return self._build_response_from_replay(http_request)
+
   def upload_file(self, file_path: Union[str, io.IOBase], upload_url: str, upload_size: int):
     if isinstance(file_path, io.IOBase):
       offset = file_path.tell()
@@ -425,6 +457,7 @@ class ReplayApiClient(ApiClient):
           method='POST', url='', data={'file_path': file_path}, headers={}
       )
     if self._should_call_api():
+      result: Union[str, HttpResponse]
       try:
         result = super().upload_file(file_path, upload_url, upload_size)
       except HTTPError as e:
@@ -453,4 +486,3 @@ class ReplayApiClient(ApiClient):
       return result
     else:
       return self._build_response_from_replay(request)
-

google/genai/_test_api_client.py

@@ -17,13 +17,13 @@ import asyncio
 import time
 from unittest.mock import MagicMock, patch
 import pytest
-from .api_client import ApiClient
+from .api_client import BaseApiClient
 
 
-@patch('genai.api_client.ApiClient._build_request')
-@patch('genai.api_client.ApiClient._request')
+@patch('genai.api_client.BaseApiClient._build_request')
+@patch('genai.api_client.BaseApiClient._request')
 def test_request_streamed_non_blocking(mock_request, mock_build_request):
-  api_client = ApiClient(api_key='test_api_key')
+  api_client = BaseApiClient(api_key='test_api_key')
   http_method = 'GET'
   path = 'test/path'
   request_dict = {'key': 'value'}
@@ -56,8 +56,8 @@ def test_request_streamed_non_blocking(mock_request, mock_build_request):
   assert end_time - start_time > 0.3
 
 
-@patch('genai.api_client.ApiClient._build_request')
-@patch('genai.api_client.ApiClient._async_request')
+@patch('genai.api_client.BaseApiClient._build_request')
+@patch('genai.api_client.BaseApiClient._async_request')
 @pytest.mark.asyncio
 async def test_async_request(mock_async_request, mock_build_request):
   api_client = ApiClient(api_key='test_api_key')
@@ -99,8 +99,8 @@ async def test_async_request(mock_async_request, mock_build_request):
   assert 0.1 <= end_time - start_time < 0.15
 
 
-@patch('genai.api_client.ApiClient._build_request')
-@patch('genai.api_client.ApiClient._async_request')
+@patch('genai.api_client.BaseApiClient._build_request')
+@patch('genai.api_client.BaseApiClient._async_request')
 @pytest.mark.asyncio
 async def test_async_request_streamed_non_blocking(
     mock_async_request, mock_build_request

google/genai/_transformers.py

@@ -20,12 +20,16 @@ from collections.abc import Iterable, Mapping
 from enum import Enum, EnumMeta
 import inspect
 import io
+import logging
 import re
 import sys
 import time
+import types as builtin_types
 import typing
 from typing import Any, GenericAlias, Optional, Union
 
+import types as builtin_types
+
 if typing.TYPE_CHECKING:
   import PIL.Image
 
@@ -34,16 +38,18 @@ import pydantic
 from . import _api_client
 from . import types
 
+logger = logging.getLogger('google_genai._transformers')
+
 if sys.version_info >= (3, 10):
-  VersionedUnionType = typing.types.UnionType
-  _UNION_TYPES = (typing.Union, typing.types.UnionType)
+  VersionedUnionType = builtin_types.UnionType
+  _UNION_TYPES = (typing.Union, builtin_types.UnionType)
 else:
   VersionedUnionType = typing._UnionGenericAlias
   _UNION_TYPES = (typing.Union,)
 
 
 def _resource_name(
-    client: _api_client.ApiClient,
+    client: _api_client.BaseApiClient,
     resource_name: str,
     *,
     collection_identifier: str,
@@ -135,7 +141,7 @@ def _resource_name(
       return resource_name
 
 
-def t_model(client: _api_client.ApiClient, model: str):
+def t_model(client: _api_client.BaseApiClient, model: str):
   if not model:
     raise ValueError('model is required.')
   if client.vertexai:
@@ -159,7 +165,7 @@ def t_model(client: _api_client.ApiClient, model: str):
       return f'models/{model}'
 
 
-def t_models_url(api_client: _api_client.ApiClient, base_models: bool) -> str:
+def t_models_url(api_client: _api_client.BaseApiClient, base_models: bool) -> str:
   if api_client.vertexai:
     if base_models:
       return 'publishers/google/models'
@@ -173,7 +179,7 @@ def t_models_url(api_client: _api_client.ApiClient, base_models: bool) -> str:
 
 
 def t_extract_models(
-    api_client: _api_client.ApiClient, response: dict
+    api_client: _api_client.BaseApiClient, response: dict
 ) -> list[types.Model]:
   if not response:
     return []
@@ -183,11 +189,18 @@ def t_extract_models(
     return response.get('tunedModels')
   elif response.get('publisherModels') is not None:
     return response.get('publisherModels')
+  elif (
+      response.get('httpHeaders') is not None
+      and response.get('jsonPayload') is None
+  ):
+    return []
   else:
-    raise ValueError('Cannot determine the models type.')
+    logger.warning('Cannot determine the models type.')
+    logger.debug('Cannot determine the models type for response: %s', response)
+    return []
 
 
-def t_caches_model(api_client: _api_client.ApiClient, model: str):
+def t_caches_model(api_client: _api_client.BaseApiClient, model: str):
   model = t_model(api_client, model)
   if not model:
     return None
@@ -203,6 +216,7 @@ def t_caches_model(api_client: _api_client.ApiClient, model: str):
 
 
 def pil_to_blob(img) -> types.Blob:
+  PngImagePlugin: Optional[builtin_types.ModuleType]
   try:
     import PIL.PngImagePlugin
 
@@ -226,10 +240,9 @@ def pil_to_blob(img) -> types.Blob:
   return types.Blob(mime_type=mime_type, data=data)
 
 
-PartType = Union[types.Part, types.PartDict, str, 'PIL.Image.Image']
-
-
-def t_part(client: _api_client.ApiClient, part: PartType) -> types.Part:
+def t_part(
+    client: _api_client.BaseApiClient, part: Optional[types.PartUnionDict]
+) -> types.Part:
   try:
     import PIL.Image
 
@@ -237,7 +250,7 @@ def t_part(client: _api_client.ApiClient, part: PartType) -> types.Part:
   except ImportError:
     PIL_Image = None
 
-  if not part:
+  if part is None:
     raise ValueError('content part is required.')
   if isinstance(part, str):
     return types.Part(text=part)
@@ -247,14 +260,19 @@ def t_part(client: _api_client.ApiClient, part: PartType) -> types.Part:
     if not part.uri or not part.mime_type:
       raise ValueError('file uri and mime_type are required.')
     return types.Part.from_uri(file_uri=part.uri, mime_type=part.mime_type)
-  else:
+  if isinstance(part, dict):
+    return types.Part.model_validate(part)
+  if isinstance(part, types.Part):
     return part
+  raise ValueError(f'Unsupported content part type: {type(part)}')
 
 
 def t_parts(
-    client: _api_client.ApiClient, parts: Union[list, PartType]
+    client: _api_client.BaseApiClient,
+    parts: Optional[Union[list[types.PartUnionDict], types.PartUnionDict]],
 ) -> list[types.Part]:
-  if parts is None:
+  #
+  if parts is None or (isinstance(parts, list) and not parts):
     raise ValueError('content parts are required.')
   if isinstance(parts, list):
     return [t_part(client, part) for part in parts]
@@ -263,7 +281,7 @@ def t_parts(
 
 
 def t_image_predictions(
-    client: _api_client.ApiClient,
+    client: _api_client.BaseApiClient,
     predictions: Optional[Iterable[Mapping[str, Any]]],
 ) -> list[types.GeneratedImage]:
   if not predictions:
@@ -282,24 +300,38 @@ def t_image_predictions(
   return images
 
 
-ContentType = Union[types.Content, types.ContentDict, PartType]
+ContentType = Union[types.Content, types.ContentDict, types.PartUnionDict]
 
 
 def t_content(
-    client: _api_client.ApiClient,
-    content: ContentType,
-):
-  if not content:
+    client: _api_client.BaseApiClient,
+    content: Optional[ContentType],
+) -> types.Content:
+  if content is None:
     raise ValueError('content is required.')
   if isinstance(content, types.Content):
     return content
   if isinstance(content, dict):
-    return types.Content.model_validate(content)
-  return types.Content(role='user', parts=t_parts(client, content))
+    try:
+      return types.Content.model_validate(content)
+    except pydantic.ValidationError:
+      possible_part = types.Part.model_validate(content)
+      return (
+          types.ModelContent(parts=[possible_part])
+          if possible_part.function_call
+          else types.UserContent(parts=[possible_part])
+      )
+  if isinstance(content, types.Part):
+    return (
+        types.ModelContent(parts=[content])
+        if content.function_call
+        else types.UserContent(parts=[content])
+    )
+  return types.UserContent(parts=content)
 
 
 def t_contents_for_embed(
-    client: _api_client.ApiClient,
+    client: _api_client.BaseApiClient,
     contents: Union[list[types.Content], list[types.ContentDict], ContentType],
 ):
   if client.vertexai and isinstance(contents, list):
@@ -314,16 +346,105 @@ def t_contents_for_embed(
 
 
 def t_contents(
-    client: _api_client.ApiClient,
-    contents: Union[list[types.Content], list[types.ContentDict], ContentType],
-):
-  if not contents:
+    client: _api_client.BaseApiClient,
+    contents: Optional[
+        Union[types.ContentListUnion, types.ContentListUnionDict]
+    ],
+) -> list[types.Content]:
+  if contents is None or (isinstance(contents, list) and not contents):
     raise ValueError('contents are required.')
-  if isinstance(contents, list):
-    return [t_content(client, content) for content in contents]
-  else:
+  if not isinstance(contents, list):
     return [t_content(client, contents)]
 
+  try:
+    import PIL.Image
+
+    PIL_Image = PIL.Image.Image
+  except ImportError:
+    PIL_Image = None
+
+  result: list[types.Content] = []
+  accumulated_parts: list[types.Part] = []
+
+  def _is_part(part: types.PartUnionDict) -> bool:
+    if (
+        isinstance(part, str)
+        or isinstance(part, types.File)
+        or (PIL_Image is not None and isinstance(part, PIL_Image))
+        or isinstance(part, types.Part)
+    ):
+      return True
+
+    if isinstance(part, dict):
+      try:
+        types.Part.model_validate(part)
+        return True
+      except pydantic.ValidationError:
+        return False
+
+    return False
+
+  def _is_user_part(part: types.Part) -> bool:
+    return not part.function_call
+
+  def _are_user_parts(parts: list[types.Part]) -> bool:
+    return all(_is_user_part(part) for part in parts)
+
+  def _append_accumulated_parts_as_content(
+      result: list[types.Content],
+      accumulated_parts: list[types.Part],
+  ):
+    if not accumulated_parts:
+      return
+    result.append(
+        types.UserContent(parts=accumulated_parts)
+        if _are_user_parts(accumulated_parts)
+        else types.ModelContent(parts=accumulated_parts)
+    )
+    accumulated_parts[:] = []
+
+  def _handle_current_part(
+      result: list[types.Content],
+      accumulated_parts: list[types.Part],
+      current_part: types.PartUnionDict,
+  ):
+    current_part = t_part(client, current_part)
+    if _is_user_part(current_part) == _are_user_parts(accumulated_parts):
+      accumulated_parts.append(current_part)
+    else:
+      _append_accumulated_parts_as_content(result, accumulated_parts)
+      accumulated_parts[:] = [current_part]
+
+  # iterator over contents
+  # if content type or content dict, append to result
+  # if consecutive part(s),
+  #   group consecutive user part(s) to a UserContent
+  #   group consecutive model part(s) to a ModelContent
+  #   append to result
+  # if list, we only accept a list of types.PartUnion
+  for content in contents:
+    if (
+        isinstance(content, types.Content)
+        # only allowed inner list is a list of types.PartUnion
+        or isinstance(content, list)
+    ):
+      _append_accumulated_parts_as_content(result, accumulated_parts)
+      if isinstance(content, list):
+        result.append(types.UserContent(parts=content))
+      else:
+        result.append(content)
+    elif (_is_part(content)): # type: ignore
+      _handle_current_part(result, accumulated_parts, content) # type: ignore
+    elif isinstance(content, dict):
+      # PactDict is already handled in _is_part
+      result.append(types.Content.model_validate(content))
+    else:
+      raise ValueError(f'Unsupported content type: {type(content)}')
+
+  _append_accumulated_parts_as_content(result, accumulated_parts)
+
+  return result
+
 
 def handle_null_fields(schema: dict[str, Any]):
   """Process null fields in the schema so it is compatible with OpenAPI.
@@ -386,7 +507,7 @@ def handle_null_fields(schema: dict[str, Any]):
 
 def process_schema(
     schema: dict[str, Any],
-    client: Optional[_api_client.ApiClient] = None,
+    client: Optional[_api_client.BaseApiClient] = None,
     defs: Optional[dict[str, Any]] = None,
     *,
     order_properties: bool = True,
@@ -549,9 +670,9 @@ def process_schema(
 
 
 def _process_enum(
-    enum: EnumMeta, client: Optional[_api_client.ApiClient] = None
+    enum: EnumMeta, client: Optional[_api_client.BaseApiClient] = None  # type: ignore
 ) -> types.Schema:
-  for member in enum:
+  for member in enum:  # type: ignore
     if not isinstance(member.value, str):
       raise TypeError(
           f'Enum member {member.name} value must be a string, got'
@@ -568,7 +689,7 @@ def _process_enum(
 
 
 def t_schema(
-    client: _api_client.ApiClient, origin: Union[types.SchemaUnionDict, Any]
+    client: _api_client.BaseApiClient, origin: Union[types.SchemaUnionDict, Any]
 ) -> Optional[types.Schema]:
   if not origin:
     return None
@@ -614,7 +735,7 @@ def t_schema(
 
 
 def t_speech_config(
-    _: _api_client.ApiClient, origin: Union[types.SpeechConfigUnionDict, Any]
+    _: _api_client.BaseApiClient, origin: Union[types.SpeechConfigUnionDict, Any]
 ) -> Optional[types.SpeechConfig]:
   if not origin:
     return None
@@ -643,7 +764,7 @@ def t_speech_config(
   raise ValueError(f'Unsupported speechConfig type: {type(origin)}')
 
 
-def t_tool(client: _api_client.ApiClient, origin) -> types.Tool:
+def t_tool(client: _api_client.BaseApiClient, origin) -> types.Tool:
   if not origin:
     return None
   if inspect.isfunction(origin) or inspect.ismethod(origin):
@@ -660,7 +781,7 @@ def t_tool(client: _api_client.ApiClient, origin) -> types.Tool:
 
 # Only support functions now.
 def t_tools(
-    client: _api_client.ApiClient, origin: list[Any]
+    client: _api_client.BaseApiClient, origin: list[Any]
 ) -> list[types.Tool]:
   if not origin:
     return []
@@ -680,11 +801,11 @@ def t_tools(
   return tools
 
 
-def t_cached_content_name(client: _api_client.ApiClient, name: str):
+def t_cached_content_name(client: _api_client.BaseApiClient, name: str):
   return _resource_name(client, name, collection_identifier='cachedContents')
 
 
-def t_batch_job_source(client: _api_client.ApiClient, src: str):
+def t_batch_job_source(client: _api_client.BaseApiClient, src: str):
   if src.startswith('gs://'):
     return types.BatchJobSource(
         format='jsonl',
@@ -699,7 +820,7 @@ def t_batch_job_source(client: _api_client.ApiClient, src: str):
     raise ValueError(f'Unsupported source: {src}')
 
 
-def t_batch_job_destination(client: _api_client.ApiClient, dest: str):
+def t_batch_job_destination(client: _api_client.BaseApiClient, dest: str):
   if dest.startswith('gs://'):
     return types.BatchJobDestination(
         format='jsonl',
@@ -714,7 +835,7 @@ def t_batch_job_destination(client: _api_client.ApiClient, dest: str):
     raise ValueError(f'Unsupported destination: {dest}')
 
 
-def t_batch_job_name(client: _api_client.ApiClient, name: str):
+def t_batch_job_name(client: _api_client.BaseApiClient, name: str):
   if not client.vertexai:
     return name
 
@@ -733,7 +854,7 @@ LRO_POLLING_TIMEOUT_SECONDS = 900.0
 LRO_POLLING_MULTIPLIER = 1.5
 
 
-def t_resolve_operation(api_client: _api_client.ApiClient, struct: dict):
+def t_resolve_operation(api_client: _api_client.BaseApiClient, struct: dict):
   if (name := struct.get('name')) and '/operations/' in name:
     operation: dict[str, Any] = struct
     total_seconds = 0.0
@@ -742,7 +863,7 @@ def t_resolve_operation(api_client: _api_client.ApiClient, struct: dict):
       if total_seconds > LRO_POLLING_TIMEOUT_SECONDS:
         raise RuntimeError(f'Operation {name} timed out.\n{operation}')
       # TODO(b/374433890): Replace with LRO module once it's available.
-      operation: dict[str, Any] = api_client.request(
+      operation = api_client.request(
           http_method='GET', path=name, request_dict={}
       )
       time.sleep(delay_seconds)
@@ -762,7 +883,7 @@ def t_resolve_operation(api_client: _api_client.ApiClient, struct: dict):
 
 
 def t_file_name(
-    api_client: _api_client.ApiClient, name: Union[str, types.File]
+    api_client: _api_client.BaseApiClient, name: Optional[Union[str, types.File]]
 ):
   # Remove the files/ prefix since it's added to the url path.
   if isinstance(name, types.File):
@@ -784,7 +905,7 @@ def t_file_name(
 
 
 def t_tuning_job_status(
-    api_client: _api_client.ApiClient, status: str
+    api_client: _api_client.BaseApiClient, status: str
 ) -> types.JobState:
   if status == 'STATE_UNSPECIFIED':
     return 'JOB_STATE_UNSPECIFIED'
@@ -802,7 +923,7 @@ def t_tuning_job_status(
 # We shouldn't use this transformer if the backend adhere to Cloud Type
 # format https://cloud.google.com/docs/discovery/type-format.
 # TODO(b/389133914,b/390320301): Remove the hack after backend fix the issue.
-def t_bytes(api_client: _api_client.ApiClient, data: bytes) -> str:
+def t_bytes(api_client: _api_client.BaseApiClient, data: bytes) -> str:
   if not isinstance(data, bytes):
     return data
   return base64.b64encode(data).decode('ascii')