google-genai 1.3.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

google/genai/_extra_utils.py

@@ -17,9 +17,9 @@
 
 import inspect
 import logging
+import sys
 import typing
 from typing import Any, Callable, Dict, Optional, Union, get_args, get_origin
-import sys
 
 import pydantic
 
@@ -37,6 +37,15 @@ _DEFAULT_MAX_REMOTE_CALLS_AFC = 10
 logger = logging.getLogger('google_genai.models')
 
 
+def _create_generate_content_config_model(
+    config: types.GenerateContentConfigOrDict,
+) -> types.GenerateContentConfig:
+  if isinstance(config, dict):
+    return types.GenerateContentConfig(**config)
+  else:
+    return config
+
+
 def format_destination(
     src: str,
     config: Optional[types.CreateBatchJobConfigOrDict] = None,
@@ -69,16 +78,12 @@ def format_destination(
 
 def get_function_map(
     config: Optional[types.GenerateContentConfigOrDict] = None,
-) -> dict[str, object]:
+) -> dict[str, Callable]:
   """Returns a function map from the config."""
-  config_model = (
-      types.GenerateContentConfig(**config)
-      if config and isinstance(config, dict)
-      else config
-  )
-  function_map = {}
-  if not config_model:
+  function_map: dict[str, Callable] = {}
+  if not config:
     return function_map
+  config_model = _create_generate_content_config_model(config)
   if config_model.tools:
     for tool in config_model.tools:
       if callable(tool):
@@ -92,6 +97,16 @@ def get_function_map(
   return function_map
 
 
+def convert_number_values_for_dict_function_call_args(
+    args: dict[str, Any],
+) -> dict[str, Any]:
+  """Converts float values in dict with no decimal to integers."""
+  return {
+      key: convert_number_values_for_function_call_args(value)
+      for key, value in args.items()
+  }
+
+
 def convert_number_values_for_function_call_args(
     args: Union[dict[str, object], list[object], object],
 ) -> Union[dict[str, object], list[object], object]:
@@ -210,25 +225,35 @@ def invoke_function_from_dict_args(
 
 def get_function_response_parts(
     response: types.GenerateContentResponse,
-    function_map: dict[str, object],
+    function_map: dict[str, Callable],
 ) -> list[types.Part]:
   """Returns the function response parts from the response."""
   func_response_parts = []
-  for part in response.candidates[0].content.parts:
-    if not part.function_call:
-      continue
-    func_name = part.function_call.name
-    func = function_map[func_name]
-    args = convert_number_values_for_function_call_args(part.function_call.args)
-    try:
-      response = {'result': invoke_function_from_dict_args(args, func)}
-    except Exception as e:  # pylint: disable=broad-except
-      response = {'error': str(e)}
-    func_response = types.Part.from_function_response(
-        name=func_name, response=response
-    )
-
-    func_response_parts.append(func_response)
+  if (
+      response.candidates is not None
+      and isinstance(response.candidates[0].content, types.Content)
+      and response.candidates[0].content.parts is not None
+  ):
+    for part in response.candidates[0].content.parts:
+      if not part.function_call:
+        continue
+      func_name = part.function_call.name
+      if func_name is not None and part.function_call.args is not None:
+        func = function_map[func_name]
+        args = convert_number_values_for_dict_function_call_args(
+            part.function_call.args
+        )
+        func_response: dict[str, Any]
+        try:
+          func_response = {
+              'result': invoke_function_from_dict_args(args, func)
+          }
+        except Exception as e:  # pylint: disable=broad-except
+          func_response = {'error': str(e)}
+        func_response_part = types.Part.from_function_response(
+            name=func_name, response=func_response
+        )
+        func_response_parts.append(func_response_part)
   return func_response_parts
 
 
@@ -236,12 +261,9 @@ def should_disable_afc(
     config: Optional[types.GenerateContentConfigOrDict] = None,
 ) -> bool:
   """Returns whether automatic function calling is enabled."""
-  config_model = (
-      types.GenerateContentConfig(**config)
-      if config and isinstance(config, dict)
-      else config
-  )
-
+  if not config:
+    return False
+  config_model = _create_generate_content_config_model(config)
   # If max_remote_calls is less or equal to 0, warn and disable AFC.
   if (
       config_model
@@ -260,8 +282,7 @@ def should_disable_afc(
 
   # Default to enable AFC if not specified.
   if (
-      not config_model
-      or not config_model.automatic_function_calling
+      not config_model.automatic_function_calling
       or config_model.automatic_function_calling.disable is None
   ):
     return False
@@ -294,20 +315,17 @@ def should_disable_afc(
 def get_max_remote_calls_afc(
     config: Optional[types.GenerateContentConfigOrDict] = None,
 ) -> int:
+  if not config:
+    return _DEFAULT_MAX_REMOTE_CALLS_AFC
   """Returns the remaining remote calls for automatic function calling."""
   if should_disable_afc(config):
     raise ValueError(
         'automatic function calling is not enabled, but SDK is trying to get'
         ' max remote calls.'
     )
-  config_model = (
-      types.GenerateContentConfig(**config)
-      if config and isinstance(config, dict)
-      else config
-  )
+  config_model = _create_generate_content_config_model(config)
   if (
-      not config_model
-      or not config_model.automatic_function_calling
+      not config_model.automatic_function_calling
       or config_model.automatic_function_calling.maximum_remote_calls is None
   ):
     return _DEFAULT_MAX_REMOTE_CALLS_AFC
@@ -317,11 +335,9 @@ def get_max_remote_calls_afc(
 def should_append_afc_history(
     config: Optional[types.GenerateContentConfigOrDict] = None,
 ) -> bool:
-  config_model = (
-      types.GenerateContentConfig(**config)
-      if config and isinstance(config, dict)
-      else config
-  )
-  if not config_model or not config_model.automatic_function_calling:
+  if not config:
+    return True
+  config_model = _create_generate_content_config_model(config)
+  if not config_model.automatic_function_calling:
     return True
   return not config_model.automatic_function_calling.ignore_call_history

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
@@ -109,7 +109,8 @@ def _redact_project_location_path(path: str) -> str:
     return path
 
 
-def _redact_request_body(body: dict[str, object]) -> dict[str, object]:
+def _redact_request_body(body: dict[str, object]):
+  """Redacts fields in the request body in place."""
   for key, value in body.items():
     if isinstance(value, str):
       body[key] = _redact_project_location_path(value)
@@ -179,7 +180,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__(
@@ -302,13 +303,24 @@ class ReplayApiClient(ApiClient):
           status_code=http_response.status_code,
           sdk_response_segments=[],
       )
-    else:
+    elif isinstance(http_response, errors.APIError):
       response = ReplayResponse(
           headers=dict(http_response.response.headers),
           body_segments=[http_response._to_replay_record()],
           status_code=http_response.code,
           sdk_response_segments=[],
       )
+    elif isinstance(http_response, bytes):
+      response = ReplayResponse(
+          headers={},
+          body_segments=[],
+          byte_segments=[http_response],
+          sdk_response_segments=[],
+      )
+    else:
+      raise ValueError(
+          'Unsupported http_response type: ' + str(type(http_response))
+      )
     self.replay_session.interactions.append(
         ReplayInteraction(request=request, response=response)
     )
@@ -457,6 +469,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:
@@ -470,6 +483,43 @@ class ReplayApiClient(ApiClient):
     else:
       return self._build_response_from_replay(request).json
 
+  async def async_upload_file(
+      self,
+      file_path: Union[str, io.IOBase],
+      upload_url: str,
+      upload_size: int,
+  ) -> str:
+    if isinstance(file_path, io.IOBase):
+      offset = file_path.tell()
+      content = file_path.read()
+      file_path.seek(offset, os.SEEK_SET)
+      request = HttpRequest(
+          method='POST',
+          url='',
+          data={'bytes': base64.b64encode(content).decode('utf-8')},
+          headers={},
+      )
+    else:
+      request = HttpRequest(
+          method='POST', url='', data={'file_path': file_path}, headers={}
+      )
+    if self._should_call_api():
+      result: Union[str, HttpResponse]
+      try:
+        result = await super().async_upload_file(
+            file_path, upload_url, upload_size
+        )
+      except HTTPError as e:
+        result = HttpResponse(
+            e.response.headers, [json.dumps({'reason': e.response.reason})]
+        )
+        result.status_code = e.response.status_code
+        raise e
+      self._record_interaction(request, HttpResponse({}, [json.dumps(result)]))
+      return result
+    else:
+      return self._build_response_from_replay(request).json
+
   def _download_file_request(self, request):
     self._initialize_replay_session_if_not_loaded()
     if self._should_call_api():
@@ -485,3 +535,22 @@ class ReplayApiClient(ApiClient):
       return result
     else:
       return self._build_response_from_replay(request)
+
+  async def async_download_file(self, path: str, http_options):
+    self._initialize_replay_session_if_not_loaded()
+    request = self._build_request(
+        'get', path=path, request_dict={}, http_options=http_options
+    )
+    if self._should_call_api():
+      try:
+        result = await super().async_download_file(path, http_options)
+      except HTTPError as e:
+        result = HttpResponse(
+            e.response.headers, [json.dumps({'reason': e.response.reason})]
+        )
+        result.status_code = e.response.status_code
+        raise e
+      self._record_interaction(request, result)
+      return result
+    else:
+      return self._build_response_from_replay(request).byte_stream[0]

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