google-genai 1.27.0__py3-none-any.whl → 1.29.0__py3-none-any.whl

google/genai/__init__.py

@@ -15,6 +15,7 @@
 
 """Google Gen AI SDK"""
 
+from . import types
 from . import version
 from .client import Client
 

google/genai/_api_client.py

@@ -20,22 +20,21 @@ The BaseApiClient is intended to be a private module and is subject to change.
 """
 
 import asyncio
-from collections.abc import Awaitable, Generator
+from collections.abc import Generator
 import copy
 from dataclasses import dataclass
-import datetime
-import http
 import inspect
 import io
 import json
 import logging
 import math
 import os
+import random
 import ssl
 import sys
 import threading
 import time
-from typing import Any, AsyncIterator, Optional, TYPE_CHECKING, Tuple, Union
+from typing import Any, AsyncIterator, Iterator, Optional, Tuple, TYPE_CHECKING, Union
 from urllib.parse import urlparse
 from urllib.parse import urlunparse
 
@@ -47,7 +46,6 @@ from google.auth.credentials import Credentials
 from google.auth.transport.requests import Request
 import httpx
 from pydantic import BaseModel
-from pydantic import Field
 from pydantic import ValidationError
 import tenacity
 
@@ -55,11 +53,11 @@ from . import _common
 from . import errors
 from . import version
 from .types import HttpOptions
-from .types import HttpOptionsDict
 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:
@@ -81,6 +79,7 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger('google_genai._api_client')
 CHUNK_SIZE = 8 * 1024 * 1024  # 8 MB chunk size
+READ_BUFFER_SIZE = 2**20
 MAX_RETRY_COUNT = 3
 INITIAL_RETRY_DELAY = 1  # second
 DELAY_MULTIPLIER = 2
@@ -236,12 +235,12 @@ class HttpResponse:
       self.headers = headers
     elif isinstance(headers, httpx.Headers):
       self.headers = {
-        key: ', '.join(headers.get_list(key))
-        for key in headers.keys()}
+          key: ', '.join(headers.get_list(key)) for key in headers.keys()
+      }
     elif type(headers).__name__ == 'CIMultiDictProxy':
       self.headers = {
-          key: ', '.join(headers.getall(key))
-        for key in headers.keys()}
+          key: ', '.join(headers.getall(key)) for key in headers.keys()
+      }
 
     self.status_code: int = 200
     self.response_stream = response_stream
@@ -263,68 +262,32 @@ class HttpResponse:
   def json(self) -> Any:
     if not self.response_stream[0]:  # Empty response
       return ''
-    return json.loads(self.response_stream[0])
+    return self._load_json_from_response(self.response_stream[0])
 
   def segments(self) -> Generator[Any, None, None]:
     if isinstance(self.response_stream, list):
       # list of objects retrieved from replay or from non-streaming API.
       for chunk in self.response_stream:
-        yield json.loads(chunk) if chunk else {}
+        yield self._load_json_from_response(chunk) if chunk else {}
     elif self.response_stream is None:
       yield from []
     else:
       # Iterator of objects retrieved from the API.
-      for chunk in self.response_stream.iter_lines():  # type: ignore[union-attr]
-        if chunk:
-          # In 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')
-          if chunk.startswith('data: '):
-            chunk = chunk[len('data: ') :]
-          yield json.loads(chunk)
+      for chunk in self._iter_response_stream():
+        yield self._load_json_from_response(chunk)
 
   async def async_segments(self) -> AsyncIterator[Any]:
     if isinstance(self.response_stream, list):
       # list of objects retrieved from replay or from non-streaming API.
       for chunk in self.response_stream:
-        yield json.loads(chunk) if chunk else {}
+        yield self._load_json_from_response(chunk) if chunk else {}
     elif self.response_stream is None:
       async for c in []:  # type: ignore[attr-defined]
         yield c
     else:
       # Iterator of objects retrieved from the API.
-      if hasattr(self.response_stream, 'aiter_lines'):
-        async for chunk in self.response_stream.aiter_lines():
-          # This is httpx.Response.
-          if chunk:
-            # 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')
-            if chunk.startswith('data: '):
-              chunk = chunk[len('data: ') :]
-            yield json.loads(chunk)
-      elif hasattr(self.response_stream, 'content'):
-        # 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.
-            chunk = chunk.decode('utf-8')
-            if chunk.startswith('data: '):
-              chunk = chunk[len('data: ') :]
-            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.')
+      async for chunk in self._aiter_response_stream():
+        yield self._load_json_from_response(chunk)
 
   def byte_segments(self) -> Generator[Union[bytes, Any], None, None]:
     if isinstance(self.byte_stream, list):
@@ -343,6 +306,130 @@ class HttpResponse:
     for attribute in dir(self):
       response_payload[attribute] = copy.deepcopy(getattr(self, attribute))
 
+  def _iter_response_stream(self) -> Iterator[str]:
+    """Iterates over chunks retrieved from the API."""
+    if not isinstance(self.response_stream, httpx.Response):
+      raise TypeError(
+          'Expected self.response_stream to be an httpx.Response object, '
+          f'but got {type(self.response_stream).__name__}.'
+      )
+
+    chunk = ''
+    balance = 0
+    for line in self.response_stream.iter_lines():
+      if not line:
+        continue
+
+      # In streaming mode, the response of JSON is prefixed with "data: " which
+      # we must strip before parsing.
+      if line.startswith('data: '):
+        yield line[len('data: '):]
+        continue
+
+      # When API returns an error message, it comes line by line. So we buffer
+      # the lines until a complete JSON string is read. A complete JSON string
+      # is found when the balance is 0.
+      for c in line:
+        if c == '{':
+          balance += 1
+        elif c == '}':
+          balance -= 1
+
+      chunk += line
+      if balance == 0:
+        yield chunk
+        chunk = ''
+
+    # If there is any remaining chunk, yield it.
+    if chunk:
+      yield chunk
+
+  async def _aiter_response_stream(self) -> AsyncIterator[str]:
+    """Asynchronously iterates over chunks retrieved from the API."""
+    if not isinstance(
+        self.response_stream, (httpx.Response, aiohttp.ClientResponse)
+    ):
+      raise TypeError(
+          'Expected self.response_stream to be an httpx.Response or'
+          ' aiohttp.ClientResponse object, but got'
+          f' {type(self.response_stream).__name__}.'
+      )
+
+    chunk = ''
+    balance = 0
+    # httpx.Response has a dedicated async line iterator.
+    if isinstance(self.response_stream, httpx.Response):
+      async for line in self.response_stream.aiter_lines():
+        if not line:
+          continue
+        # In streaming mode, the response of JSON is prefixed with "data: "
+        # which we must strip before parsing.
+        if line.startswith('data: '):
+          yield line[len('data: '):]
+          continue
+
+        # When API returns an error message, it comes line by line. So we buffer
+        # the lines until a complete JSON string is read. A complete JSON string
+        # is found when the balance is 0.
+        for c in line:
+          if c == '{':
+            balance += 1
+          elif c == '}':
+            balance -= 1
+
+        chunk += line
+        if balance == 0:
+          yield chunk
+          chunk = ''
+
+    # aiohttp.ClientResponse uses a content stream that we read line by line.
+    elif isinstance(self.response_stream, aiohttp.ClientResponse):
+      while True:
+        # Read a line from the stream. This returns bytes.
+        line_bytes = await self.response_stream.content.readline()
+        if not line_bytes:
+          break
+        # Decode the bytes and remove trailing whitespace and newlines.
+        line = line_bytes.decode('utf-8').rstrip()
+        if not line:
+          continue
+
+        # In streaming mode, the response of JSON is prefixed with "data: "
+        # which we must strip before parsing.
+        if line.startswith('data: '):
+          yield line[len('data: '):]
+          continue
+
+        # When API returns an error message, it comes line by line. So we buffer
+        # the lines until a complete JSON string is read. A complete JSON string
+        # is found when the balance is 0.
+        for c in line:
+          if c == '{':
+            balance += 1
+          elif c == '}':
+            balance -= 1
+
+        chunk += line
+        if balance == 0:
+          yield chunk
+          chunk = ''
+
+    # If there is any remaining chunk, yield it.
+    if chunk:
+      yield chunk
+
+    if hasattr(self, '_session') and self._session:
+      await self._session.close()
+
+  @classmethod
+  def _load_json_from_response(cls, response: Any) -> Any:
+    """Loads JSON from the response, or raises an error if the parsing fails."""
+    try:
+      return json.loads(response)
+    except json.JSONDecodeError as e:
+      raise errors.UnknownApiResponseError(
+          f'Failed to parse response as JSON. Raw response: {response}'
+      ) from e
 
 # Default retry options.
 # The config is based on https://cloud.google.com/storage/docs/retry-strategy.
@@ -363,7 +450,7 @@ _RETRY_HTTP_STATUS_CODES = (
 )
 
 
-def _retry_args(options: Optional[HttpRetryOptions]) -> dict[str, Any]:
+def _retry_args(options: Optional[HttpRetryOptions]) -> _common.StringDict:
   """Returns the retry args for the given http retry options.
 
   Args:
