google-genai 1.60.0__py3-none-any.whl → 1.62.0__py3-none-any.whl

google/genai/errors.py

@@ -18,6 +18,7 @@
 from typing import Any, Callable, Optional, TYPE_CHECKING, Union
 import httpx
 import json
+import websockets
 from . import _common
 
 
@@ -69,14 +70,26 @@ class APIError(Exception):
     return obj
 
   def _get_status(self, response_json: Any) -> Any:
-    return response_json.get(
-        'status', response_json.get('error', {}).get('status', None)
-    )
+    try:
+      status = response_json.get(
+          'status', response_json.get('error', {}).get('status', None)
+      )
+      return status
+    except AttributeError:
+      # If response_json is not a dict, return close code to handle the case
+      # when encountering a websocket error.
+      return None
 
   def _get_message(self, response_json: Any) -> Any:
-    return response_json.get(
-        'message', response_json.get('error', {}).get('message', None)
-    )
+    try:
+      message = response_json.get(
+          'message', response_json.get('error', {}).get('message', None)
+      )
+      return message
+    except AttributeError:
+      # If response_json is not a dict, return it as None.
+      # This is to handle the case when encountering a websocket error.
+      return None
 
   def _get_code(self, response_json: Any) -> Any:
     return response_json.get(

google/genai/files.py

@@ -101,6 +101,17 @@ def _GetFileParameters_to_mldev(
   return to_object
 
 
+def _InternalRegisterFilesParameters_to_mldev(
+    from_object: Union[dict[str, Any], object],
+    parent_object: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
+  to_object: dict[str, Any] = {}
+  if getv(from_object, ['uris']) is not None:
+    setv(to_object, ['uris'], getv(from_object, ['uris']))
+
+  return to_object
+
+
 def _ListFilesConfig_to_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -152,17 +163,6 @@ def _ListFilesResponse_from_mldev(
   return to_object
 
 
-def _RegisterFilesParameters_to_mldev(
-    from_object: Union[dict[str, Any], object],
-    parent_object: Optional[dict[str, Any]] = None,
-) -> dict[str, Any]:
-  to_object: dict[str, Any] = {}
-  if getv(from_object, ['uris']) is not None:
-    setv(to_object, ['uris'], getv(from_object, ['uris']))
-
-  return to_object
-
-
 def _RegisterFilesResponse_from_mldev(
     from_object: Union[dict[str, Any], object],
     parent_object: Optional[dict[str, Any]] = None,
@@ -438,7 +438,7 @@ class Files(_api_module.BaseModule):
       uris: list[str],
       config: Optional[types.RegisterFilesConfigOrDict] = None,
   ) -> types.RegisterFilesResponse:
-    parameter_model = types._RegisterFilesParameters(
+    parameter_model = types._InternalRegisterFilesParameters(
         uris=uris,
         config=config,
     )
@@ -449,7 +449,7 @@ class Files(_api_module.BaseModule):
           'This method is only supported in the Gemini Developer client.'
       )
     else:
-      request_dict = _RegisterFilesParameters_to_mldev(parameter_model)
+      request_dict = _InternalRegisterFilesParameters_to_mldev(parameter_model)
       request_url_dict = request_dict.get('_url')
       if request_url_dict:
         path = 'files:register'.format_map(request_url_dict)
@@ -977,7 +977,7 @@ class AsyncFiles(_api_module.BaseModule):
       uris: list[str],
       config: Optional[types.RegisterFilesConfigOrDict] = None,
   ) -> types.RegisterFilesResponse:
-    parameter_model = types._RegisterFilesParameters(
+    parameter_model = types._InternalRegisterFilesParameters(
         uris=uris,
         config=config,
     )
@@ -988,7 +988,7 @@ class AsyncFiles(_api_module.BaseModule):
           'This method is only supported in the Gemini Developer client.'
       )
     else:
-      request_dict = _RegisterFilesParameters_to_mldev(parameter_model)
+      request_dict = _InternalRegisterFilesParameters_to_mldev(parameter_model)
       request_url_dict = request_dict.get('_url')
       if request_url_dict:
         path = 'files:register'.format_map(request_url_dict)

google/genai/live.py

@@ -26,7 +26,7 @@ import warnings
 
 import google.auth
 import pydantic
-from websockets import ConnectionClosed
+import websockets
 
 from . import _api_module
 from . import _common
@@ -41,6 +41,7 @@ from ._common import set_value_by_path as setv
 from .live_music import AsyncLiveMusic
 from .models import _Content_to_mldev
 
+ConnectionClosed = websockets.ConnectionClosed
 
 try:
   from websockets.asyncio.client import ClientConnection
@@ -534,6 +535,14 @@ class AsyncSession:
       raw_response = await self._ws.recv(decode=False)
     except TypeError:
       raw_response = await self._ws.recv()  # type: ignore[assignment]
+    except ConnectionClosed as e:
+      if e.rcvd:
+        code = e.rcvd.code
+        reason = e.rcvd.reason
+      else:
+        code = 1006
+        reason = websockets.frames.CLOSE_CODE_EXPLANATIONS.get(code, 'Abnormal closure.')
+      errors.APIError.raise_error(code, reason, None)
     if raw_response:
       try:
         response = json.loads(raw_response)
@@ -545,8 +554,11 @@ class AsyncSession:
     if self._api_client.vertexai:
       response_dict = live_converters._LiveServerMessage_from_vertex(response)
     else:
-      response_dict = response
+      response_dict = live_converters._LiveServerMessage_from_mldev(response)
 
+    if not response_dict and response:
+      # Error handling.
+      errors.APIError.raise_error(response.get('code'), response, None)
     return types.LiveServerMessage._from_response(
         response=response_dict, kwargs=parameter_model.model_dump()
     )
@@ -1093,6 +1105,14 @@ class AsyncLive(_api_module.BaseModule):
         raw_response = await ws.recv(decode=False)
       except TypeError:
         raw_response = await ws.recv()  # type: ignore[assignment]
+      except ConnectionClosed as e:
+        if e.rcvd:
+          code = e.rcvd.code
+          reason = e.rcvd.reason
+        else:
+          code = 1006
+          reason = 'Abnormal closure.'
+        errors.APIError.raise_error(code, reason, None)
       if raw_response:
         try:
           response = json.loads(raw_response)

google/genai/live_music.py

@@ -19,15 +19,18 @@ import contextlib
 import json
 import logging
 from typing import AsyncIterator
+import websockets
 
 from . import _api_module
 from . import _common
 from . import _live_converters as live_converters
 from . import _transformers as t
+from . import errors
 from . import types
 from ._api_client import BaseApiClient
 from ._common import set_value_by_path as setv
 
+ConnectionClosed = websockets.ConnectionClosed
 
 try:
   from websockets.asyncio.client import ClientConnection
@@ -122,6 +125,14 @@ class AsyncMusicSession:
       raw_response = await self._ws.recv(decode=False)
     except TypeError:
       raw_response = await self._ws.recv()  # type: ignore[assignment]
+    except ConnectionClosed as e:
+      if e.rcvd:
+        code = e.rcvd.code
+        reason = e.rcvd.reason
+      else:
+        code = 1006
+        reason = websockets.frames.CLOSE_CODE_EXPLANATIONS.get(code, 'Abnormal closure.')
+      errors.APIError.raise_error(code, reason, None)
     if raw_response:
       try:
         response = json.loads(raw_response)
@@ -134,7 +145,9 @@ class AsyncMusicSession:
       raise NotImplementedError('Live music generation is not supported in Vertex AI.')
     else:
       response_dict = response
-
+    if not response_dict and response:
+      # Error handling.
+      errors.APIError.raise_error(response.get('code'), response, None)
     return types.LiveMusicServerMessage._from_response(
         response=response_dict, kwargs=parameter_model.model_dump()
     )