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

google/genai/tests/live/test_live.py

@@ -0,0 +1,2177 @@
+# 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.
+#
+
+
+"""Tests for live.py."""
+
+import contextlib
+import json
+import os
+import ssl
+import typing
+from typing import Any, AsyncIterator
+from unittest import mock
+from unittest.mock import AsyncMock
+from unittest.mock import Mock
+from unittest.mock import patch
+import warnings
+
+import certifi
+from google.oauth2.credentials import Credentials
+import pytest
+from websockets import client
+
+from .. import pytest_helper
+from ... import _api_client as api_client
+from ... import _common
+from ... import Client
+from ... import client as gl_client
+from ... import live
+from ... import types
+try:
+    import aiohttp
+    AIOHTTP_NOT_INSTALLED = False
+except ImportError:
+    AIOHTTP_NOT_INSTALLED = True
+    aiohttp = mock.MagicMock()
+
+
+if typing.TYPE_CHECKING:
+  from mcp import types as mcp_types
+  from mcp import ClientSession as McpClientSession
+else:
+  mcp_types: typing.Type = Any
+  McpClientSession: typing.Type = Any
+  try:
+    from mcp import types as mcp_types
+    from mcp import ClientSession as McpClientSession
+  except ImportError:
+    mcp_types = None
+    McpClientSession = None
+
+
+requires_aiohttp = pytest.mark.skipif(
+    AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
+)
+
+function_declarations = [{
+    'name': 'get_current_weather',
+    'description': 'Get the current weather in a city',
+    'parameters': {
+        'type': 'OBJECT',
+        'properties': {
+            'location': {
+                'type': 'STRING',
+                'description': 'The location to get the weather for',
+            },
+            'unit': {
+                'type': 'STRING',
+                'enum': ['C', 'F'],
+            },
+        },
+    },
+}]
+
+
+def get_current_weather(location: str, unit: str):
+  """Get the current weather in a city."""
+  return 15 if unit == 'C' else 59
+
+
+def mock_api_client(vertexai=False, credentials=None, http_options=None):
+  api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
+  if not vertexai:
+    api_client.api_key = 'TEST_API_KEY'
+    api_client.location = None
+    api_client.project = None
+    api_client.custom_base_url = None
+  else:
+    api_client.api_key = None
+    if http_options:
+      http_options = (
+          types.HttpOptions(**http_options)
+          if isinstance(http_options, dict)
+          else http_options
+      )
+      api_client.custom_base_url = http_options.base_url
+      api_client.location = None
+      api_client.project = None
+    else:
+      api_client.location = 'us-central1'
+      api_client.project = 'test_project'
+      api_client.custom_base_url = None
+
+  api_client._host = lambda: 'test_host'
+  api_client._credentials = credentials
+  api_client._http_options = types.HttpOptions.model_validate(
+      {'headers': {}}
+  )  # Ensure headers exist
+  api_client.vertexai = vertexai
+  api_client._api_client = api_client
+  ctx = ssl.create_default_context(
+      cafile=os.environ.get("SSL_CERT_FILE", certifi.where()),
+      capath=os.environ.get("SSL_CERT_DIR"),
+  )
+  api_client._websocket_ssl_ctx = {'ssl': ctx}
+  return api_client
+
+
+@pytest.fixture
+def mock_websocket():
+  websocket = AsyncMock(spec=client.ClientConnection)
+  websocket.send = AsyncMock()
+  websocket.recv = AsyncMock(
+      return_value='{"serverContent": {"turnComplete": true}}'
+  )  # Default response
+  websocket.close = AsyncMock()
+  return websocket
+
+
+async def get_connect_message(api_client, model, config=None):
+  if config is None:
+    config = {}
+  mock_ws = AsyncMock()
+  mock_ws.send = AsyncMock()
+  mock_ws.recv = AsyncMock(
+      return_value=(
+          b'{\n  "setupComplete": {"sessionId": "test_session_id"}\n}\n'
+      )
+  )
+
+  mock_google_auth_default = Mock(return_value=(None, None))
+  mock_creds = Mock(token='test_token')
+  mock_google_auth_default.return_value = (mock_creds, None)
+
+  @contextlib.asynccontextmanager
+  async def mock_connect(uri, additional_headers=None, **kwargs):
+    yield mock_ws
+
+  @patch('google.auth.default', new=mock_google_auth_default)
+  @patch.object(live, 'ws_connect', new=mock_connect)
+  async def _test_connect():
+    live_module = live.AsyncLive(api_client)
+    async with live_module.connect(
+        model=model,
+        config=config,
+    ):
+      pass
+
+    mock_ws.send.assert_called_once()
+    return json.loads(mock_ws.send.call_args[0][0])
+
+  return await _test_connect()
+
+
+async def _async_iterator_to_list(async_iter):
+  return [value async for value in async_iter]
+
+
+def test_mldev_from_env(monkeypatch):
+  api_key = 'google_api_key'
+  monkeypatch.setenv('GOOGLE_API_KEY', api_key)
+
+  client = Client()
+
+  assert not client.aio.live._api_client.vertexai
+  assert client.aio.live._api_client.api_key == api_key
+  assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
+  assert client.aio.live._api_client._http_options.headers['x-goog-api-key'] == api_key
+
+
+@requires_aiohttp
+def test_vertex_from_env(monkeypatch):
+  project_id = 'fake_project_id'
+  location = 'fake-location'
+  monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
+  monkeypatch.setenv('GOOGLE_CLOUD_PROJECT', project_id)
+  monkeypatch.setenv('GOOGLE_CLOUD_LOCATION', location)
+
+  client = Client()
+
+  assert client.aio.live._api_client.custom_base_url is None
+  assert client.aio.live._api_client.vertexai
+  assert client.aio.live._api_client.project == project_id
+  assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
+  assert 'x-goog-api-key' not in client.aio.live._api_client._http_options.headers
+
+
+def test_vertex_api_key_from_env(monkeypatch):
+  api_key = 'google_api_key'
+  monkeypatch.setenv('GOOGLE_GENAI_USE_VERTEXAI', 'true')
+  monkeypatch.setenv('GOOGLE_API_KEY', api_key)
+
+  # Due to proj/location taking precedence, need to clear proj/location env
+  # variables. Tests in client/test_client_initialization.py provide
+  # comprehensive coverage for proj/location and api key precedence.
+  monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "")
+  monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "")
+
+  client = Client()
+
+  assert client.aio.live._api_client.vertexai
+  assert client.aio.live._api_client.api_key == api_key
+  assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
+  assert client.aio.live._api_client._http_options.headers['x-goog-api-key'] == api_key
+
+
+def test_websocket_base_url():
+  base_url = 'https://test.com'
+  api_client = gl_client.BaseApiClient(
+      api_key='google_api_key',
+      http_options={'base_url': base_url},
+  )
+  assert api_client._websocket_base_url() == 'wss://test.com'
+
+
+def test_websocket_base_url_no_auth_with_custom_base_url():
+  base_url = 'https://test-api-gateway-proxy.com'
+  api_client = gl_client.BaseApiClient(
+      vertexai=True,
+      http_options={
+          'base_url': base_url,
+          'headers': {'Authorization': 'Bearer test_token'},
+      },
+  )
+  # Note that our test environment does have project/location set. So we
+  # need to explicitly set them to None here.
+  api_client.project = None
+  api_client.location = None
+
+  # Fully pass the custom base url if no API key or project/location.
+  assert api_client._websocket_base_url() == base_url
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_send_text(mock_websocket, vertexai):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  await session.send(input='test')
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'client_content' in sent_data
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_send_content_dict(
+    mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  client_content = {
+      'content': [{'parts': [{'text': 'test'}]}],
+      'turn_complete': True,
+  }
+  await session.send(input=client_content)
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'client_content' in sent_data
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_send_content(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  client_content = types.LiveClientContent(
+      turns=[types.Content(parts=[types.Part(text='test')])], turn_complete=True
+  )
+  await session.send(input=client_content)
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'client_content' in sent_data
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_send_bytes(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  realtime_input = {'data': b'000000', 'mime_type': 'audio/pcm'}
+
+  await session.send(input=realtime_input)
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'realtime_input' in sent_data
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_send_blob(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  realtime_input = types.Blob(data=b'000000', mime_type='audio/pcm')
+
+  await session.send(input=realtime_input)
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'realtime_input' in sent_data
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_send_realtime_input(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  realtime_input = types.LiveClientRealtimeInput(
+      media_chunks=[types.Blob(data='MDAwMDAw', mime_type='audio/pcm')]
+  )
+  await session.send(input=realtime_input)
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'realtime_input' in sent_data
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_send_tool_response(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+
+  if vertexai:
+    tool_response = types.LiveClientToolResponse(
+        function_responses=[
+            types.FunctionResponse(
+                name='get_current_weather',
+                response={'temperature': 14.5, 'unit': 'C'},
+            )
+        ]
+    )
+  else:
+    tool_response = types.LiveClientToolResponse(
+        function_responses=[
+            types.FunctionResponse(
+                name='get_current_weather',
+                response={'temperature': 14.5, 'unit': 'C'},
+                id='some-id',
+            )
+        ]
+    )
+  await session.send(input=tool_response)
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'tool_response' in sent_data
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_send_input_none(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  await session.send(input=None)
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'client_content' in sent_data
+  assert sent_data['client_content']['turn_complete']
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_send_error(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  with pytest.raises(ValueError):
+    await session.send(input=[{'invalid_key': 'invalid_value'}])
+
+  with pytest.raises(ValueError):
+    await session.send(input={'invalid_key': 'invalid_value'})
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_receive( mock_websocket, vertexai):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  responses = session.receive()
+  responses = await _async_iterator_to_list(responses)
+  assert isinstance(responses[0], types.LiveServerMessage)
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_receive_error(
+     mock_websocket, vertexai
+):
+  mock_websocket.recv = AsyncMock(return_value='invalid json')
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  with pytest.raises(ValueError):
+    await session.receive().__anext__()
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_receive_text(
+     mock_websocket, vertexai
+):
+  mock_websocket.recv = AsyncMock(
+      side_effect=[
+          '{"serverContent": {"modelTurn": {"parts":[{"text": "test"}]}}}',
+          '{"serverContent": {"turnComplete": true}}',
+      ]
+  )
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  messages = session.receive()
+  messages = await _async_iterator_to_list(messages)
+  assert isinstance(messages[0], types.LiveServerMessage)
+  assert messages[0].server_content.model_turn.parts[0].text == 'test'
+  assert messages[1].server_content.turn_complete == True
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_receive_audio(
+     mock_websocket, vertexai
+):
+  mock_websocket.recv = AsyncMock(
+      side_effect=[
+          (
+              '{"serverContent": {"modelTurn": {"parts":[{"inlineData":'
+              ' {"data": "MDAwMDAw", "mimeType": "audio/pcm" }}]}}}'
+          ),
+          '{"serverContent": {"turnComplete": true}}',
+      ]
+  )
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  messages = session.receive()
+  messages = await _async_iterator_to_list(messages)
+  assert isinstance(messages[0], types.LiveServerMessage)
+  assert (
+      messages[0].server_content.model_turn.parts[0].inline_data.mime_type
+      == 'audio/pcm'
+  )
+  assert (
+      messages[0].server_content.model_turn.parts[0].inline_data.data
+      == b'000000'
+  )
+
+  with pytest.raises(RuntimeError):
+    await _async_iterator_to_list(session.receive())
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_receive_tool_call(
+     mock_websocket, vertexai
+):
+  mock_websocket.recv = AsyncMock(
+      side_effect=[
+          (
+              '{"toolCall": {"functionCalls": [{"name":'
+              ' "get_current_weather", "args": {"location": "San Francisco",'
+              ' "unit": "C"}}]}}'
+          ),
+          '{"serverContent": {"turnComplete": true}}',
+      ]
+  )
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  messages = session.receive()
+  messages = await _async_iterator_to_list(messages)
+  assert isinstance(messages[0], types.LiveServerMessage)
+  assert messages[0].tool_call.function_calls[0].name == 'get_current_weather'
+  assert (
+      messages[0].tool_call.function_calls[0].args['location']
+      == 'San Francisco'
+  )
+  assert messages[0].tool_call.function_calls[0].args['unit'] == 'C'
+
+  with pytest.raises(RuntimeError):
+    await _async_iterator_to_list(session.receive())
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_receive_transcription(
+     mock_websocket, vertexai
+):
+  mock_websocket.recv = AsyncMock(
+      side_effect=[
+          '{"serverContent": {"inputTranscription": {"text": "test_input", "finished": true}}}',
+          '{"serverContent": {"outputTranscription": {"text": "test_output", "finished": false}}}',
+          '{"serverContent": {"turnComplete": true}}',
+      ]
+  )
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  messages = session.receive()
+  messages = await _async_iterator_to_list(messages)
+  assert isinstance(messages[0], types.LiveServerMessage)
+  assert messages[0].server_content.input_transcription.text == 'test_input'
+  assert messages[0].server_content.input_transcription.finished == True
+
+  assert isinstance(messages[1], types.LiveServerMessage)
+  assert messages[1].server_content.output_transcription.text == 'test_output'
+  assert messages[1].server_content.output_transcription.finished == False
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_go_away(
+    mock_websocket, vertexai
+):
+  mock_websocket.recv = AsyncMock(
+      side_effect=[
+          '{"goAway": {"timeLeft": "10s"}}',
+          '{"serverContent": {"turnComplete": true}}',
+      ]
+  )
+  expected_result = types.LiveServerMessage(
+      go_away=types.LiveServerGoAway(time_left='10s'),
+  )
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  messages = session.receive()
+  messages = await _async_iterator_to_list(messages)
+  message = messages[0]
+
+  assert isinstance(message, types.LiveServerMessage)
+  assert message == expected_result
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_resumption_update(
+    mock_websocket, vertexai
+):
+  mock_websocket.recv = AsyncMock(
+      side_effect=[
+          """{
+                "sessionResumptionUpdate": {
+                    "newHandle": "test_handle",
+                    "resumable": "true",
+                    "lastConsumedClientMessageIndex": "123456789"
+                }
+          }""",
+          '{"serverContent": {"turnComplete": true}}',
+      ]
+  )
+
+  expected_result = types.LiveServerMessage(
+      session_resumption_update=types.LiveServerSessionResumptionUpdate(
+          new_handle='test_handle',
+          resumable=True,
+          last_consumed_client_message_index=123456789
+      ),
+  )
+
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  messages = session.receive()
+  messages = await _async_iterator_to_list(messages)
+  message = messages[0]
+
+  assert isinstance(message, types.LiveServerMessage)
+  assert message == expected_result
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_start_stream(
+     mock_websocket, vertexai
+):
+
+  session = live.AsyncSession(
+      mock_api_client(vertexai=vertexai), mock_websocket
+  )
+
+  async def mock_stream():
+    yield b'data1'
+    yield b'data2'
+
+  async for message in session.start_stream(
+      stream=mock_stream(), mime_type='audio/pcm'
+  ):
+    assert isinstance(message, types.LiveServerMessage)
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_receive_vad_signal(mock_websocket, vertexai):
+  # Simulate the server sending a VAD signal message
+  mock_websocket.recv = mock.AsyncMock(
+      side_effect=[
+          '{"voiceActivityDetectionSignal": {"vadSignalType": "VAD_SIGNAL_TYPE_SOS"}}',
+          '{"serverContent": {"turnComplete": true}}',  # To close the receiver loop
+      ]
+  )
+
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+
+  messages = await _async_iterator_to_list(session.receive())
+
+  # Check if the first message contains the VAD signal
+  assert len(messages) > 0
+  vad_message = messages[0]
+  assert isinstance(vad_message, types.LiveServerMessage)
+  assert vad_message.voice_activity_detection_signal is not None
+  assert (
+      vad_message.voice_activity_detection_signal.vad_signal_type
+      == types.VadSignalType.VAD_SIGNAL_TYPE_SOS
+  )
+
+  # Check that the session can close cleanly
+  assert messages[-1].server_content.turn_complete is True
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_close( mock_websocket, vertexai):
+  session = live.AsyncSession(
+      mock_api_client(vertexai=vertexai), mock_websocket
+  )
+  await session.close()
+  mock_websocket.close.assert_called_once()
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_no_config(vertexai):
+  with warnings.catch_warnings():
+    # Make sure there are no warnings cause by default values.
+    warnings.simplefilter('error')
+    result = await get_connect_message(
+        mock_api_client(vertexai=vertexai),
+        model='test_model'
+    )
+  expected_result = {'setup': {}}
+  if vertexai:
+    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
+    expected_result['setup']['generationConfig'] = {}
+    expected_result['setup']['generationConfig']['responseModalities'] = ["AUDIO"]
+  else:
+    expected_result['setup']['model'] = 'models/test_model'
+  assert result == expected_result
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_speech_config(vertexai):
+
+  expected_result = {
+      'setup': {
+          'model': 'models/test_model',
+          'generationConfig': {
+              'speechConfig': {
+                  'voice_config': {
+                      'prebuilt_voice_config': {'voice_name': 'en-default'}
+                  },
+                  'language_code': 'en-US',
+              },
+              'enableAffectiveDialog': True,
+              'temperature': 0.7,
+              'topP': 0.8,
+              'topK': 9.0,
+              'maxOutputTokens': 10,
+              'mediaResolution': 'MEDIA_RESOLUTION_MEDIUM',
+              'seed': 13,
+          },
+          'proactivity': {'proactive_audio': True},
+          'systemInstruction': {
+              'parts': [
+                  {
+                      'text': 'test instruction',
+                  },
+              ],
+              'role': 'user',
+          },
+      }
+  }
+  if vertexai:
+    expected_result['setup']['model'] = (
+        'projects/test_project/locations/us-central1/'
+        'publishers/google/models/test_model'
+    )
+    expected_result['setup']['generationConfig']['responseModalities'] = [
+        'AUDIO'
+    ]
+    expected_result['setup']['generationConfig']['speechConfig'] = {
+        'voiceConfig': {
+            'prebuilt_voice_config': {'voice_name': 'en-default'}
+        },
+        'languageCode': 'en-US',
+    }
+  else:
+    expected_result['setup']['model'] = 'models/test_model'
+
+  # Config is a dict
+  config_dict = {
+      'speech_config': {
+          'voice_config': {
+              'prebuilt_voice_config': {'voice_name': 'en-default'}
+          },
+          'language_code': 'en-US',
+      },
+      'enable_affective_dialog': True,
+      'proactivity': {'proactive_audio': True},
+      'temperature': 0.7,
+      'top_p': 0.8,
+      'top_k': 9,
+      'max_output_tokens': 10,
+      'seed': 13,
+      'system_instruction': 'test instruction',
+      'media_resolution': 'MEDIA_RESOLUTION_MEDIUM',
+  }
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai), model='test_model', config=config_dict
+  )
+  assert types.LiveClientMessage._from_response(
+      response=result, kwargs=None
+  ) == types.LiveClientMessage._from_response(
+      response=expected_result, kwargs=None
+  )
+  # Config is a LiveConnectConfig
+  config = types.LiveConnectConfig(
+      speech_config=types.SpeechConfig(
+          voice_config=types.VoiceConfig(
+              prebuilt_voice_config=types.PrebuiltVoiceConfig(
+                  voice_name='en-default'
+              )
+          ),
+          language_code='en-US',
+      ),
+      enable_affective_dialog=True,
+      proactivity=types.ProactivityConfig(proactive_audio=True),
+      temperature=0.7,
+      top_p=0.8,
+      top_k=9,
+      max_output_tokens=10,
+      media_resolution=types.MediaResolution.MEDIA_RESOLUTION_MEDIUM,
+      seed=13,
+      system_instruction='test instruction',
+  )
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+  assert types.LiveClientMessage._from_response(
+      response=result, kwargs=None
+  ) == types.LiveClientMessage._from_response(
+      response=expected_result, kwargs=None
+  )
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_error_if_multispeaker_voice_config(vertexai):
+
+  # Config is a dict
+  config_dict = {
+      'speech_config': {
+          'multi_speaker_voice_config': {
+              'speaker_voice_configs': [
+                  {
+                      'speaker': 'Alice',
+                      'voice_config': {
+                          'prebuilt_voice_config': {'voice_name': 'leda'}
+                      },
+                  },
+                  {
+                      'speaker': 'Bob',
+                      'voice_config': {
+                          'prebuilt_voice_config': {'voice_name': 'kore'}
+                      },
+                  },
+              ],
+          },
+      },
+      'temperature': 0.7,
+      'top_p': 0.8,
+      'top_k': 9,
+      'max_output_tokens': 10,
+      'seed': 13,
+      'system_instruction': 'test instruction',
+      'media_resolution': 'MEDIA_RESOLUTION_MEDIUM',
+  }
+  with pytest.raises(ValueError, match='.*multi_speaker_voice_config.*'):
+    result = await get_connect_message(
+        mock_api_client(vertexai=vertexai),
+        model='test_model',
+        config=config_dict,
+    )
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_replicated_voice_config(vertexai):
+  # Config is a dict
+  config_dict = {
+      'speech_config': {
+          'voice_config': {
+              'replicated_voice_config': {
+                  'mime_type': 'audio/pcm',
+                  'voice_sample_audio': bytes([0, 0, 0]),
+              },
+          },
+      },
+  }
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model',
+      config=config_dict,
+  )
+  if vertexai:
+    try:
+      replicated_voice_config = result['setup']['generationConfig'][
+          'speechConfig'
+      ]['voiceConfig']['replicatedVoiceConfig']
+    except KeyError:
+      replicated_voice_config = result['setup']['generationConfig'][
+          'speechConfig'
+      ]['voiceConfig']['replicated_voice_config']
+    assert replicated_voice_config == {
+        'mime_type': 'audio/pcm',
+        'voice_sample_audio': 'AAAA',
+    }
+  else:
+    return
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_explicit_vad(vertexai):
+  # Config is a dict
+  config_dict = {'explicit_vad_signal': True}
+  with pytest_helper.exception_if_mldev(
+      mock_api_client(vertexai=vertexai), ValueError
+  ):
+    result = await get_connect_message(
+        mock_api_client(vertexai=vertexai), model='test_model', config=config_dict
+    )
+  if not vertexai:
+    return
+  assert result['setup']['explicitVadSignal'] == True
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_explicit_vad_config(vertexai):
+  api_client = mock_api_client(vertexai=vertexai)
+
+  # Config is a dict
+  config_dict = {'explicit_vad_signal': True}
+  with pytest_helper.exception_if_mldev(api_client, ValueError):
+    result = await get_connect_message(
+        mock_api_client(vertexai=vertexai),
+        model='test_model',
+        config=config_dict,
+    )
+  if not vertexai:
+    return
+  assert result['setup']['explicitVadSignal'] == True
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_system_instruction_as_content_type(
+    vertexai,
+):
+  config_dict = {
+      'system_instruction': {
+          'parts': [{'text': 'test instruction'}],
+          'role': 'user',
+      },
+  }
+  config = types.LiveConnectConfig(**config_dict)
+  expected_result = {
+      'setup': {
+          'model': 'test_model',
+          'systemInstruction': {
+              'parts': [{'text': 'test instruction'}],
+              'role': 'user',
+          },
+      }
+  }
+  if vertexai:
+    expected_result['setup'][
+        'model'
+    ] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
+    expected_result['setup']['generationConfig'] = {}
+    expected_result['setup']['generationConfig']['responseModalities'] = [
+        'AUDIO'
+    ]
+  else:
+    expected_result['setup']['model'] = 'models/test_model'
+
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+  assert result == expected_result
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_config_tools_google_search(vertexai):
+  config_dict = {
+      'response_modalities': ['TEXT'],
+      'system_instruction': 'test instruction',
+      'generation_config': {'temperature': 0.7},
+      'tools': [{'google_search': {}}],
+  }
+
+  config = types.LiveConnectConfig(**config_dict)
+  expected_result = {
+      'setup': {
+          'generationConfig': {
+              'temperature': 0.7,
+              'responseModalities': ['TEXT'],
+          },
+          'systemInstruction': {
+              'parts': [{'text': 'test instruction'}],
+              'role': 'user',
+          },
+          'tools': [{'googleSearch': {}}],
+      }
+  }
+  if vertexai:
+    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
+  else:
+    expected_result['setup']['model'] = 'models/test_model'
+
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config_dict
+  )
+
+  assert result == expected_result
+
+  # Test config is a LiveConnectConfig
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+
+  assert result == expected_result
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_config_tools_with_no_mcp(vertexai):
+  config_dict = {
+      'response_modalities': ['TEXT'],
+      'system_instruction': 'test instruction',
+      'generation_config': {'temperature': 0.7},
+      'tools': [{'google_search': {}}],
+  }
+
+  config = types.LiveConnectConfig(**config_dict)
+  expected_result = {
+      'setup': {
+          'generationConfig': {
+              'temperature': 0.7,
+              'responseModalities': ['TEXT'],
+          },
+          'systemInstruction': {
+              'parts': [{'text': 'test instruction'}],
+              'role': 'user',
+          },
+          'tools': [{'googleSearch': {}}],
+      }
+  }
+  if vertexai:
+    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
+  else:
+    expected_result['setup']['model'] = 'models/test_model'
+
+  @patch.object(live, "McpClientSession", new=None)
+  @patch.object(live, "McpTool", new=None)
+  async def get_connect_message_no_mcp(config):
+    return await get_connect_message(
+        mock_api_client(vertexai=vertexai),
+        model='test_model', config=config
+    )
+
+  result = await get_connect_message_no_mcp(config_dict)
+  assert result == expected_result
+
+  result = await get_connect_message_no_mcp(config_dict)
+  assert result == expected_result
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_context_window_compression(
+     vertexai
+):
+  config = types.LiveConnectConfig(
+      generation_config=types.GenerationConfig(temperature=0.7),
+      response_modalities=['TEXT'],
+      system_instruction=types.Content(
+          parts=[types.Part(text='test instruction')], role='user'
+      ),
+      context_window_compression=types.ContextWindowCompressionConfig(
+          trigger_tokens=1000,
+          sliding_window=types.SlidingWindow(target_tokens=10),
+      ),
+  )
+  expected_result = {
+      'setup': {
+          'generationConfig': {
+              'temperature': 0.7,
+              'responseModalities': ['TEXT'],
+          },
+          'systemInstruction': {
+              'parts': [{'text': 'test instruction'}],
+              'role': 'user',
+          },
+           'contextWindowCompression': {
+              'trigger_tokens': 1000,
+              'sliding_window': {'target_tokens': 10},
+          }
+      }
+  }
+  if vertexai:
+    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
+  else:
+    expected_result['setup']['model'] = 'models/test_model'
+
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+  assert result == expected_result
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_config_tools_function_declaration(
+     vertexai
+):
+  config_dict = {
+      'generation_config': {'temperature': 0.7},
+      'tools': [{'function_declarations': function_declarations}],
+  }
+  config = types.LiveConnectConfig(**config_dict)
+  expected_result = {
+      'setup': {
+          'model': 'test_model',
+          'tools': [{
+              'functionDeclarations': [{
+                  'parameters': {
+                      'type': 'OBJECT',
+                      'properties': {
+                          'location': {
+                              'type': 'STRING',
+                              'description': (
+                                  'The location to get the weather for'
+                              ),
+                          },
+                          'unit': {'type': 'STRING', 'enum': ['C', 'F']},
+                      },
+                  },
+                  'name': 'get_current_weather',
+                  'description': 'Get the current weather in a city',
+              }],
+          }],
+      }
+  }
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+
+  assert result['setup']['tools'][0]['functionDeclarations'][0][
+      'description'
+  ] == (
+      expected_result['setup']['tools'][0]['functionDeclarations'][0][
+          'description'
+      ]
+  )
+
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+
+  assert result['setup']['tools'][0]['functionDeclarations'][0][
+      'description'
+  ] == (
+      expected_result['setup']['tools'][0]['functionDeclarations'][0][
+          'description'
+      ]
+  )
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_config_tools_function_directly(
+     vertexai
+):
+  config_dict = {
+      'generation_config': {'temperature': 0.7},
+      'tools': [get_current_weather],
+  }
+  config = types.LiveConnectConfig(**config_dict)
+  expected_result = {
+      'setup': {
+          'model': 'test_model',
+          'tools': [{
+              'functionDeclarations': [{
+                  'parameters': {
+                      'type': 'OBJECT',
+                      'properties': {
+                          'location': {
+                              'type': 'STRING',
+                              'description': (
+                                  'The location to get the weather for'
+                              ),
+                          },
+                          'unit': {'type': 'STRING', 'enum': ['C', 'F']},
+                      },
+                  },
+                  'name': 'get_current_weather',
+                  'description': 'Get the current weather in a city.',
+              }],
+          }],
+      }
+  }
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+
+  assert result['setup']['tools'][0]['functionDeclarations'][0][
+      'description'
+  ] == (
+      expected_result['setup']['tools'][0]['functionDeclarations'][0][
+          'description'
+      ]
+  )
+
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+
+  assert result['setup']['tools'][0]['functionDeclarations'][0][
+      'description'
+  ] == (
+      expected_result['setup']['tools'][0]['functionDeclarations'][0][
+          'description'
+      ]
+  )
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_tools_function_behavior(vertexai):
+  api_client = mock_api_client(vertexai=vertexai)
+
+  declaration = types.FunctionDeclaration.from_callable(
+      client=api_client, callable=get_current_weather
+  )
+  declaration.behavior = types.Behavior.NON_BLOCKING
+  config_dict = {
+      'generation_config': {'temperature': 0.7},
+      'tools': [{'function_declarations': [declaration]}],
+  }
+  config = types.LiveConnectConfig(**config_dict)
+
+  with pytest_helper.exception_if_vertex(api_client, ValueError):
+    result = await get_connect_message(
+        mock_api_client(vertexai=vertexai), model='test_model', config=config
+    )
+  if vertexai:
+    return
+
+  assert (
+      result['setup']['tools'][0]['functionDeclarations'][0]['behavior']
+      == 'NON_BLOCKING'
+  )
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_config_mcp_tools(
+    vertexai,
+):
+  if mcp_types is None:
+    return
+
+  expected_result_googleai = {
+      'setup': {
+          'model': 'models/test_model',
+          'tools': [{
+              'functionDeclarations': [{
+                  'parameters': {
+                      'type': 'OBJECT',
+                      'properties': {
+                          'location': {
+                              'type': 'STRING',
+                          },
+                      },
+                  },
+                  'name': 'get_weather',
+                  'description': 'Get the weather in a city.',
+              }],
+          }],
+      }
+  }
+  expected_result_vertexai = {
+      'setup': {
+          'generationConfig': {
+              'responseModalities': [
+                  'AUDIO',
+              ],
+          },
+          'model': (
+              'projects/test_project/locations/us-central1/publishers/google/models/test_model'
+          ),
+          'tools': [{
+              'functionDeclarations': [{
+                  'parameters': {
+                      'type': 'OBJECT',
+                      'properties': {
+                          'location': {
+                              'type': 'STRING',
+                          },
+                      },
+                  },
+                  'name': 'get_weather',
+                  'description': 'Get the weather in a city.',
+              }],
+          }],
+      }
+  }
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model',
+      config={
+          'tools': [
+              mcp_types.Tool(
+                  name='get_weather',
+                  description='Get the weather in a city.',
+                  inputSchema={
+                      'type': 'object',
+                      'properties': {'location': {'type': 'string'}},
+                  },
+              )
+          ],
+      },
+  )
+
+  assert (
+      result == expected_result_vertexai
+      if vertexai
+      else expected_result_googleai
+  )
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_config_mcp_session(
+    vertexai,
+):
+  if mcp_types is None:
+    return
+
+  class MockMcpClientSession(McpClientSession):
+
+    def __init__(self):
+      self._read_stream = None
+      self._write_stream = None
+
+    async def list_tools(self):
+      return mcp_types.ListToolsResult(
+          tools=[
+              mcp_types.Tool(
+                  name='get_weather',
+                  description='Get the weather in a city.',
+                  inputSchema={
+                      'type': 'object',
+                      'properties': {'location': {'type': 'string'}},
+                  },
+              ),
+          ]
+      )
+
+  expected_result_googleai = {
+      'setup': {
+          'model': 'models/test_model',
+          'tools': [{
+              'functionDeclarations': [{
+                  'parameters': {
+                      'type': 'OBJECT',
+                      'properties': {
+                          'location': {
+                              'type': 'STRING',
+                          },
+                      },
+                  },
+                  'name': 'get_weather',
+                  'description': 'Get the weather in a city.',
+              }],
+          }],
+      }
+  }
+  expected_result_vertexai = {
+      'setup': {
+          'generationConfig': {
+              'responseModalities': [
+                  'AUDIO',
+              ],
+          },
+          'model': (
+              'projects/test_project/locations/us-central1/publishers/google/models/test_model'
+          ),
+          'tools': [{
+              'functionDeclarations': [{
+                  'parameters': {
+                      'type': 'OBJECT',
+                      'properties': {
+                          'location': {
+                              'type': 'STRING',
+                          },
+                      },
+                  },
+                  'name': 'get_weather',
+                  'description': 'Get the weather in a city.',
+              }],
+          }],
+      }
+  }
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model',
+      config={
+          'tools': [MockMcpClientSession()],
+      },
+  )
+
+  assert (
+      result == expected_result_vertexai
+      if vertexai
+      else expected_result_googleai
+  )
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_config_tools_code_execution(
+     vertexai
+):
+  config_dict = {
+      'tools': [{'code_execution': {}}],
+  }
+  config = types.LiveConnectConfig(**config_dict)
+  expected_result = {
+      'setup': {
+          'model': 'test_model',
+          'tools': [{
+              'codeExecution': {},
+          }],
+      }
+  }
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+
+  assert result['setup']['tools'][0] == expected_result['setup']['tools'][0]
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_realtime_input_config(vertexai):
+  config_dict = {
+      'realtime_input_config': {
+          'automatic_activity_detection': {
+              'disabled': True,
+              'start_of_speech_sensitivity': 'START_SENSITIVITY_HIGH',
+              'end_of_speech_sensitivity': 'END_SENSITIVITY_HIGH',
+              'prefix_padding_ms': 20,
+              'silence_duration_ms': 100,
+          },
+          'activity_handling': 'NO_INTERRUPTION',
+          'turn_coverage': 'TURN_INCLUDES_ALL_INPUT',
+      }
+  }
+
+  config = types.LiveConnectConfig(**config_dict)
+  expected_result = {
+      'setup': {
+          'model': 'test_model',
+          'realtimeInputConfig': config_dict['realtime_input_config'],
+      }
+  }
+
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+
+  assert (
+      result['setup']['realtimeInputConfig']
+      == expected_result['setup']['realtimeInputConfig']
+  )
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_input_transcription(vertexai):
+  config_dict = {
+      'input_audio_transcription': {},
+  }
+  config = types.LiveConnectConfig(**config_dict)
+  expected_result = {
+      'setup': {
+          'model': 'test_model',
+          'inputAudioTranscription': {},
+      }
+  }
+
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai), model='test_model', config=config
+  )
+
+  assert (
+      result['setup']['inputAudioTranscription']
+      == expected_result['setup']['inputAudioTranscription']
+  )
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_output_transcription(vertexai):
+  config_dict = {
+      'output_audio_transcription': {},
+  }
+  config = types.LiveConnectConfig(**config_dict)
+  expected_result = {
+      'setup': {
+          'model': 'test_model',
+          'outputAudioTranscription': {},
+      }
+  }
+
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+
+  assert (
+      result['setup']['outputAudioTranscription']
+      == expected_result['setup']['outputAudioTranscription']
+  )
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_media_resolution(vertexai):
+  config_dict = {
+      'media_resolution': 'MEDIA_RESOLUTION_LOW',
+  }
+  config = types.LiveConnectConfig(**config_dict)
+  expected_result = {
+      'setup': {
+          'model': 'test_model',
+          'generationConfig': {'mediaResolution':'MEDIA_RESOLUTION_LOW'},
+      }
+  }
+
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model', config=config
+  )
+
+  assert (
+      result['setup']['generationConfig']['mediaResolution']
+      == expected_result['setup']['generationConfig']['mediaResolution']
+  )
+
+
+@pytest.mark.parametrize('vertexai', [True])
+@pytest.mark.asyncio
+async def test_bidi_setup_publishers(
+     vertexai
+):
+  expected_result = {
+      'setup': {
+         'generationConfig': {
+             'responseModalities': [
+                 'AUDIO',
+             ],
+         },
+         'model': 'projects/test_project/locations/us-central1/publishers/google/models/test_model',
+      }
+  }
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='publishers/google/models/test_model')
+
+  assert result == expected_result
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_generation_config_warning(
+     vertexai
+):
+  with pytest.warns(
+      DeprecationWarning,
+      match='Setting `LiveConnectConfig.generation_config` is deprecated'
+  ):
+    result = await get_connect_message(
+        mock_api_client(vertexai=vertexai),
+        model='models/test_model',
+        config={'generation_config': {'temperature': 0.7}})
+
+  assert result['setup']['generationConfig']['temperature'] == 0.7
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_session_resumption(vertexai):
+  config_dict = {
+      'session_resumption': {'handle': 'test_handle'},
+  }
+  config = types.LiveConnectConfig(**config_dict)
+
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model',
+      config=config
+  )
+  expected_result = {
+      'setup': {
+          'sessionResumption': {
+              'handle': 'test_handle',
+          },
+      }
+  }
+  if vertexai:
+    expected_result['setup']['generationConfig'] = {
+        'responseModalities': [
+            'AUDIO',
+        ],
+    }
+    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
+  else:
+    expected_result['setup']['model'] = 'models/test_model'
+  assert result == expected_result
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_transparent_session_resumption(vertexai):
+  api_client = mock_api_client(vertexai=vertexai)
+  config_dict = {
+      'session_resumption': {'handle': 'test_handle', 'transparent': True},
+  }
+  config = types.LiveConnectConfig(**config_dict)
+
+  with pytest_helper.exception_if_mldev(api_client, ValueError):
+    result = await get_connect_message(
+        mock_api_client(vertexai=vertexai),
+        model='test_model',
+        config=config
+    )
+
+  expected_result = {
+      'setup': {
+          'sessionResumption': {
+              'handle': 'test_handle',
+              'transparent': True,
+          },
+      }
+  }
+  if vertexai:
+    expected_result['setup']['generationConfig'] = {
+        'responseModalities': [
+            'AUDIO',
+        ],
+    }
+    expected_result['setup']['model'] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
+  else:
+    return
+
+  assert result == expected_result
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_str( mock_websocket, vertexai):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  result = session._parse_client_message('test')
+  assert 'client_content' in result
+  assert result == {
+      'client_content': {
+          'turn_complete': False,
+          'turns': [{'role': 'user', 'parts': [{'text': 'test'}]}],
+      }
+  }
+  # _parse_client_message returns a TypedDict, so we should be able to
+  # construct a LiveClientMessage from it
+  assert types.LiveClientMessage(**result)
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_thinking_config(vertexai):
+  config_dict = {
+      'thinking_config': {
+          'include_thoughts': True,
+          'thinking_budget': 1024,
+      }
+  }
+
+  expected_gen_config = {
+      'thinkingConfig': {
+          'include_thoughts': True,
+          'thinking_budget': 1024,
+      }
+  }
+
+  if vertexai:
+    expected_gen_config['responseModalities'] = ['AUDIO']
+
+  expected_result = {
+      'setup': {
+          'generationConfig': expected_gen_config,
+      }
+  }
+
+  if vertexai:
+    expected_result['setup'][
+        'model'
+    ] = 'projects/test_project/locations/us-central1/publishers/google/models/test_model'
+  else:
+    expected_result['setup']['model'] = 'models/test_model'
+
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai), model='test_model', config=config_dict
+  )
+  assert result == expected_result
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_blob( mock_websocket, vertexai):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  result = session._parse_client_message(
+      types.Blob(data=bytes([0, 0, 0]), mime_type='text/plain')
+  )
+  assert 'realtime_input' in result
+  assert result == {
+      'realtime_input': {
+          'media_chunks': [{'mime_type': 'text/plain', 'data': 'AAAA'}],
+      }
+  }
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_blob_dict(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+
+  blob = types.Blob(data=bytes([0, 0, 0]), mime_type='text/plain')
+  blob_dict = blob.model_dump()
+  result = session._parse_client_message(blob_dict)
+  assert 'realtime_input' in result
+  assert result == {
+      'realtime_input': {
+          'media_chunks': [{'mime_type': 'text/plain', 'data': 'AAAA'}],
+      }
+  }
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_client_content(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  result = session._parse_client_message(
+      types.LiveClientContent(
+          turn_complete=False,
+          turns=[types.Content(parts=[types.Part(text='test')], role='user')],
+      )
+  )
+  assert 'client_content' in result
+  assert result == {
+      'client_content': {
+          'turn_complete': False,
+          'turns': [{'role': 'user', 'parts': [{'text': 'test'}]}],
+      }
+  }
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_client_content_blob(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  client_content = types.LiveClientContent(
+      turn_complete=False,
+      turns=[
+          types.Content(
+              parts=[
+                  types.Part(
+                      inline_data=types.Blob(
+                          data=bytes([0, 0, 0]), mime_type='text/plain'
+                      )
+                  )
+              ],
+              role='user',
+          )
+      ],
+  )
+  result = session._parse_client_message(client_content)
+  assert 'client_content' in result
+  assert (
+      type(
+          result['client_content']['turns'][0]['parts'][0]['inline_data'][
+              'data'
+          ]
+      )
+      == str
+  )
+  assert result == {
+      'client_content': {
+          'turn_complete': False,
+          'turns': [{
+              'role': 'user',
+              'parts': [
+                  {'inline_data': {'mime_type': 'text/plain', 'data': 'AAAA'}}
+              ],
+          }],
+      }
+  }
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_client_content_dict(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  client_content = types.LiveClientContent(
+      turn_complete=False,
+      turns=[
+          types.Content(
+              parts=[
+                  types.Part(
+                      inline_data=types.Blob(
+                          data=bytes([0, 0, 0]), mime_type='text/plain'
+                      )
+                  )
+              ],
+              role='user',
+          )
+      ],
+  )
+  result = session._parse_client_message(
+      client_content.model_dump(mode='json', exclude_none=True)
+  )
+  assert 'client_content' in result
+  assert (
+      type(
+          result['client_content']['turns'][0]['parts'][0]['inline_data'][
+              'data'
+          ]
+      )
+      == str
+  )
+  assert result == {
+      'client_content': {
+          'turn_complete': False,
+          'turns': [{
+              'role': 'user',
+              'parts': [
+                  {'inline_data': {'mime_type': 'text/plain', 'data': 'AAAA'}}
+              ],
+          }],
+      }
+  }
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_realtime_input(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  input = types.LiveClientRealtimeInput(
+      media_chunks=[types.Blob(data=bytes([0, 0, 0]), mime_type='text/plain')]
+  )
+  result = session._parse_client_message(input)
+  assert 'realtime_input' in result
+  assert result == {
+      'realtime_input': {
+          'media_chunks': [{'mime_type': 'text/plain', 'data': 'AAAA'}],
+      }
+  }
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_realtime_input_dict(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  input = types.LiveClientRealtimeInput(
+      media_chunks=[types.Blob(data=bytes([0, 0, 0]), mime_type='text/plain')]
+  )
+  result = session._parse_client_message(
+      input.model_dump(mode='json', exclude_none=True)
+  )
+  assert 'realtime_input' in result
+  assert result == {
+      'realtime_input': {
+          'media_chunks': [{'mime_type': 'text/plain', 'data': 'AAAA'}],
+      }
+  }
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_tool_response(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  input = types.LiveClientToolResponse(
+      function_responses=[
+          types.FunctionResponse(
+              id='test_id',
+              name='test_name',
+              response={'result': 'test_response'},
+          )
+      ]
+  )
+  result = session._parse_client_message(input)
+  assert 'tool_response' in result
+  assert result == {
+      'tool_response': {
+          'function_responses': [
+              {
+                  'id': 'test_id',
+                  'name': 'test_name',
+                  'response': {
+                      'result': 'test_response',
+                  },
+              },
+          ],
+      }
+  }
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_function_response(
+    mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  input = types.FunctionResponse(
+    id='test_id',
+    name='test_name',
+    response={
+        'result': 'test_response',
+        'user_name': 'test_user_name',
+        'userEmail': 'test_user_email',
+    },
+  )
+  result = session._parse_client_message(input)
+  assert 'tool_response' in result
+  assert result == {
+      'tool_response': {
+          'function_responses': [
+              {
+                  'id': 'test_id',
+                  'name': 'test_name',
+                  'response': {
+                      'result': 'test_response',
+                      'user_name': 'test_user_name',
+                      'userEmail': 'test_user_email',
+                  },
+              },
+          ],
+      }
+  }
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_tool_response_dict_with_only_response(
+     mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  input = {
+    'id': 'test_id',
+    'name': 'test_name',
+    'response': {
+        'result': 'test_response',
+    }
+  }
+  result = session._parse_client_message(input)
+  assert 'tool_response' in result
+  assert result == {
+      'tool_response': {
+          'function_responses': [
+              {
+                  'id': 'test_id',
+                  'name': 'test_name',
+                  'response': {
+                      'result': 'test_response',
+                  },
+              },
+          ],
+      }
+  }
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+def test_parse_client_message_realtime_tool_response(
+    mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  input = types.LiveClientToolResponse(
+      function_responses=[
+          types.FunctionResponse(
+              id='test_id',
+              name='test_name',
+              response={'result': 'test_response'},
+          )
+      ]
+  )
+
+  result = session._parse_client_message(
+      input.model_dump(mode='json', exclude_none=True)
+  )
+  assert 'tool_response' in result
+  assert result == {
+      'tool_response': {
+          'function_responses': [
+              {
+                  'id': 'test_id',
+                  'name': 'test_name',
+                  'response': {
+                      'result': 'test_response',
+                  },
+              },
+          ],
+      }
+  }
+
+
+@pytest.mark.asyncio
+async def test_connect_with_provided_credentials(mock_websocket):
+  # custom oauth2 credentials
+  credentials = Credentials(token='provided_fake_token')
+  # mock api client
+  client = mock_api_client(vertexai=True, credentials=credentials)
+  capture = {}
+
+  @contextlib.asynccontextmanager
+  async def mock_connect(uri, additional_headers=None, **kwargs):
+    capture['headers'] = additional_headers
+    yield mock_websocket
+
+  @patch.object(live, 'ws_connect', new=mock_connect)
+  async def _test_connect():
+    live_module = live.AsyncLive(client)
+    async with live_module.connect(model='test-model'):
+      pass
+
+    assert 'Authorization' in capture['headers']
+    assert capture['headers']['Authorization'] == 'Bearer provided_fake_token'
+
+  await _test_connect()
+
+
+@pytest.mark.asyncio
+async def test_connect_with_default_credentials(mock_websocket):
+  # mock api client
+  client = mock_api_client(vertexai=True, credentials=None)
+  # mock google auth cred
+  mock_google_auth_default = Mock(return_value=(None, None))
+  mock_creds = Mock(token='default_test_token')
+  mock_google_auth_default.return_value = (mock_creds, None)
+  capture = {}
+
+  @contextlib.asynccontextmanager
+  async def mock_connect(uri, additional_headers=None, **kwargs):
+    capture['headers'] = additional_headers
+    yield mock_websocket
+
+  @patch('google.auth.default', new=mock_google_auth_default)
+  @patch.object(live, 'ws_connect', new=mock_connect)
+  async def _test_connect():
+    live_module = live.AsyncLive(client)
+    async with live_module.connect(model='test-model'):
+      pass
+
+    assert 'Authorization' in capture['headers']
+    assert capture['headers']['Authorization'] == 'Bearer default_test_token'
+
+  await _test_connect()
+
+
+@pytest.mark.asyncio
+async def test_connect_with_custom_base_url(mock_websocket):
+  # mock api client
+  client = gl_client.BaseApiClient(
+      vertexai=True,
+      http_options={
+          'base_url': 'https://custom-base-url.com',
+          'headers': {'Authorization': 'Bearer custom_test_token'},
+      }
+  )
+  # No ADC credentials.
+  capture = {}
+
+  @contextlib.asynccontextmanager
+  async def mock_connect(uri, additional_headers=None, **kwargs):
+    capture['uri'] = uri
+    capture['headers'] = additional_headers
+    yield mock_websocket
+
+  @patch.object(live, 'ws_connect', new=mock_connect)
+  async def _test_connect():
+    live_module = live.AsyncLive(client)
+    async with live_module.connect(model='test-model'):
+      pass
+
+    assert 'Authorization' in capture['headers']
+    assert capture['headers']['Authorization'] == 'Bearer custom_test_token'
+    assert capture['uri'] == 'https://custom-base-url.com'
+
+  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'{\n  "setupComplete": {"sessionId": "test_session_id"}\n}\n'
+    )
+  )
+  capture = {}
+
+  @contextlib.asynccontextmanager
+  async def mock_connect(uri, additional_headers=None, **kwargs):
+    capture['uri'] = uri
+    capture['headers'] = additional_headers
+    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 (
+      'Authorization' in capture['headers']
+  ), 'Authorization key is missing from headers'
+  assert (
+      capture['headers']['Authorization'] == 'Token auth_tokens/TEST_AUTH_TOKEN'
+  )
+  assert 'BidiGenerateContentConstrained' in capture['uri']
+
+
+@pytest.mark.parametrize('vertexai', [False])
+@pytest.mark.asyncio
+async def test_bidi_setup_to_api_with_api_key(mock_websocket, vertexai):
+  api_client_mock = mock_api_client(vertexai=vertexai)
+  api_client_mock._http_options = types.HttpOptions.model_validate(
+      {'headers': {'x-goog-api-key': 'TEST_API_KEY'}}
+  )
+  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'{\n  "setupComplete": {"sessionId": "test_session_id"}\n}\n'
+    )
+  )
+  capture = {}
+
+  @contextlib.asynccontextmanager
+  async def mock_connect(uri, additional_headers=None, **kwargs):
+    capture['uri'] = uri
+    capture['headers'] = additional_headers
+    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 'x-goog-api-key' in capture['headers'], "x-goog-api-key is missing from headers"
+  assert capture['headers']['x-goog-api-key'] == 'TEST_API_KEY'
+  assert 'BidiGenerateContent' in capture['uri']
+