@@ -533,17 +620,30 @@ class BaseApiClient:
             + ' precedence over the API key from the environment variables.'
         )
         self.api_key = None
-      if not self.project and not self.api_key:
+
+      # Skip fetching project from ADC if base url is provided in http options.
+      if (
+          not self.project
+          and not self.api_key
+          and not validated_http_options.base_url
+      ):
         credentials, self.project = _load_auth(project=None)
         if not self._credentials:
           self._credentials = credentials
-      if not ((self.project and self.location) or self.api_key):
+
+      has_sufficient_auth = (self.project and self.location) or self.api_key
+
+      if (not has_sufficient_auth and not validated_http_options.base_url):
+        # Skip sufficient auth check if base url is provided in http options.
         raise ValueError(
             'Project and location or API key must be set when using the Vertex '
             'AI API.'
         )
       if self.api_key or self.location == 'global':
         self._http_options.base_url = f'https://aiplatform.googleapis.com/'
+      elif validated_http_options.base_url and not has_sufficient_auth:
+        # Avoid setting default base url and api version if base_url provided.
+        self._http_options.base_url = validated_http_options.base_url
       else:
         self._http_options.base_url = (
             f'https://{self.location}-aiplatform.googleapis.com/'
@@ -592,7 +692,7 @@ class BaseApiClient:
   @staticmethod
   def _ensure_httpx_ssl_ctx(
       options: HttpOptions,
-  ) -> Tuple[dict[str, Any], dict[str, Any]]:
+  ) -> Tuple[_common.StringDict, _common.StringDict]:
     """Ensures the SSL context is present in the HTTPX client args.
 
     Creates a default SSL context if one is not provided.
@@ -626,9 +726,9 @@ class BaseApiClient:
       )
 
     def _maybe_set(
-        args: Optional[dict[str, Any]],
+        args: Optional[_common.StringDict],
         ctx: ssl.SSLContext,
-    ) -> dict[str, Any]:
+    ) -> _common.StringDict:
       """Sets the SSL context in the client args if not set.
 
       Does not override the SSL context if it is already set.
