google-genai 1.33.0__py3-none-any.whl → 1.53.0__py3-none-any.whl

google/genai/live.py

@@ -21,7 +21,7 @@ import contextlib
 import json
 import logging
 import typing
-from typing import Any, AsyncIterator, Dict, Optional, Sequence, Union, get_args
+from typing import Any, AsyncIterator, Optional, Sequence, Union, get_args
 import warnings
 
 import google.auth
@@ -40,7 +40,6 @@ from ._common import get_value_by_path as getv
 from ._common import set_value_by_path as setv
 from .live_music import AsyncLiveMusic
 from .models import _Content_to_mldev
-from .models import _Content_to_vertex
 
 
 try:
@@ -51,6 +50,11 @@ except ModuleNotFoundError:
   from websockets.client import ClientConnection  # type: ignore
   from websockets.client import connect as ws_connect  # type: ignore
 
+try:
+  from google.auth.transport import requests
+except ImportError:
+  requests = None  # type: ignore[assignment]
+
 if typing.TYPE_CHECKING:
   from mcp import ClientSession as McpClientSession
   from mcp.types import Tool as McpTool
@@ -82,9 +86,15 @@ _FUNCTION_RESPONSE_REQUIRES_ID = (
 class AsyncSession:
   """[Preview] AsyncSession."""
 
-  def __init__(self, api_client: BaseApiClient, websocket: ClientConnection):
+  def __init__(
+      self,
+      api_client: BaseApiClient,
+      websocket: ClientConnection,
+      session_id: Optional[str] = None,
+  ):
     self._api_client = api_client
     self._ws = websocket
+    self.session_id = session_id
 
   async def send(
       self,
@@ -217,8 +227,8 @@ class AsyncSession:
     )
 
     if self._api_client.vertexai:
-      client_content_dict = live_converters._LiveClientContent_to_vertex(
-          from_object=client_content
+      client_content_dict = _common.convert_to_dict(
+          client_content, convert_keys=True
       )
     else:
       client_content_dict = live_converters._LiveClientContent_to_mldev(
@@ -404,12 +414,12 @@ class AsyncSession:
     """
     tool_response = t.t_tool_response(function_responses)
     if self._api_client.vertexai:
-      tool_response_dict = live_converters._LiveClientToolResponse_to_vertex(
-          from_object=tool_response
+      tool_response_dict = _common.convert_to_dict(
+          tool_response, convert_keys=True
       )
     else:
-      tool_response_dict = live_converters._LiveClientToolResponse_to_mldev(
-          from_object=tool_response
+      tool_response_dict = _common.convert_to_dict(
+          tool_response, convert_keys=True
       )
       for response in tool_response_dict.get('functionResponses', []):
         if response.get('id') is None:
@@ -535,7 +545,7 @@ class AsyncSession:
     if self._api_client.vertexai:
       response_dict = live_converters._LiveServerMessage_from_vertex(response)
     else:
-      response_dict = live_converters._LiveServerMessage_from_mldev(response)
+      response_dict = response
 
     return types.LiveServerMessage._from_response(
         response=response_dict, kwargs=parameter_model.model_dump()
@@ -649,7 +659,7 @@ class AsyncSession:
           content_input_parts.append(item)
       if self._api_client.vertexai:
         contents = [
-            _Content_to_vertex(item, to_object)
+            _common.convert_to_dict(item, convert_keys=True)
             for item in t.t_contents(content_input_parts)
         ]
       else:
@@ -975,7 +985,8 @@ class AsyncLive(_api_module.BaseModule):
       api_key = self._api_client.api_key
       version = self._api_client._http_options.api_version
       uri = f'{base_url}/ws/google.cloud.aiplatform.{version}.LlmBidiService/BidiGenerateContent'
-      headers = self._api_client._http_options.headers or {}
+      original_headers = self._api_client._http_options.headers
+      headers = original_headers.copy() if original_headers is not None else {}
 
       request_dict = _common.convert_to_dict(
           live_converters._LiveConnectParameters_to_vertex(
@@ -992,27 +1003,51 @@ class AsyncLive(_api_module.BaseModule):
 
       request = json.dumps(request_dict)
     else:
-      if not self._api_client._credentials:
-        # Get bearer token through Application Default Credentials.
-        creds, _ = google.auth.default(  # type: ignore
-            scopes=['https://www.googleapis.com/auth/cloud-platform']
+      version = self._api_client._http_options.api_version
+      has_sufficient_auth = (
+          self._api_client.project and self._api_client.location
+      )
+      if self._api_client.custom_base_url and not has_sufficient_auth:
+        # API gateway proxy can use the auth in custom headers, not url.
+        # Enable custom url if auth is not sufficient.
+        uri = self._api_client.custom_base_url
+        # Keep the model as is.
+        transformed_model = model
+        # Do not get credentials for custom url.
+        original_headers = self._api_client._http_options.headers
+        headers = (
+            original_headers.copy() if original_headers is not None else {}
         )
+
       else:
-        creds = self._api_client._credentials
-      # creds.valid is False, and creds.token is None
-      # Need to refresh credentials to populate those
-      if not (creds.token and creds.valid):
-        auth_req = google.auth.transport.requests.Request()  # type: ignore
-        creds.refresh(auth_req)
-      bearer_token = creds.token
-      original_headers = self._api_client._http_options.headers
-      headers = original_headers.copy() if original_headers is not None else {}
-      headers['Authorization'] = f'Bearer {bearer_token}'
-      version = self._api_client._http_options.api_version
-      uri = f'{base_url}/ws/google.cloud.aiplatform.{version}.LlmBidiService/BidiGenerateContent'
+        uri = f'{base_url}/ws/google.cloud.aiplatform.{version}.LlmBidiService/BidiGenerateContent'
+
+        if not self._api_client._credentials:
+          # Get bearer token through Application Default Credentials.
+          creds, _ = google.auth.default(  # type: ignore
+              scopes=['https://www.googleapis.com/auth/cloud-platform']
+          )
+        else:
+          creds = self._api_client._credentials
+        # creds.valid is False, and creds.token is None
+        # Need to refresh credentials to populate those
+        if not (creds.token and creds.valid):
+          if requests is None:
+            raise ValueError('The requests module is required to refresh google-auth credentials. Please install with `pip install google-auth[requests]`')
+          auth_req = requests.Request()  # type: ignore
+          creds.refresh(auth_req)
+        bearer_token = creds.token
+
+        original_headers = self._api_client._http_options.headers
+        headers = (
+            original_headers.copy() if original_headers is not None else {}
+        )
+        if not headers.get('Authorization'):
+          headers['Authorization'] = f'Bearer {bearer_token}'
+
       location = self._api_client.location
       project = self._api_client.project
-      if transformed_model.startswith('publishers/'):
+      if transformed_model.startswith('publishers/') and project and location:
         transformed_model = (
             f'projects/{project}/locations/{location}/' + transformed_model
         )
@@ -1054,11 +1089,34 @@ class AsyncLive(_api_module.BaseModule):
       await ws.send(request)
       try:
         # websockets 14.0+
-        logger.info(await ws.recv(decode=False))
+        raw_response = await ws.recv(decode=False)
       except TypeError:
-        logger.info(await ws.recv())
+        raw_response = await ws.recv()  # type: ignore[assignment]
+      if raw_response:
+        try:
+          response = json.loads(raw_response)
+        except json.decoder.JSONDecodeError:
+          raise ValueError(f'Failed to parse response: {raw_response!r}')
+      else:
+        response = {}
 
-      yield AsyncSession(api_client=self._api_client, websocket=ws)
+      if self._api_client.vertexai:
+        response_dict = live_converters._LiveServerMessage_from_vertex(response)
+      else:
+        response_dict = response
+
+      setup_response = types.LiveServerMessage._from_response(
+          response=response_dict, kwargs=parameter_model.model_dump()
+      )
+      if setup_response.setup_complete:
+        session_id = setup_response.setup_complete.session_id
+      else:
+        session_id = None
+      yield AsyncSession(
+          api_client=self._api_client,
+          websocket=ws,
+          session_id=session_id,
+      )
 
 
 async def _t_live_connect_config(

google/genai/live_music.py

@@ -22,13 +22,11 @@ from typing import AsyncIterator
 
 from . import _api_module
 from . import _common
+from . import _live_converters as live_converters
 from . import _transformers as t
 from . import types
 from ._api_client import BaseApiClient
 from ._common import set_value_by_path as setv
-from . import _live_converters as live_converters
-from .models import _Content_to_mldev
-from .models import _Content_to_vertex
 
 
 try:
@@ -44,46 +42,47 @@ logger = logging.getLogger('google_genai.live_music')
 class AsyncMusicSession:
   """[Experimental] AsyncMusicSession."""
 
-  def __init__(
-      self, api_client: BaseApiClient, websocket: ClientConnection
-  ):
+  def __init__(self, api_client: BaseApiClient, websocket: ClientConnection):
     self._api_client = api_client
     self._ws = websocket
 
   async def set_weighted_prompts(
-      self,
-  	  prompts: list[types.WeightedPrompt]
+      self, prompts: list[types.WeightedPrompt]
   ) -> None:
     if self._api_client.vertexai:
-      raise NotImplementedError('Live music generation is not supported in Vertex AI.')
-    else:
-      client_content_dict = live_converters._LiveMusicClientContent_to_mldev(
-          from_object={'weighted_prompts': prompts}
+      raise NotImplementedError(
+          'Live music generation is not supported in Vertex AI.'
       )
+    else:
+      client_content_dict = {
+          'weightedPrompts': [
+              _common.convert_to_dict(prompt, convert_keys=True)
+              for prompt in prompts
+          ]
+      }
+
     await self._ws.send(json.dumps({'clientContent': client_content_dict}))
 
   async def set_music_generation_config(
-      self,
-  	  config: types.LiveMusicGenerationConfig
+      self, config: types.LiveMusicGenerationConfig
   ) -> None:
     if self._api_client.vertexai:
-      raise NotImplementedError('Live music generation is not supported in Vertex AI.')
-    else:
-      config_dict = live_converters._LiveMusicGenerationConfig_to_mldev(
-          from_object=config
+      raise NotImplementedError(
+          'Live music generation is not supported in Vertex AI.'
       )
+    else:
+      config_dict = _common.convert_to_dict(config, convert_keys=True)
     await self._ws.send(json.dumps({'musicGenerationConfig': config_dict}))
 
   async def _send_control_signal(
-      self,
-      playback_control: types.LiveMusicPlaybackControl
+      self, playback_control: types.LiveMusicPlaybackControl
   ) -> None:
     if self._api_client.vertexai:
-      raise NotImplementedError('Live music generation is not supported in Vertex AI.')
-    else:
-      playback_control_dict = live_converters._LiveMusicClientMessage_to_mldev(
-          from_object={'playback_control': playback_control}
+      raise NotImplementedError(
+          'Live music generation is not supported in Vertex AI.'
       )
+    else:
+      playback_control_dict = {'playbackControl': playback_control.value}
       await self._ws.send(json.dumps(playback_control_dict))
 
   async def play(self) -> None:
@@ -134,9 +133,7 @@ class AsyncMusicSession:
     if self._api_client.vertexai:
       raise NotImplementedError('Live music generation is not supported in Vertex AI.')
     else:
-      response_dict = live_converters._LiveMusicServerMessage_from_mldev(
-          response
-      )
+      response_dict = response
 
     return types.LiveMusicServerMessage._from_response(
         response=response_dict, kwargs=parameter_model.model_dump()

google/genai/local_tokenizer.py

@@ -25,11 +25,16 @@ from . import _common
 from . import _local_tokenizer_loader as loader
 from . import _transformers as t
 from . import types
-from . import types
-from ._transformers import t_contents
 
 logger = logging.getLogger("google_genai.local_tokenizer")
 
+__all__ = [
+    "_parse_hex_byte",
+    "_token_str_to_bytes",
+    "LocalTokenizer",
+    "_TextsAccumulator",
+]
+
 
 class _TextsAccumulator:
   """Accumulates countable texts from `Content` and `Tool` objects.
@@ -303,9 +308,20 @@ class LocalTokenizer:
 
     Args:
       contents: The contents to tokenize.
+      config: The configuration for counting tokens.
 
     Returns:
       A `CountTokensResult` containing the total number of tokens.
+
+    Usage:
+
+    .. code-block:: python
+
+      from google import genai
+      tokenizer = genai.LocalTokenizer(model_name='gemini-2.0-flash-001')
+      result = tokenizer.count_tokens("What is your name?")
+      print(result)
+      # total_tokens=5
     """
     processed_contents = t.t_contents(contents)
     text_accumulator = _TextsAccumulator()
@@ -330,7 +346,24 @@ class LocalTokenizer:
       self,
       contents: Union[types.ContentListUnion, types.ContentListUnionDict],
   ) -> types.ComputeTokensResult:
-    """Computes the tokens ids and string pieces in the input."""
+    """Computes the tokens ids and string pieces in the input.
+
+    Args:
+      contents: The contents to tokenize.
+
+    Returns:
+      A `ComputeTokensResult` containing the token information.
+
+    Usage:
+
+    .. code-block:: python
+
+      from google import genai
+      tokenizer = genai.LocalTokenizer(model_name='gemini-2.0-flash-001')
+      result = tokenizer.compute_tokens("What is your name?")
+      print(result)
+      # tokens_info=[TokensInfo(token_ids=[279, 329, 1313, 2508, 13], tokens=[b' What', b' is', b' your', b' name', b'?'], role='user')]
+    """
     processed_contents = t.t_contents(contents)
     text_accumulator = _TextsAccumulator()
     for content in processed_contents: