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

google/genai/live.py

@@ -31,7 +31,7 @@ from . import _transformers as t
 from . import client
 from . import errors
 from . import types
-from ._api_client import ApiClient
+from ._api_client import BaseApiClient
 from ._common import experimental_warning
 from ._common import get_value_by_path as getv
 from ._common import set_value_by_path as setv
@@ -49,11 +49,12 @@ from .models import _Tool_to_mldev
 from .models import _Tool_to_vertex
 
 try:
-  from websockets.asyncio.client import ClientConnection
-  from websockets.asyncio.client import connect
+  from websockets.asyncio.client import ClientConnection # type: ignore
+  from websockets.asyncio.client import connect # type: ignore
 except ModuleNotFoundError:
-  from websockets.client import ClientConnection
-  from websockets.client import connect
+  # This try/except is for TAP, mypy complains about it which is why we have the type: ignore
+  from websockets.client import ClientConnection # type: ignore
+  from websockets.client import connect # type: ignore
 
 logger = logging.getLogger('google_genai.live')
 
@@ -66,7 +67,9 @@ _FUNCTION_RESPONSE_REQUIRES_ID = (
 class AsyncSession:
   """AsyncSession. The live module is experimental."""
 
-  def __init__(self, api_client: client.ApiClient, websocket: ClientConnection):
+  def __init__(
+      self, api_client: client.BaseApiClient, websocket: ClientConnection
+  ):
     self._api_client = api_client
     self._ws = websocket
 
@@ -217,7 +220,7 @@ class AsyncSession:
       response_dict = self._LiveServerMessage_from_mldev(response)
 
     return types.LiveServerMessage._from_response(
-        response_dict, parameter_model
+        response=response_dict, kwargs=parameter_model
     )
 
   async def _send_loop(
@@ -238,7 +241,7 @@ class AsyncSession:
       self,
       from_object: Union[dict, object],
   ) -> Dict[str, Any]:
-    to_object = {}
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['modelTurn']) is not None:
       setv(
           to_object,
@@ -258,7 +261,7 @@ class AsyncSession:
       self,
       from_object: Union[dict, object],
   ) -> Dict[str, Any]:
-    to_object = {}
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['functionCalls']) is not None:
       setv(
           to_object,
@@ -271,7 +274,7 @@ class AsyncSession:
       self,
       from_object: Union[dict, object],
   ) -> Dict[str, Any]:
-    to_object = {}
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['functionCalls']) is not None:
       setv(
           to_object,
@@ -284,7 +287,7 @@ class AsyncSession:
       self,
       from_object: Union[dict, object],
   ) -> Dict[str, Any]:
-    to_object = {}
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['serverContent']) is not None:
       setv(
           to_object,
@@ -311,7 +314,7 @@ class AsyncSession:
       self,
       from_object: Union[dict, object],
   ) -> Dict[str, Any]:
-    to_object = {}
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['modelTurn']) is not None:
       setv(
           to_object,
@@ -331,7 +334,7 @@ class AsyncSession:
       self,
       from_object: Union[dict, object],
   ) -> Dict[str, Any]:
