google-genai 1.15.0__py3-none-any.whl → 1.16.1__py3-none-any.whl

google/genai/live.py

@@ -20,7 +20,8 @@ import base64
 import contextlib
 import json
 import logging
-from typing import Any, AsyncIterator, Dict, Optional, Sequence, Union, cast, get_args
+import typing
+from typing import Any, AsyncIterator, Dict, Optional, Sequence, Union, get_args
 import warnings
 
 import google.auth
@@ -29,24 +30,46 @@ from websockets import ConnectionClosed
 
 from . import _api_module
 from . import _common
+from . import _live_converters as live_converters
+from . import _mcp_utils
 from . import _transformers as t
 from . import client
 from . import types
 from ._api_client import BaseApiClient
 from ._common import get_value_by_path as getv
 from ._common import set_value_by_path as setv
-from . import _live_converters as live_converters
+from .live_music import AsyncLiveMusic
 from .models import _Content_to_mldev
 from .models import _Content_to_vertex
 
 
 try:
   from websockets.asyncio.client import ClientConnection
-  from websockets.asyncio.client import connect
+  from websockets.asyncio.client import connect as ws_connect
 except ModuleNotFoundError:
   # 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
+  from websockets.client import connect as ws_connect  # type: ignore
+
+if typing.TYPE_CHECKING:
+  from mcp import ClientSession as McpClientSession
+  from mcp.types import Tool as McpTool
+  from ._adapters import McpToGenAiToolAdapter
+  from ._mcp_utils import mcp_to_gemini_tool
+else:
+  McpClientSession: typing.Type = Any
+  McpTool: typing.Type = Any
+  McpToGenAiToolAdapter: typing.Type = Any
+  try:
+    from mcp import ClientSession as McpClientSession
+    from mcp.types import Tool as McpTool
+    from ._adapters import McpToGenAiToolAdapter
+    from ._mcp_utils import mcp_to_gemini_tool
+  except ImportError:
+    McpClientSession = None
+    McpTool = None
+    McpToGenAiToolAdapter = None
+    mcp_to_gemini_tool = None
 
 logger = logging.getLogger('google_genai.live')
 
@@ -56,12 +79,13 @@ _FUNCTION_RESPONSE_REQUIRES_ID = (
 )
 
 
+_DUMMY_KEY = 'dummy_key'
+
+
 class AsyncSession:
   """[Preview] AsyncSession."""
 
-  def __init__(
-      self, api_client: BaseApiClient, websocket: ClientConnection
-  ):
+  def __init__(self, api_client: BaseApiClient, websocket: ClientConnection):
     self._api_client = api_client
     self._ws = websocket
 
@@ -123,7 +147,7 @@ class AsyncSession:
           Union[
               types.Content,
               types.ContentDict,
-              list[Union[types.Content, types.ContentDict]]
+              list[Union[types.Content, types.ContentDict]],
           ]
       ] = None,
       turn_complete: bool = True,
