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

google/genai/live.py

@@ -20,7 +20,7 @@ import base64
 import contextlib
 import json
 import logging
-from typing import AsyncIterator, Optional, Sequence, Union
+from typing import Any, AsyncIterator, Dict, Optional, Sequence, Union
 
 import google.auth
 from websockets import ConnectionClosed
@@ -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,12 +49,14 @@ 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')
 
 _FUNCTION_RESPONSE_REQUIRES_ID = (
     'FunctionResponse request must have an `id` field from the'
@@ -65,22 +67,26 @@ _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
 
   async def send(
       self,
       *,
-      input: Union[
-          types.ContentListUnion,
-          types.ContentListUnionDict,
-          types.LiveClientContentOrDict,
-          types.LiveClientRealtimeInputOrDict,
-          types.LiveClientToolResponseOrDict,
-          types.FunctionResponseOrDict,
-          Sequence[types.FunctionResponseOrDict],
-      ],
+      input: Optional[
+          Union[
+              types.ContentListUnion,
+              types.ContentListUnionDict,
+              types.LiveClientContentOrDict,
+              types.LiveClientRealtimeInputOrDict,
+              types.LiveClientToolResponseOrDict,
+              types.FunctionResponseOrDict,
+              Sequence[types.FunctionResponseOrDict],
+          ]
+      ] = None,
       end_of_turn: Optional[bool] = False,
   ):
     """Send input to the model.
@@ -214,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(
@@ -234,8 +240,8 @@ class AsyncSession:
   def _LiveServerContent_from_mldev(
       self,
       from_object: Union[dict, object],
-  ) -> dict:
-    to_object = {}
+  ) -> Dict[str, Any]:
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['modelTurn']) is not None:
       setv(
           to_object,
@@ -254,8 +260,8 @@ class AsyncSession:
   def _LiveToolCall_from_mldev(
       self,
       from_object: Union[dict, object],
-  ) -> dict:
-    to_object = {}
+  ) -> Dict[str, Any]:
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['functionCalls']) is not None:
       setv(
           to_object,
@@ -267,8 +273,8 @@ class AsyncSession:
   def _LiveToolCall_from_vertex(
       self,
       from_object: Union[dict, object],
-  ) -> dict:
-    to_object = {}
+  ) -> Dict[str, Any]:
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['functionCalls']) is not None:
       setv(
           to_object,
@@ -280,8 +286,8 @@ class AsyncSession:
   def _LiveServerMessage_from_mldev(
       self,
       from_object: Union[dict, object],
-  ) -> dict:
-    to_object = {}
+  ) -> Dict[str, Any]:
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['serverContent']) is not None:
       setv(
           to_object,
@@ -307,8 +313,8 @@ class AsyncSession:
   def _LiveServerContent_from_vertex(
       self,
       from_object: Union[dict, object],
-  ) -> dict:
-    to_object = {}
+  ) -> Dict[str, Any]:
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['modelTurn']) is not None:
       setv(
           to_object,
@@ -327,8 +333,8 @@ class AsyncSession:
   def _LiveServerMessage_from_vertex(
       self,
       from_object: Union[dict, object],
-  ) -> dict:
-    to_object = {}
+  ) -> Dict[str, Any]:
+    to_object: dict[str, Any] = {}
     if getv(from_object, ['serverContent']) is not None:
       setv(
           to_object,
@@ -354,18 +360,23 @@ class AsyncSession:
 
   def _parse_client_message(
       self,
-      input: Union[
-          types.ContentListUnion,
-          types.ContentListUnionDict,
-          types.LiveClientContentOrDict,
-          types.LiveClientRealtimeInputOrDict,
-          types.LiveClientRealtimeInputOrDict,
-          types.LiveClientToolResponseOrDict,
-          types.FunctionResponseOrDict,
-          Sequence[types.FunctionResponseOrDict],
-      ],
+      input: Optional[
+          Union[
+              types.ContentListUnion,
+              types.ContentListUnionDict,
+              types.LiveClientContentOrDict,
+              types.LiveClientRealtimeInputOrDict,
+              types.LiveClientToolResponseOrDict,
+              types.FunctionResponseOrDict,
+              Sequence[types.FunctionResponseOrDict],
+          ]
+      ] = None,
       end_of_turn: Optional[bool] = False,
-  ) -> dict:
+  ) -> Dict[str, Any]:
+
+    if not input:
+      logging.info('No input provided. Assume it is the end of turn.')
+      return {'client_content': {'turn_complete': True}}
     if isinstance(input, str):
       input = [input]
     elif isinstance(input, dict) and 'data' in input:
@@ -374,7 +385,6 @@ class AsyncSession:
         input['data'] = decoded_data
       input = [input]
     elif isinstance(input, types.Blob):
-      input.data = base64.b64encode(input.data).decode('utf-8')
       input = [input]
     elif isinstance(input, dict) and 'name' in input and 'response' in input:
       # ToolResponse.FunctionResponse
@@ -392,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)
@@ -411,7 +421,7 @@ class AsyncSession:
       if any((isinstance(b, dict) and 'data' in b) for b in input):
         pass
       elif any(isinstance(b, types.Blob) for b in input):
-        input = [b.model_dump(exclude_none=True) for b in input]
+        input = [b.model_dump(exclude_none=True, mode='json') for b in input]
       else:
         raise ValueError(
             f'Unsupported input type "{type(input)}" or input content "{input}"'
@@ -419,11 +429,21 @@ class AsyncSession:
 
       client_message = {'realtime_input': {'media_chunks': input}}
 
-    elif isinstance(input, dict) and 'content' in input:
-      # TODO(b/365983264) Add validation checks for content_update input_dict.
-      client_message = {'client_content': input}
+    elif isinstance(input, dict):
+      if 'content' in input or 'turns' in input:
+        # TODO(b/365983264) Add validation checks for content_update input_dict.
+        client_message = {'client_content': input}
+      elif 'media_chunks' in input:
+        client_message = {'realtime_input': input}
+      elif 'function_responses' in input:
+        client_message = {'tool_response': input}
+      else:
+        raise ValueError(
+          f'Unsupported input type "{type(input)}" or input content "{input}"')
     elif isinstance(input, types.LiveClientRealtimeInput):
-      client_message = {'realtime_input': input.model_dump(exclude_none=True)}
+      client_message = {
+          'realtime_input': input.model_dump(exclude_none=True, mode='json')
+      }
       if isinstance(
           client_message['realtime_input']['media_chunks'][0]['data'], bytes
       ):
@@ -436,20 +456,26 @@ class AsyncSession:
         ]
 
     elif isinstance(input, types.LiveClientContent):
-      client_message = {'client_content': input.model_dump(exclude_none=True)}
+      client_message = {
+          'client_content': input.model_dump(exclude_none=True, mode='json')
+      }
     elif isinstance(input, types.LiveClientToolResponse):
       # ToolResponse.FunctionResponse
       if not (self._api_client.vertexai) and not (
           input.function_responses[0].id
       ):
         raise ValueError(_FUNCTION_RESPONSE_REQUIRES_ID)
-      client_message = {'tool_response': input.model_dump(exclude_none=True)}
+      client_message = {
+          'tool_response': input.model_dump(exclude_none=True, mode='json')
+      }
     elif isinstance(input, types.FunctionResponse):
       if not (self._api_client.vertexai) and not (input.id):
         raise ValueError(_FUNCTION_RESPONSE_REQUIRES_ID)
       client_message = {
           'tool_response': {
-              'function_responses': [input.model_dump(exclude_none=True)]
+              'function_responses': [
+                  input.model_dump(exclude_none=True, mode='json')
+              ]
           }
       }
     elif isinstance(input, Sequence) and isinstance(
@@ -460,7 +486,7 @@ class AsyncSession:
       client_message = {
           'tool_response': {
               'function_responses': [
-                  c.model_dump(exclude_none=True) for c in input
+                  c.model_dump(exclude_none=True, mode='json') for c in input
               ]
           }
       }
@@ -480,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,
         )
@@ -521,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']))
           ],
       )
 
@@ -554,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.
@@ -590,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,
         )
@@ -604,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']))
           ],
       )
 
@@ -661,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(
@@ -682,26 +712,28 @@ class AsyncLive(_api_module.BaseModule):
       auth_req = google.auth.transport.requests.Request()
       creds.refresh(auth_req)
       bearer_token = creds.token
-      headers = {
-          'Content-Type': 'application/json',
+      headers = self._api_client._http_options['headers']
+      headers.update({
           'Authorization': 'Bearer {}'.format(bearer_token),
-      }
+      })
       version = self._api_client._http_options['api_version']
       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)
-      logging.info(await ws.recv(decode=False))
+      logger.info(await ws.recv(decode=False))
 
       yield AsyncSession(api_client=self._api_client, websocket=ws)