-    to_object = {}
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['serverContent']) is not None:
       setv(
           to_object,
@@ -399,7 +402,7 @@ class AsyncSession:
             raise ValueError(_FUNCTION_RESPONSE_REQUIRES_ID)
       client_message = {'tool_response': {'function_responses': input}}
     elif isinstance(input, Sequence) and any(isinstance(c, str) for c in input):
-      to_object = {}
+      to_object: dict[str, Any] = {}
       if self._api_client.vertexai:
         contents = [
             _Content_to_vertex(self._api_client, item, to_object)
@@ -503,39 +506,35 @@ class AsyncLive(_api_module.BaseModule):
   """AsyncLive. The live module is experimental."""
 
   def _LiveSetup_to_mldev(
-      self, model: str, config: Optional[types.LiveConnectConfigOrDict] = None
+      self, model: str, config: Optional[types.LiveConnectConfig] = None
   ):
-    if isinstance(config, types.LiveConnectConfig):
-      from_object = config.model_dump(exclude_none=True)
-    else:
-      from_object = config
 
-    to_object = {}
-    if getv(from_object, ['generation_config']) is not None:
+    to_object: dict[str, Any] = {}
+    if getv(config, ['generation_config']) is not None:
       setv(
           to_object,
           ['generationConfig'],
           _GenerateContentConfig_to_mldev(
               self._api_client,
-              getv(from_object, ['generation_config']),
+              getv(config, ['generation_config']),
               to_object,
           ),
       )
-    if getv(from_object, ['response_modalities']) is not None:
+    if getv(config, ['response_modalities']) is not None:
       if getv(to_object, ['generationConfig']) is not None:
-        to_object['generationConfig']['responseModalities'] = from_object[
-            'response_modalities'
-        ]
+        to_object['generationConfig']['responseModalities'] = getv(
+            config, ['response_modalities']
+        )
       else:
         to_object['generationConfig'] = {
-            'responseModalities': from_object['response_modalities']
+            'responseModalities': getv(config, ['response_modalities'])
         }
-    if getv(from_object, ['speech_config']) is not None:
+    if getv(config, ['speech_config']) is not None:
       if getv(to_object, ['generationConfig']) is not None:
         to_object['generationConfig']['speechConfig'] = _SpeechConfig_to_mldev(
             self._api_client,
             t.t_speech_config(
-                self._api_client, getv(from_object, ['speech_config'])
+                self._api_client, getv(config, ['speech_config'])
             ),
             to_object,
         )
@@ -544,31 +543,33 @@ class AsyncLive(_api_module.BaseModule):
             'speechConfig': _SpeechConfig_to_mldev(
                 self._api_client,
                 t.t_speech_config(
-                    self._api_client, getv(from_object, ['speech_config'])
+                    self._api_client, getv(config, ['speech_config'])
                 ),
                 to_object,
             )
         }
 
-    if getv(from_object, ['system_instruction']) is not None:
+    if getv(config, ['system_instruction']) is not None:
       setv(
           to_object,
           ['systemInstruction'],
           _Content_to_mldev(
               self._api_client,
               t.t_content(
-                  self._api_client, getv(from_object, ['system_instruction'])
+                  self._api_client, getv(config, ['system_instruction'])
               ),
               to_object,
           ),
       )
-    if getv(from_object, ['tools']) is not None:
+    if getv(config, ['tools']) is not None:
       setv(
           to_object,
           ['tools'],
           [
-              _Tool_to_mldev(self._api_client, item, to_object)
-              for item in getv(from_object, ['tools'])
+              _Tool_to_mldev(
+                  self._api_client, t.t_tool(self._api_client, item), to_object
+              )
+              for item in t.t_tools(self._api_client, getv(config, ['tools']))
           ],
       )
 
@@ -577,33 +578,29 @@ class AsyncLive(_api_module.BaseModule):
     return return_value
 
   def _LiveSetup_to_vertex(
-      self, model: str, config: Optional[types.LiveConnectConfigOrDict] = None
+      self, model: str, config: Optional[types.LiveConnectConfig] = None
   ):
-    if isinstance(config, types.LiveConnectConfig):
-      from_object = config.model_dump(exclude_none=True)
-    else:
-      from_object = config
 
-    to_object = {}
+    to_object: dict[str, Any] = {}
 
-    if getv(from_object, ['generation_config']) is not None:
+    if getv(config, ['generation_config']) is not None:
       setv(
           to_object,
           ['generationConfig'],
           _GenerateContentConfig_to_vertex(
               self._api_client,
-              getv(from_object, ['generation_config']),
+              getv(config, ['generation_config']),
               to_object,
           ),
       )
-    if getv(from_object, ['response_modalities']) is not None:
+    if getv(config, ['response_modalities']) is not None:
       if getv(to_object, ['generationConfig']) is not None:
-        to_object['generationConfig']['responseModalities'] = from_object[
-            'response_modalities'
-        ]
+        to_object['generationConfig']['responseModalities'] = getv(
+            config, ['response_modalities']
+        )
       else:
         to_object['generationConfig'] = {
-            'responseModalities': from_object['response_modalities']
+            'responseModalities': getv(config, ['response_modalities'])
         }
     else:
       # Set default to AUDIO to align with MLDev API.
@@ -613,12 +610,12 @@ class AsyncLive(_api_module.BaseModule):
         to_object.update(
             {'generationConfig': {'responseModalities': ['AUDIO']}}
         )
-    if getv(from_object, ['speech_config']) is not None:
+    if getv(config, ['speech_config']) is not None:
       if getv(to_object, ['generationConfig']) is not None:
         to_object['generationConfig']['speechConfig'] = _SpeechConfig_to_vertex(
             self._api_client,
             t.t_speech_config(
-                self._api_client, getv(from_object, ['speech_config'])
+                self._api_client, getv(config, ['speech_config'])
             ),
             to_object,
         )
@@ -627,30 +624,32 @@ class AsyncLive(_api_module.BaseModule):
             'speechConfig': _SpeechConfig_to_vertex(
                 self._api_client,
                 t.t_speech_config(
-                    self._api_client, getv(from_object, ['speech_config'])
+                    self._api_client, getv(config, ['speech_config'])
                 ),
                 to_object,
             )
         }