@@ -656,7 +756,7 @@ class BaseApiClient:
     )
 
   @staticmethod
-  def _ensure_aiohttp_ssl_ctx(options: HttpOptions) -> dict[str, Any]:
+  def _ensure_aiohttp_ssl_ctx(options: HttpOptions) -> _common.StringDict:
     """Ensures the SSL context is present in the async client args.
 
     Creates a default SSL context if one is not provided.
@@ -684,9 +784,9 @@ class BaseApiClient:
       )
 
     def _maybe_set(
-        args: Optional[dict[str, Any]],
+        args: Optional[_common.StringDict],
         ctx: ssl.SSLContext,
-    ) -> dict[str, Any]:
+    ) -> _common.StringDict:
       """Sets the SSL context in the client args if not set.
 
       Does not override the SSL context if it is already set.
@@ -714,7 +814,7 @@ class BaseApiClient:
     return _maybe_set(async_args, ctx)
 
   @staticmethod
-  def _ensure_websocket_ssl_ctx(options: HttpOptions) -> dict[str, Any]:
+  def _ensure_websocket_ssl_ctx(options: HttpOptions) -> _common.StringDict:
     """Ensures the SSL context is present in the async client args.
 
     Creates a default SSL context if one is not provided.
@@ -742,9 +842,9 @@ class BaseApiClient:
       )
 
     def _maybe_set(
-        args: Optional[dict[str, Any]],
+        args: Optional[_common.StringDict],
         ctx: ssl.SSLContext,
-    ) -> dict[str, Any]:
+    ) -> _common.StringDict:
       """Sets the SSL context in the client args if not set.
 
       Does not override the SSL context if it is already set.