@@ -264,7 +288,7 @@ class AsyncSession:
           print(f'{msg.text}')
     ```
     """
-    kwargs:dict[str, Any] = {}
+    kwargs: dict[str, Any] = {}
     if media is not None:
       kwargs['media'] = media
     if audio is not None:
@@ -506,9 +530,13 @@ class AsyncSession:
       response = {}
 
     if self._api_client.vertexai:
-      response_dict = live_converters._LiveServerMessage_from_vertex(self._api_client, response)
+      response_dict = live_converters._LiveServerMessage_from_vertex(
+          self._api_client, response
+      )
     else:
-      response_dict = live_converters._LiveServerMessage_from_mldev(self._api_client, response)
+      response_dict = live_converters._LiveServerMessage_from_mldev(
+          self._api_client, response
+      )
 
     return types.LiveServerMessage._from_response(
         response=response_dict, kwargs=parameter_model.model_dump()
@@ -522,7 +550,7 @@ class AsyncSession:
   ) -> None:
     async for data in data_stream:
       model_input = types.LiveClientRealtimeInput(
-        media_chunks=[types.Blob(data=data, mime_type=mime_type)]
+          media_chunks=[types.Blob(data=data, mime_type=mime_type)]
       )
       await self.send(input=model_input)
       # Give a chance for the receive loop to process responses.
@@ -560,9 +588,8 @@ class AsyncSession:
         raise ValueError(
             f'Unsupported input type "{type(input)}" or input content "{input}"'
         )
-      if (
-          isinstance(blob_input, types.Blob)
-          and isinstance(blob_input.data, bytes)
+      if isinstance(blob_input, types.Blob) and isinstance(
+          blob_input.data, bytes
       ):
         formatted_input = [
             blob_input.model_dump(mode='json', exclude_none=True)
@@ -841,6 +868,13 @@ class AsyncSession:
 class AsyncLive(_api_module.BaseModule):
   """[Preview] AsyncLive."""
 
+  def __init__(self, api_client: BaseApiClient):
+    super().__init__(api_client)
+    self._music = AsyncLiveMusic(api_client)
+
+  @property
+  def music(self) -> AsyncLiveMusic:
+    return self._music
 
   @contextlib.asynccontextmanager
   async def connect(
@@ -860,30 +894,70 @@ class AsyncLive(_api_module.BaseModule):
       client = genai.Client(api_key=API_KEY)
       config = {}
       async with client.aio.live.connect(model='...', config=config) as session:
-        await session.send(input='Hello world!', end_of_turn=True)
+        await session.send_client_content(
+          turns=types.Content(
+            role='user',
+            parts=[types.Part(text='hello!')]
+          ),
+          turn_complete=True
+        )
         async for message in session.receive():
           print(message)
+
+    Args:
+      model: The model to use for the live session.
+      config: The configuration for the live session.
+      **kwargs: additional keyword arguments.
+
+    Yields:
+      An AsyncSession object.
     """
+    async with self._connect(
+        model=model,
+        config=config,
+    ) as session:
+      yield session
+
+  @contextlib.asynccontextmanager
+  async def _connect(
+      self,
+      *,
+      model: Optional[str] = None,
+      config: Optional[types.LiveConnectConfigOrDict] = None,
+      uri: Optional[str] = None,
+  ) -> AsyncIterator[AsyncSession]:
+
+    # TODO(b/404946570): Support per request http options.
+    if isinstance(config, dict):
+      config = types.LiveConnectConfig(**config)
+    if config and config.http_options and uri is None:
+      raise ValueError(
+          'google.genai.client.aio.live.connect() does not support'
+          ' http_options at request-level in LiveConnectConfig yet. Please use'
+          ' the client-level http_options configuration instead.'
+      )
+
     base_url = self._api_client._websocket_base_url()
     if isinstance(base_url, bytes):
       base_url = base_url.decode('utf-8')
-    transformed_model = t.t_model(self._api_client, model)
+    transformed_model = t.t_model(self._api_client, model)  # type: ignore
 
-    parameter_model = _t_live_connect_config(self._api_client, config)
+    parameter_model = await _t_live_connect_config(self._api_client, 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}'
+      if uri is None:
+        uri = f'{base_url}/ws/google.ai.generativelanguage.{version}.GenerativeService.BidiGenerateContent?key={api_key}'
       headers = self._api_client._http_options.headers
 
       request_dict = _common.convert_to_dict(
           live_converters._LiveConnectParameters_to_mldev(
               api_client=self._api_client,
               from_object=types.LiveConnectParameters(
-                model=transformed_model,
-                config=parameter_model,
-              ).model_dump(exclude_none=True)
+                  model=transformed_model,
+                  config=parameter_model,
+              ).model_dump(exclude_none=True),
           )
       )
       del request_dict['config']
@@ -894,15 +968,15 @@ class AsyncLive(_api_module.BaseModule):
     else:
       if not self._api_client._credentials:
         # Get bearer token through Application Default Credentials.
-        creds, _ = google.auth.default(  # type: ignore[no-untyped-call]
-          scopes=['https://www.googleapis.com/auth/cloud-platform']
+        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):
-        auth_req = google.auth.transport.requests.Request()  # type: ignore[no-untyped-call]
+        auth_req = google.auth.transport.requests.Request()  # type: ignore
         creds.refresh(auth_req)
       bearer_token = creds.token
       headers = self._api_client._http_options.headers
@@ -922,34 +996,126 @@ class AsyncLive(_api_module.BaseModule):
           live_converters._LiveConnectParameters_to_vertex(
               api_client=self._api_client,
               from_object=types.LiveConnectParameters(
-                model=transformed_model,
-                config=parameter_model,
-              ).model_dump(exclude_none=True)
+                  model=transformed_model,
+                  config=parameter_model,
+              ).model_dump(exclude_none=True),
           )
       )
       del request_dict['config']
 
-      if getv(request_dict, ['setup', 'generationConfig', 'responseModalities']) is None:
-        setv(request_dict, ['setup', 'generationConfig', 'responseModalities'], ['AUDIO'])
+      if (
+          getv(
+              request_dict, ['setup', 'generationConfig', 'responseModalities']
+          )
+          is None
+      ):
+        setv(
+            request_dict,
+            ['setup', 'generationConfig', 'responseModalities'],
+            ['AUDIO'],
+        )
 
       request = json.dumps(request_dict)
 
+    if parameter_model.tools and _mcp_utils.has_mcp_tool_usage(
+        parameter_model.tools
+    ):
+      if headers is None:
+        headers = {}
+      _mcp_utils.set_mcp_usage_header(headers)
     try:
-      async with connect(uri, additional_headers=headers) as ws:
+      async with ws_connect(uri, additional_headers=headers) as ws:
         await ws.send(request)
         logger.info(await ws.recv(decode=False))
 
         yield AsyncSession(api_client=self._api_client, websocket=ws)
     except TypeError:
       # Try with the older websockets API
-      async with connect(uri, extra_headers=headers) as ws:
+      async with ws_connect(uri, extra_headers=headers) as ws:
         await ws.send(request)
         logger.info(await ws.recv())
 
         yield AsyncSession(api_client=self._api_client, websocket=ws)
 
 
-def _t_live_connect_config(
+@_common.experimental_warning(
+    "The SDK's Live API connection with ephemeral token implementation is"
+    ' experimental, and may change in future versions.',
+)
+@contextlib.asynccontextmanager
+async def live_ephemeral_connect(
+    access_token: str,
+    model: Optional[str] = None,
+    config: Optional[types.LiveConnectConfigOrDict] = None,
+) -> AsyncIterator[AsyncSession]:
+  """[Experimental] Connect to the live server using ephermeral token (Gemini Developer API only).
+
+  Note: the live API is currently in experimental.
+
+  Usage:
+
+  .. code-block:: python
+    from google import genai
+
+    config = {}
+    async with genai.live_ephemeral_connect(
+        access_token='auth_tokens/12345',
+        model='...',
+        config=config,
+        http_options=types.HttpOptions(api_version='v1beta'),
+    ) as session:
+      await session.send_client_content(
+          turns=types.Content(
+              role='user',
+              parts=[types.Part(text='hello!')]
+          ),
+          turn_complete=True
+      )
+
+      async for message in session.receive():
+        print(message)
+
+    Args:
+      access_token: The access token to use for the Live session. It can be
+        generated by the `client.tokens.create` method.
+      model: The model to use for the Live session.
+      config: The configuration for the Live session.
+
+    Yields:
+      An AsyncSession object.
+  """
+  if isinstance(config, dict):
+    config = types.LiveConnectConfig(**config)
+
+  http_options = config.http_options if config else None
+
+  base_url = (
+      http_options.base_url
+      if http_options and http_options.base_url
+      else 'https://generativelanguage.googleapis.com/'
+  )
+  api_version = (
+      http_options.api_version
+      if http_options and http_options.api_version
+      else 'v1beta'
+  )
+  internal_client = client.Client(
+      api_key=_DUMMY_KEY,  # Can't be None during initialization
+      http_options=types.HttpOptions(
+          base_url=base_url,
+          api_version=api_version,
+      ),
+  )
+  websocket_base_url = internal_client._api_client._websocket_base_url()
+  uri = f'{websocket_base_url}/ws/google.ai.generativelanguage.{api_version}.GenerativeService.BidiGenerateContentConstrained?access_token={access_token}'
+
+  async with internal_client.aio.live._connect(
+      model=model, config=config, uri=uri
+  ) as session:
+    yield session
+
+
+async def _t_live_connect_config(
     api_client: BaseApiClient,
     config: Optional[types.LiveConnectConfigOrDict],
 ) -> types.LiveConnectConfig:
@@ -975,7 +1141,24 @@ def _t_live_connect_config(
     parameter_model = config
     parameter_model.system_instruction = system_instruction
 
-  if parameter_model.generation_config is not None:
+  # Create a copy of the config model with the tools field cleared as they will
+  # be replaced with the MCP tools converted to GenAI tools.
+  parameter_model_copy = parameter_model.model_copy(update={'tools': None})
+  if parameter_model.tools:
+    parameter_model_copy.tools = []
+    for tool in parameter_model.tools:
+      if McpClientSession is not None and isinstance(tool, McpClientSession):
+        mcp_to_genai_tool_adapter = McpToGenAiToolAdapter(
+            tool, await tool.list_tools()
+        )
+        # Extend the config with the MCP session tools converted to GenAI tools.
+        parameter_model_copy.tools.extend(mcp_to_genai_tool_adapter.tools)
+      elif McpTool is not None and isinstance(tool, McpTool):
+        parameter_model_copy.tools.append(mcp_to_gemini_tool(tool))
+      else:
+        parameter_model_copy.tools.append(tool)
+
+  if parameter_model_copy.generation_config is not None:
     warnings.warn(
         'Setting `LiveConnectConfig.generation_config` is deprecated, '
         'please set the fields on `LiveConnectConfig` directly. This will '
@@ -984,4 +1167,4 @@ def _t_live_connect_config(
         stacklevel=4,
     )
 
-  return parameter_model
+  return parameter_model_copy

google/genai/live_music.py

@@ -0,0 +1,201 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""[Experimental] Live Music API client."""
+
+import contextlib
+import json
+import logging
+from typing import AsyncIterator
+
+from . import _api_module
+from . import _common
+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:
+  from websockets.asyncio.client import ClientConnection
+  from websockets.asyncio.client import connect
+except ModuleNotFoundError:
+  from websockets.client import ClientConnection  # type: ignore
+  from websockets.client import connect  # type: ignore
+
+logger = logging.getLogger('google_genai.live_music')
+
+
+class AsyncMusicSession:
+  """[Experimental] AsyncMusicSession."""
+
+  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]
+  ) -> 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(
+          api_client=self._api_client, from_object={'weighted_prompts': prompts}
+      )
+    await self._ws.send(json.dumps({'clientContent': client_content_dict}))
+
+  async def set_music_generation_config(
+      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(
+          api_client=self._api_client, from_object=config
+      )
+    await self._ws.send(json.dumps({'musicGenerationConfig': config_dict}))
+
+  async def _send_control_signal(
+      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(
+          api_client=self._api_client, from_object={'playback_control': playback_control}
+      )
+      await self._ws.send(json.dumps(playback_control_dict))
+
+  async def play(self) -> None:
+    """Sends playback signal to start the music stream."""
+    return await self._send_control_signal(types.LiveMusicPlaybackControl.PLAY)
+
+  async def pause(self) -> None:
+    """Sends a playback signal to pause the music stream."""
+    return await self._send_control_signal(types.LiveMusicPlaybackControl.PAUSE)
+
+  async def stop(self) -> None:
+    """Sends a playback signal to stop the music stream.
+
+    Resets the music generation context while retaining the current config.
+    """
+    return await self._send_control_signal(types.LiveMusicPlaybackControl.STOP)
+
+  async def reset_context(self) -> None:
+    """Reset the context (prompts retained) without stopping the music generation."""
+    return await self._send_control_signal(
+        types.LiveMusicPlaybackControl.RESET_CONTEXT
+    )
+
+  async def receive(self) -> AsyncIterator[types.LiveMusicServerMessage]:
+    """Receive model responses from the server.
+
+    Yields:
+      The audio chunks from the server.
+    """
+    # TODO(b/365983264) Handle intermittent issues for the user.
+    while result := await self._receive():
+      yield result
+
+  async def _receive(self) -> types.LiveMusicServerMessage:
+    parameter_model = types.LiveMusicServerMessage()
+    try:
+      raw_response = await self._ws.recv(decode=False)
+    except TypeError:
+      raw_response = await self._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 = {}
+
+    if self._api_client.vertexai:
+      raise NotImplementedError('Live music generation is not supported in Vertex AI.')
+    else:
+      response_dict = live_converters._LiveMusicServerMessage_from_mldev(
+          self._api_client, response
+      )
+
+    return types.LiveMusicServerMessage._from_response(
+        response=response_dict, kwargs=parameter_model.model_dump()
+    )
+
+  async def close(self) -> None:
+    """Closes the bi-directional stream and terminates the session."""
+    await self._ws.close()
+
+
+class AsyncLiveMusic(_api_module.BaseModule):
+  """[Experimental] Live music module.
+
+  Live music can be accessed via `client.aio.live.music`.
+  """
+
+  @_common.experimental_warning(
+      'Realtime music generation is experimental and may change in future versions.'
+  )
+  @contextlib.asynccontextmanager
+  async def connect(self, *, model: str) -> AsyncIterator[AsyncMusicSession]:
+    """[Experimental] Connect to the live music server."""
+    base_url = self._api_client._websocket_base_url()
+    if isinstance(base_url, bytes):
+      base_url = base_url.decode('utf-8')
+    transformed_model = t.t_model(self._api_client, model)
+
+    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.BidiGenerateMusic?key={api_key}'
+      headers = self._api_client._http_options.headers
+
+      # Only mldev supported
+      request_dict = _common.convert_to_dict(
+          live_converters._LiveMusicConnectParameters_to_mldev(
+              api_client=self._api_client,
+              from_object=types.LiveMusicConnectParameters(
+                model=transformed_model,
+              ).model_dump(exclude_none=True)
+          )
+      )
+
+      setv(request_dict, ['setup', 'model'], transformed_model)
+
+      request = json.dumps(request_dict)
+    else:
+      raise NotImplementedError('Live music generation is not supported in Vertex AI.')
+
+    try:
+      async with connect(uri, additional_headers=headers) as ws:
+        await ws.send(request)
+        logger.info(await ws.recv(decode=False))
+
+        yield AsyncMusicSession(api_client=self._api_client, websocket=ws)
+    except TypeError:
+      # Try with the older websockets API
+      async with connect(uri, extra_headers=headers) as ws:
+        await ws.send(request)
+        logger.info(await ws.recv())
+
+        yield AsyncMusicSession(api_client=self._api_client, websocket=ws)