-    if getv(from_object, ['system_instruction']) is not None:
+    if getv(config, ['system_instruction']) is not None:
       setv(
           to_object,
           ['systemInstruction'],
           _Content_to_vertex(
               self._api_client,
               t.t_content(
-                  self._api_client, getv(from_object, ['system_instruction'])
+                  self._api_client, getv(config, ['system_instruction'])
               ),
               to_object,
           ),
       )
-    if getv(from_object, ['tools']) is not None:
+    if getv(config, ['tools']) is not None:
       setv(
           to_object,
           ['tools'],
           [
-              _Tool_to_vertex(self._api_client, item, to_object)
-              for item in getv(from_object, ['tools'])
+              _Tool_to_vertex(
+                  self._api_client, t.t_tool(self._api_client, item), to_object
+              )
+              for item in t.t_tools(self._api_client, getv(config, ['tools']))
           ],
       )
 
@@ -684,16 +683,24 @@ class AsyncLive(_api_module.BaseModule):
           print(message)
     """
     base_url = self._api_client._websocket_base_url()
+    transformed_model = t.t_model(self._api_client, model)
+    # Ensure the config is a LiveConnectConfig.
+    parameter_model = types.LiveConnectConfig(**config) if isinstance(
+        config, dict
+    ) else config
+
     if self._api_client.api_key:
       api_key = self._api_client.api_key
       version = self._api_client._http_options['api_version']
       uri = f'{base_url}/ws/google.ai.generativelanguage.{version}.GenerativeService.BidiGenerateContent?key={api_key}'
       headers = self._api_client._http_options['headers']
-
-      transformed_model = t.t_model(self._api_client, model)
-      request = json.dumps(
-          self._LiveSetup_to_mldev(model=transformed_model, config=config)
+      request_dict = _common.convert_to_dict(
+          self._LiveSetup_to_mldev(
+              model=transformed_model,
+              config=parameter_model,
+          )
       )
+      request = json.dumps(request_dict)
     else:
       # Get bearer token through Application Default Credentials.
       creds, _ = google.auth.default(
@@ -713,15 +720,17 @@ class AsyncLive(_api_module.BaseModule):
       uri = f'{base_url}/ws/google.cloud.aiplatform.{version}.LlmBidiService/BidiGenerateContent'
       location = self._api_client.location
       project = self._api_client.project
-      transformed_model = t.t_model(self._api_client, model)
       if transformed_model.startswith('publishers/'):
         transformed_model = (
             f'projects/{project}/locations/{location}/' + transformed_model
         )
-
-      request = json.dumps(
-          self._LiveSetup_to_vertex(model=transformed_model, config=config)
+      request_dict = _common.convert_to_dict(
+          self._LiveSetup_to_vertex(
+              model=transformed_model,
+              config=parameter_model,
+          )
       )
+      request = json.dumps(request_dict)
 
     async with connect(uri, additional_headers=headers) as ws:
       await ws.send(request)