google-genai 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

google/genai/live.py

@@ -68,6 +68,7 @@ class AsyncSession:
 
   async def send(
       self,
+      *,
       input: Union[
           types.ContentListUnion,
           types.ContentListUnionDict,
@@ -80,6 +81,25 @@ class AsyncSession:
       ],
       end_of_turn: Optional[bool] = False,
   ):
+    """Send input to the model.
+
+    The method will send the input request to the server.
+
+    Args:
+      input: The input request to the model.
+      end_of_turn: Whether the input is the last message in a turn.
+
+    Example usage:
+
+    .. code-block:: python
+
+      client = genai.Client(api_key=API_KEY)
+
+      async with client.aio.live.connect(model='...') as session:
+        await session.send(input='Hello world!', end_of_turn=True)
+        async for message in session.receive():
+          print(message)
+    """
     client_message = self._parse_client_message(input, end_of_turn)
     await self._ws.send(json.dumps(client_message))
 
@@ -113,7 +133,7 @@ class AsyncSession:
       yield result
 
   async def start_stream(
-      self, stream: AsyncIterator[bytes], mime_type: str
+      self, *, stream: AsyncIterator[bytes], mime_type: str
   ) -> AsyncIterator[types.LiveServerMessage]:
     """start a live session from a data stream.
 
@@ -199,7 +219,7 @@ class AsyncSession:
   ):
     async for data in data_stream:
       input = {'data': data, 'mimeType': mime_type}
-      await self.send(input)
+      await self.send(input=input)
       # Give a chance for the receive loop to process responses.
       await asyncio.sleep(10**-12)
     # Give a chance for the receiver to process the last response.
@@ -221,6 +241,8 @@ class AsyncSession:
       )
     if getv(from_object, ['turnComplete']) is not None:
       setv(to_object, ['turn_complete'], getv(from_object, ['turnComplete']))
+    if getv(from_object, ['interrupted']) is not None:
+      setv(to_object, ['interrupted'], getv(from_object, ['interrupted']))
     return to_object
 
   def _LiveToolCall_from_mldev(
@@ -292,6 +314,8 @@ class AsyncSession:
       )
     if getv(from_object, ['turnComplete']) is not None:
       setv(to_object, ['turn_complete'], getv(from_object, ['turnComplete']))
+    if getv(from_object, ['interrupted']) is not None:
+      setv(to_object, ['interrupted'], getv(from_object, ['interrupted']))
     return to_object
 
   def _LiveServerMessage_from_vertex(
@@ -338,7 +362,7 @@ class AsyncSession:
   ) -> dict:
     if isinstance(input, str):
       input = [input]
-    elif (isinstance(input, dict) and 'data' in input):
+    elif isinstance(input, dict) and 'data' in input:
       if isinstance(input['data'], bytes):
         decoded_data = base64.b64encode(input['data']).decode('utf-8')
         input['data'] = decoded_data
@@ -405,7 +429,9 @@ class AsyncSession:
       client_message = {'client_content': input.model_dump(exclude_none=True)}
     elif isinstance(input, types.LiveClientToolResponse):
       # ToolResponse.FunctionResponse
-      if not (self._api_client.vertexai) and not (input.function_responses[0].id):
+      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)}
     elif isinstance(input, types.FunctionResponse):
@@ -457,7 +483,7 @@ class AsyncLive(_common.BaseModule):
           to_object,
           ['generationConfig'],
           _GenerateContentConfig_to_mldev(
-              self.api_client,
+              self._api_client,
               getv(from_object, ['generation_config']),
               to_object,
           ),
@@ -474,17 +500,18 @@ class AsyncLive(_common.BaseModule):
     if getv(from_object, ['speech_config']) is not None:
       if getv(to_object, ['generationConfig']) is not None:
         to_object['generationConfig']['speechConfig'] = _SpeechConfig_to_mldev(
-            self.api_client,
+            self._api_client,
             t.t_speech_config(
-                self.api_client, getv(from_object, ['speech_config'])),
+                self._api_client, getv(from_object, ['speech_config'])
+            ),
             to_object,
         )
       else:
         to_object['generationConfig'] = {
             'speechConfig': _SpeechConfig_to_mldev(
-                self.api_client,
+                self._api_client,
                 t.t_speech_config(
-                    self.api_client, getv(from_object, ['speech_config'])
+                    self._api_client, getv(from_object, ['speech_config'])
                 ),
                 to_object,
             )
@@ -495,9 +522,9 @@ class AsyncLive(_common.BaseModule):
           to_object,
           ['systemInstruction'],
           _Content_to_mldev(
-              self.api_client,
+              self._api_client,
               t.t_content(
-                  self.api_client, getv(from_object, ['system_instruction'])
+                  self._api_client, getv(from_object, ['system_instruction'])
               ),
               to_object,
           ),
@@ -507,7 +534,7 @@ class AsyncLive(_common.BaseModule):
           to_object,
           ['tools'],
           [
-              _Tool_to_mldev(self.api_client, item, to_object)
+              _Tool_to_mldev(self._api_client, item, to_object)
               for item in getv(from_object, ['tools'])
           ],
       )
@@ -531,7 +558,7 @@ class AsyncLive(_common.BaseModule):
           to_object,
           ['generationConfig'],
           _GenerateContentConfig_to_vertex(
-              self.api_client,
+              self._api_client,
               getv(from_object, ['generation_config']),
               to_object,
           ),
@@ -556,17 +583,18 @@ class AsyncLive(_common.BaseModule):
     if getv(from_object, ['speech_config']) is not None:
       if getv(to_object, ['generationConfig']) is not None:
         to_object['generationConfig']['speechConfig'] = _SpeechConfig_to_vertex(
-            self.api_client,
+            self._api_client,
             t.t_speech_config(
-                self.api_client, getv(from_object, ['speech_config'])),
+                self._api_client, getv(from_object, ['speech_config'])
+            ),
             to_object,
         )
       else:
         to_object['generationConfig'] = {
             'speechConfig': _SpeechConfig_to_vertex(
-                self.api_client,
+                self._api_client,
                 t.t_speech_config(
-                    self.api_client, getv(from_object, ['speech_config'])
+                    self._api_client, getv(from_object, ['speech_config'])
                 ),
                 to_object,
             )
@@ -576,9 +604,9 @@ class AsyncLive(_common.BaseModule):
           to_object,
           ['systemInstruction'],
           _Content_to_vertex(
-              self.api_client,
+              self._api_client,
               t.t_content(
-                  self.api_client, getv(from_object, ['system_instruction'])
+                  self._api_client, getv(from_object, ['system_instruction'])
               ),
               to_object,
           ),
@@ -588,7 +616,7 @@ class AsyncLive(_common.BaseModule):
           to_object,
           ['tools'],
           [
-              _Tool_to_vertex(self.api_client, item, to_object)
+              _Tool_to_vertex(self._api_client, item, to_object)
               for item in getv(from_object, ['tools'])
           ],
       )
@@ -599,7 +627,10 @@ class AsyncLive(_common.BaseModule):
 
   @contextlib.asynccontextmanager
   async def connect(
-      self, model: str, config: Optional[types.LiveConnectConfigOrDict] = None
+      self,
+      *,
+      model: str,
+      config: Optional[types.LiveConnectConfigOrDict] = None,
   ) -> AsyncSession:
     """Connect to the live server.
 