@@ -864,7 +964,7 @@ class BaseApiClient:
         self.vertexai
         and not path.startswith('projects/')
         and not query_vertex_base_models
-        and not self.api_key
+        and (self.project or self.location)
     ):
       path = f'projects/{self.project}/locations/{self.location}/' + path
 
@@ -920,7 +1020,8 @@ class BaseApiClient:
       stream: bool = False,
   ) -> HttpResponse:
     data: Optional[Union[str, bytes]] = None
-    if self.vertexai and not self.api_key:
+    # If using proj/location, fetch ADC
+    if self.vertexai and (self.project or self.location):
       http_request.headers['Authorization'] = f'Bearer {self._access_token()}'
       if self._credentials and self._credentials.quota_project_id:
         http_request.headers['x-goog-user-project'] = (
@@ -963,8 +1064,21 @@ class BaseApiClient:
   def _request(
       self,
       http_request: HttpRequest,
+      http_options: Optional[HttpOptionsOrDict] = None,
       stream: bool = False,
   ) -> HttpResponse:
+    if http_options:
+      parameter_model = (
+          HttpOptions(**http_options)
+          if isinstance(http_options, dict)
+          else http_options
+      )
+      # Support per request retry options.
+      if parameter_model.retry_options:
+        retry_kwargs = _retry_args(parameter_model.retry_options)
+        retry = tenacity.Retrying(**retry_kwargs)
+        return retry(self._request_once, http_request, stream)  # type: ignore[no-any-return]
+
     return self._retry(self._request_once, http_request, stream)  # type: ignore[no-any-return]
 
   async def _async_request_once(
@@ -972,7 +1086,8 @@ class BaseApiClient:
   ) -> HttpResponse:
     data: Optional[Union[str, bytes]] = None
 
-    if self.vertexai and not self.api_key:
+    # If using proj/location, fetch ADC
+    if self.vertexai and (self.project or self.location):
       http_request.headers['Authorization'] = (
           f'Bearer {await self._async_access_token()}'
       )
@@ -993,15 +1108,43 @@ class BaseApiClient:
         session = aiohttp.ClientSession(
             headers=http_request.headers,
             trust_env=True,
+            read_bufsize=READ_BUFFER_SIZE,
         )
-        response = await session.request(
-            method=http_request.method,
-            url=http_request.url,
-            headers=http_request.headers,
-            data=data,
-            timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
-            **self._async_client_session_request_args,
-        )
+        try:
+          response = await session.request(
+              method=http_request.method,
+              url=http_request.url,
+              headers=http_request.headers,
+              data=data,
+              timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
+              **self._async_client_session_request_args,
+          )
+        except (
+            aiohttp.ClientConnectorError,
+            aiohttp.ClientConnectorDNSError,
+            aiohttp.ClientOSError,
+            aiohttp.ServerDisconnectedError,
+        ) as e:
+          await asyncio.sleep(1 + random.randint(0, 9))
+          logger.info('Retrying due to aiohttp error: %s' % e)
+          # Retrieve the SSL context from the session.
+          self._async_client_session_request_args = (
+              self._ensure_aiohttp_ssl_ctx(self._http_options)
+          )
+          # Instantiate a new session with the updated SSL context.
+          session = aiohttp.ClientSession(
+              headers=http_request.headers,
+              trust_env=True,
+              read_bufsize=READ_BUFFER_SIZE,
+          )
+          response = await session.request(
+              method=http_request.method,
+              url=http_request.url,
+              headers=http_request.headers,
+              data=data,
+              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, session=session)
@@ -1022,20 +1165,50 @@ class BaseApiClient:
         return HttpResponse(client_response.headers, client_response)
     else:
       if self._use_aiohttp():
-        async with aiohttp.ClientSession(
-            headers=http_request.headers,
-            trust_env=True,
-        ) as session:
-          response = await session.request(
-              method=http_request.method,
-              url=http_request.url,
+        try:
+          async with aiohttp.ClientSession(
               headers=http_request.headers,
-              data=data,
-              timeout=aiohttp.ClientTimeout(connect=http_request.timeout),
-              **self._async_client_session_request_args,
+              trust_env=True,
+              read_bufsize=READ_BUFFER_SIZE,
+          ) as session:
+            response = await session.request(
+                method=http_request.method,
+                url=http_request.url,
+                headers=http_request.headers,
+                data=data,
+                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, [await response.text()])
+        except (
+            aiohttp.ClientConnectorError,
+            aiohttp.ClientConnectorDNSError,
+            aiohttp.ClientOSError,
+            aiohttp.ServerDisconnectedError,
+        ) as e:
+          await asyncio.sleep(1 + random.randint(0, 9))
+          logger.info('Retrying due to aiohttp error: %s' % e)
+          # Retrieve the SSL context from the session.
+          self._async_client_session_request_args = (
+              self._ensure_aiohttp_ssl_ctx(self._http_options)
           )
-          await errors.APIError.raise_for_async_response(response)
-          return HttpResponse(response.headers, [await response.text()])
+          # Instantiate a new session with the updated SSL context.
+          async with aiohttp.ClientSession(
+              headers=http_request.headers,
+              trust_env=True,
+              read_bufsize=READ_BUFFER_SIZE,
+          ) as session:
+            response = await session.request(
+                method=http_request.method,
+                url=http_request.url,
+                headers=http_request.headers,
+                data=data,
+                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, [await response.text()])
       else:
         # aiohttp is not available. Fall back to httpx.
         client_response = await self._async_httpx_client.request(
@@ -1051,13 +1224,25 @@ class BaseApiClient:
   async def _async_request(
       self,
       http_request: HttpRequest,
+      http_options: Optional[HttpOptionsOrDict] = None,
       stream: bool = False,
   ) -> HttpResponse:
+    if http_options:
+      parameter_model = (
+          HttpOptions(**http_options)
+          if isinstance(http_options, dict)
+          else http_options
+      )
+      # Support per request retry options.
+      if parameter_model.retry_options:
+        retry_kwargs = _retry_args(parameter_model.retry_options)
+        retry = tenacity.AsyncRetrying(**retry_kwargs)
+        return await retry(self._async_request_once, http_request, stream)  # type: ignore[no-any-return]
     return await self._async_retry(  # type: ignore[no-any-return]
         self._async_request_once, http_request, stream
     )
 
-  def get_read_only_http_options(self) -> dict[str, Any]:
+  def get_read_only_http_options(self) -> _common.StringDict:
     if isinstance(self._http_options, BaseModel):
       copied = self._http_options.model_dump()
     else:
@@ -1074,7 +1259,7 @@ class BaseApiClient:
     http_request = self._build_request(
         http_method, path, request_dict, http_options
     )
-    response = self._request(http_request, stream=False)
+    response = self._request(http_request, http_options, stream=False)
     response_body = (
         response.response_stream[0] if response.response_stream else ''
     )
@@ -1091,7 +1276,7 @@ class BaseApiClient:
         http_method, path, request_dict, http_options
     )
 
-    session_response = self._request(http_request, stream=True)
+    session_response = self._request(http_request, http_options, stream=True)
     for chunk in session_response.segments():
       yield SdkHttpResponse(
           headers=session_response.headers, body=json.dumps(chunk)
@@ -1108,7 +1293,9 @@ class BaseApiClient:
         http_method, path, request_dict, http_options
     )
 
-    result = await self._async_request(http_request=http_request, stream=False)
+    result = await self._async_request(
+        http_request=http_request, http_options=http_options, stream=False
+    )
     response_body = result.response_stream[0] if result.response_stream else ''
     return SdkHttpResponse(headers=result.headers, body=response_body)
 
@@ -1340,6 +1527,7 @@ class BaseApiClient:
       async with aiohttp.ClientSession(
           headers=self._http_options.headers,
           trust_env=True,
+          read_bufsize=READ_BUFFER_SIZE,
       ) as session:
         while True:
           if isinstance(file, io.IOBase):
@@ -1523,6 +1711,7 @@ class BaseApiClient:
       async with aiohttp.ClientSession(
           headers=http_request.headers,
           trust_env=True,
+          read_bufsize=READ_BUFFER_SIZE,
       ) as session:
         response = await session.request(
             method=http_request.method,

google/genai/_automatic_function_calling_util.py

@@ -41,6 +41,39 @@ _py_builtin_type_to_schema_type = {
 }
 
 
+def _raise_for_unsupported_param(
+    param: inspect.Parameter, func_name: str, exception: Union[Exception, type[Exception]]
+) -> None:
+  raise ValueError(
+      f'Failed to parse the parameter {param} of function {func_name} for'
+      ' automatic function calling.Automatic function calling works best with'
+      ' simpler function signature schema, consider manually parsing your'
+      f' function declaration for function {func_name}.'
+  ) from exception
+
+
+def _handle_params_as_deferred_annotations(param: inspect.Parameter, annotation_under_future: dict[str, Any], name: str) -> inspect.Parameter:
+  """Catches the case when type hints are stored as strings."""
+  if isinstance(param.annotation, str):
+    param = param.replace(annotation=annotation_under_future[name])
+  return param
+
+
+def _add_unevaluated_items_to_fixed_len_tuple_schema(
+    json_schema: dict[str, Any]
+) -> dict[str, Any]:
+  if (
+      json_schema.get('maxItems')
+      and (
+          json_schema.get('prefixItems')
+          and len(json_schema['prefixItems']) == json_schema['maxItems']
+      )
+      and json_schema.get('type') == 'array'
+  ):
+    json_schema['unevaluatedItems'] = False
+  return json_schema
+
+
 def _is_builtin_primitive_or_compound(
     annotation: inspect.Parameter.annotation,  # type: ignore[valid-type]
 ) -> bool:
@@ -92,7 +125,7 @@ def _is_default_value_compatible(
   return False
 
 
-def _parse_schema_from_parameter(
+def _parse_schema_from_parameter(  # type: ignore[return]
     api_option: Literal['VERTEX_AI', 'GEMINI_API'],
     param: inspect.Parameter,
     func_name: str,
@@ -267,12 +300,7 @@ def _parse_schema_from_parameter(
       )
     schema.required = _get_required_fields(schema)
     return schema
-  raise ValueError(
-      f'Failed to parse the parameter {param} of function {func_name} for'
-      ' automatic function calling.Automatic function calling works best with'
-      ' simpler function signature schema, consider manually parsing your'
-      f' function declaration for function {func_name}.'
-  )
+  _raise_for_unsupported_param(param, func_name, ValueError)
 
 
 def _get_required_fields(schema: types.Schema) -> Optional[list[str]]: