feat: Support ephemeral auth tokens as API keys for live connections in Python.

PiperOrigin-RevId: 765372194

Author: MarkDaoust

google/genai/__init__.py

@@ -17,9 +17,8 @@
 
 from . import version
 from .client import Client
-from .live import live_ephemeral_connect
 
 
 __version__ = version.__version__
 
-__all__ = ['Client', 'live_ephemeral_connect']
+__all__ = ['Client']

google/genai/_api_client.py

@@ -63,6 +63,11 @@
 INITIAL_RETRY_DELAY = 1  # second
 DELAY_MULTIPLIER = 2
 
+
+class EphemeralTokenAPIKeyError(ValueError):
+  """Error raised when the API key is invalid."""
+
+
 def _append_library_version_headers(headers: dict[str, str]) -> None:
   """Appends the telemetry header to the headers dict."""
   library_label = f'google-genai-sdk/{version.__version__}'
@@ -625,6 +630,11 @@ def _build_request(
         versioned_path,
     )
 
+    if self.api_key and self.api_key.startswith('auth_tokens/'):
+      raise EphemeralTokenAPIKeyError(
+          'Ephemeral tokens can only be used with the live API.'
+      )
+
     timeout_in_seconds = _get_timeout_in_seconds(patched_http_options.timeout)
 
     if patched_http_options.headers is None:

google/genai/live.py

@@ -34,6 +34,7 @@
 from . import _mcp_utils
 from . import _transformers as t
 from . import client
+from . import errors
 from . import types
 from ._api_client import BaseApiClient
 from ._common import get_value_by_path as getv
@@ -78,10 +79,6 @@
     ' response of a ToolCall.FunctionalCalls in Google AI.'
 )
 
-
-_DUMMY_KEY = 'dummy_key'
-
-
 class AsyncSession:
   """[Preview] AsyncSession."""
 
@@ -912,25 +909,10 @@ async def connect(
     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:
+    if config and config.http_options:
       raise ValueError(
           'google.genai.client.aio.live.connect() does not support'
           ' http_options at request-level in LiveConnectConfig yet. Please use'
@@ -945,10 +927,22 @@ async def _connect(
     parameter_model = await _t_live_connect_config(self._api_client, config)
 
     if self._api_client.api_key and not self._api_client.vertexai:
-      api_key = self._api_client.api_key
       version = self._api_client._http_options.api_version
-      if uri is None:
-        uri = f'{base_url}/ws/google.ai.generativelanguage.{version}.GenerativeService.BidiGenerateContent?key={api_key}'
+      api_key = self._api_client.api_key
+      method = 'BidiGenerateContent'
+      key_name = 'key'
+      if api_key.startswith('auth_tokens/'):
+        warnings.warn(
+            message=(
+                "The SDK's ephemeral token support is experimental, and may"
+                ' change in future versions.'
+            ),
+            category=errors.ExperimentalWarning,
+        )
+        method = 'BidiGenerateContentConstrained'
+        key_name = 'access_token'
+
+      uri = f'{base_url}/ws/google.ai.generativelanguage.{version}.GenerativeService.{method}?{key_name}={api_key}'
       headers = self._api_client._http_options.headers
 
       request_dict = _common.convert_to_dict(
@@ -969,8 +963,7 @@ async def _connect(
       # Headers already contains api key for express mode.
       api_key = self._api_client.api_key
       version = self._api_client._http_options.api_version
-      if uri is None:
-        uri = f'{base_url}/ws/google.cloud.aiplatform.{version}.LlmBidiService/BidiGenerateContent'
+      uri = f'{base_url}/ws/google.cloud.aiplatform.{version}.LlmBidiService/BidiGenerateContent'
       headers = self._api_client._http_options.headers
 
       request_dict = _common.convert_to_dict(
@@ -1060,83 +1053,6 @@ async def _connect(
         yield AsyncSession(api_client=self._api_client, websocket=ws)
 
 
-@_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],

google/genai/tests/live/test_live.py

@@ -1867,3 +1867,34 @@ async def _test_connect():
         )
 
     await _test_connect()
+
+
+@pytest.mark.parametrize('vertexai', [False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_auth_tokens(mock_websocket, vertexai):
+    api_client_mock = mock_api_client(vertexai=vertexai)
+    api_client_mock.api_key = 'auth_tokens/TEST_AUTH_TOKEN'
+    result = await get_connect_message(
+        api_client_mock,
+        model='test_model'
+    )
+
+    mock_ws = AsyncMock()
+    mock_ws.send = AsyncMock()
+    mock_ws.recv = AsyncMock(return_value=b'some response')
+    uri_capture = {}  # Capture the uri here
+
+    @contextlib.asynccontextmanager
+    async def mock_connect(uri, additional_headers=None):
+        uri_capture['uri'] = uri # Capture the uri
+        yield mock_ws
+
+    with patch.object(live, 'ws_connect', new=mock_connect):
+      live_module = live.AsyncLive(api_client_mock)
+      async with live_module.connect(
+          model='test_model',
+      ):
+        pass
+
+    assert 'access_token=auth_tokens/TEST_AUTH_TOKEN' in uri_capture['uri']
+    assert 'BidiGenerateContentConstrained' in uri_capture['uri']