@@ -609,19 +640,19 @@ class AsyncLive(_common.BaseModule):
 
       client = genai.Client(api_key=API_KEY)
       config = {}
-      async with client.aio.live.connect(model='gemini-1.0-pro-002', config=config) as session:
+      async with client.aio.live.connect(model='...', config=config) as session:
         await session.send(input='Hello world!', end_of_turn=True)
-        async for message in session:
+        async for message in session.receive():
           print(message)
     """
-    base_url = self.api_client._websocket_base_url()
-    if self.api_client.api_key:
-      api_key = self.api_client.api_key
-      version = self.api_client._http_options['api_version']
+    base_url = self._api_client._websocket_base_url()
+    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']
+      headers = self._api_client._http_options['headers']
 
-      transformed_model = t.t_model(self.api_client, model)
+      transformed_model = t.t_model(self._api_client, model)
       request = json.dumps(
           self._LiveSetup_to_mldev(model=transformed_model, config=config)
       )
@@ -640,11 +671,11 @@ class AsyncLive(_common.BaseModule):
           'Content-Type': 'application/json',
           'Authorization': 'Bearer {}'.format(bearer_token),
       }
-      version = self.api_client._http_options['api_version']
+      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)
+      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
@@ -658,4 +689,4 @@ class AsyncLive(_common.BaseModule):
       await ws.send(request)
       logging.info(await ws.recv(decode=False))
 
-      yield AsyncSession(api_client=self.api_client, websocket=ws)
+      yield AsyncSession(api_client=self._api_client, websocket=ws)