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

google/genai/tests/live/test_live_music.py

@@ -0,0 +1,362 @@
+# 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_music.py."""
+import contextlib
+import json
+from typing import AsyncIterator
+from unittest import mock
+from unittest.mock import AsyncMock
+from unittest.mock import Mock
+from unittest.mock import patch
+import warnings
+
+from google.oauth2.credentials import Credentials
+import pytest
+from websockets import client
+
+from ... import _api_client as api_client
+from ... import _common
+from ... import Client
+from ... import client as gl_client
+from ... import live
+from ... import live_music
+from ... import types
+from .. import pytest_helper
+try:
+    import aiohttp
+    AIOHTTP_NOT_INSTALLED = False
+except ImportError:
+    AIOHTTP_NOT_INSTALLED = True
+    aiohttp = mock.MagicMock()
+
+
+requires_aiohttp = pytest.mark.skipif(
+    AIOHTTP_NOT_INSTALLED, reason="aiohttp is not installed, skipping test."
+)
+
+
+def mock_api_client(vertexai=False, credentials=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
+  else:
+    api_client.api_key = None
+    api_client.location = 'us-central1'
+    api_client.project = 'test_project'
+
+  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
+  return api_client
+
+
+@pytest.fixture
+def mock_websocket():
+  websocket = AsyncMock(spec=client.ClientConnection)
+  websocket.send = AsyncMock()
+  websocket.recv = AsyncMock(
+      return_value=b"""{
+  "serverContent": {
+    "audioChunks": [
+      {
+        "data": "Z2VsYmFuYW5h",
+        "mimeType": "audio/l16;rate=48000;channels=2",
+        "sourceMetadata": {
+          "clientContent": {
+            "weightedPrompts": [
+              {
+                "text": "Jazz",
+                "weight": 1
+              }
+            ]
+          },
+          "musicGenerationConfig": {
+            "seed": -957124937,
+            "bpm": 140,
+            "scale": "A_FLAT_MAJOR_F_MINOR"
+          }
+        }
+      }
+    ]
+  }
+}"""
+  )  # Default response
+  websocket.close = AsyncMock()
+  return websocket
+
+
+async def get_connect_message(api_client, model):
+  mock_ws = AsyncMock()
+  mock_ws.send = AsyncMock()
+  mock_ws.recv = AsyncMock(return_value=b'some response')
+
+  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):
+    yield mock_ws
+
+  @patch('google.auth.default', new=mock_google_auth_default)
+  @patch.object(live_music, 'connect', new=mock_connect)
+  async def _test_connect():
+    live_module = live.AsyncLive(api_client)
+    async with live_module.music.connect(
+        model=model,
+    ):
+      pass
+
+    mock_ws.send.assert_called_once()
+    return json.loads(mock_ws.send.call_args[0][0])
+
+  return await _test_connect()
+
+
+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.music._api_client.vertexai
+  assert client.aio.live.music._api_client.api_key == api_key
+  assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
+
+
+@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.music._api_client.vertexai
+  assert client.aio.live.music._api_client.project == project_id
+  assert isinstance(client.aio.live._api_client, api_client.BaseApiClient)
+
+
+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'
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_send_weighted_prompts(
+    mock_websocket, vertexai
+):
+  session = live_music.AsyncMusicSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  if vertexai:
+    with pytest.raises(NotImplementedError):
+      await session.set_weighted_prompts(prompts=[types.WeightedPrompt(text='Jazz', weight=1)])
+    return
+  await session.set_weighted_prompts(prompts=[types.WeightedPrompt(text='Jazz', weight=1)])
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'clientContent' in sent_data
+  assert sent_data['clientContent']['weightedPrompts'][0]['text'] == 'Jazz'
+  assert sent_data['clientContent']['weightedPrompts'][0]['weight'] == 1
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_send_config(
+    mock_websocket, vertexai
+):
+  session = live_music.AsyncMusicSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  if vertexai:
+    with pytest.raises(NotImplementedError):
+      await session.set_music_generation_config(
+          config=types.LiveMusicGenerationConfig(
+              bpm=140,
+              music_generation_mode=types.MusicGenerationMode.VOCALIZATION,
+          )
+      )
+    return
+  await session.set_music_generation_config(
+      config=types.LiveMusicGenerationConfig(
+          bpm=140,
+          music_generation_mode=types.MusicGenerationMode.VOCALIZATION,
+      )
+  )
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'musicGenerationConfig' in sent_data
+  assert sent_data['musicGenerationConfig']['bpm'] == 140
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_control_signal_play(
+    mock_websocket, vertexai
+):
+  session = live_music.AsyncMusicSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  if vertexai:
+    with pytest.raises(NotImplementedError):
+      await session.play()
+    return
+  await session.play()
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'playbackControl' in sent_data
+  assert 'PLAY' in sent_data['playbackControl']
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_control_signal_pause(
+    mock_websocket, vertexai
+):
+  session = live_music.AsyncMusicSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  if vertexai:
+    with pytest.raises(NotImplementedError):
+      await session.pause()
+    return
+  await session.pause()
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'playbackControl' in sent_data
+  assert 'PAUSE' in sent_data['playbackControl']
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_control_signal_stop(
+    mock_websocket, vertexai
+):
+  session = live_music.AsyncMusicSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  if vertexai:
+    with pytest.raises(NotImplementedError):
+      await session.stop()
+    return
+  await session.stop()
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'playbackControl' in sent_data
+  assert 'STOP' in sent_data['playbackControl']
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_control_signal_reset_context(
+    mock_websocket, vertexai
+):
+  session = live_music.AsyncMusicSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  if vertexai:
+    with pytest.raises(NotImplementedError):
+      await session.reset_context()
+    return
+  await session.reset_context()
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'playbackControl' in sent_data
+  assert 'RESET_CONTEXT' in sent_data['playbackControl']
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_async_session_receive( mock_websocket, vertexai):
+  session = live_music.AsyncMusicSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  if vertexai:
+    with pytest.raises(NotImplementedError):
+      async for _ in session.receive():
+        pass
+    return
+  async for response in session.receive():
+    assert isinstance(response, types.LiveMusicServerMessage)
+    audio_chunk = response.server_content.audio_chunks[0]
+    # Data contains decoded b64 audio
+    assert audio_chunk.data == b'gelbanana'
+    assert audio_chunk.mime_type == 'audio/l16;rate=48000;channels=2'
+    assert audio_chunk.source_metadata.client_content.weighted_prompts[0].text == 'Jazz'
+    assert audio_chunk.source_metadata.client_content.weighted_prompts[0].weight == 1
+    assert audio_chunk.source_metadata.music_generation_config.bpm == 140
+    break
+
+
+@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_music.AsyncMusicSession(
+      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_close( mock_websocket, vertexai):
+  session = live_music.AsyncMusicSession(
+      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_setup_to_api(vertexai):
+  if vertexai:
+    with pytest.raises(NotImplementedError):
+      await get_connect_message(
+          mock_api_client(vertexai=vertexai),
+          model='test_model'
+      )
+    return
+  result = await get_connect_message(
+      mock_api_client(vertexai=vertexai),
+      model='test_model'
+  )
+  expected_result = {'setup': {}}
+  if vertexai:
+    # Vertex is not supported yet
+    assert False
+  else:
+    expected_result['setup']['model'] = 'models/test_model'
+  assert result == expected_result

google/genai/tests/live/test_live_response.py

@@ -0,0 +1,163 @@
+# 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 response handling."""
+import json
+from typing import cast
+from unittest import mock
+from unittest.mock import AsyncMock
+
+import pytest
+
+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
+
+
+def mock_api_client(vertexai=False):
+  """Creates a mock BaseApiClient."""
+  mock_client = mock.MagicMock(spec=gl_client.BaseApiClient)
+  if not vertexai:
+    mock_client.api_key = 'TEST_API_KEY'
+    mock_client.location = None
+    mock_client.project = None
+  else:
+    mock_client.api_key = None
+    mock_client.location = 'us-central1'
+    mock_client.project = 'test_project'
+
+  mock_client._host = lambda: 'test_host'
+  mock_client._http_options = types.HttpOptions.model_validate(
+      {'headers': {}}
+  )
+  mock_client.vertexai = vertexai
+  return mock_client
+
+
+@pytest.fixture
+def mock_websocket():
+  """Provides a mock websocket connection."""
+  # Use live.ClientConnection if that's the specific type hint in AsyncSession
+  websocket = AsyncMock(spec=live.ClientConnection)
+  websocket.send = AsyncMock()
+  # Set default recv value, will be overridden in the test
+  websocket.recv = AsyncMock(return_value='{}')
+  websocket.close = AsyncMock()
+  return websocket
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_receive_server_content(mock_websocket, vertexai):
+
+  raw_response_json = json.dumps({
+      "usageMetadata": {
+          "promptTokenCount": 15,
+          "responseTokenCount": 25,
+          "candidatesTokenCount": 50,
+          "totalTokenCount": 200,
+          "responseTokensDetails": [
+              {
+                  "tokenCount": 20,
+                  "modality": "TEXT",
+              }
+          ],
+          "candidatesTokensDetails": [
+              {
+                  "tokenCount": 10,
+                  "modality": "TEXT",
+              }
+          ],
+      },
+      "serverContent": {
+          "modelTurn": {
+              "parts": [{"text": "This is a simple response."}]
+          },
+          "turnComplete": True,
+          "groundingMetadata": {
+              "web_search_queries": ["test query"],
+              "groundingChunks": [{
+                  "web": {
+                      "domain": "google.com",
+                      "title": "Search results",
+                  }
+              }]
+          }
+      }
+  })
+  mock_websocket.recv.return_value = raw_response_json
+
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  result = await session._receive()
+
+  # Assert the results
+  assert isinstance(result, types.LiveServerMessage)
+
+  assert (
+      result.server_content.model_turn.parts[0].text
+      == "This is a simple response."
+  )
+  assert result.server_content.turn_complete
+  assert result.server_content.grounding_metadata.web_search_queries == ["test query"]
+  assert result.server_content.grounding_metadata.grounding_chunks[0].web.domain == "google.com"
+  assert result.server_content.grounding_metadata.grounding_chunks[0].web.title == "Search results"
+  # Verify usageMetadata was parsed
+  assert isinstance(result.usage_metadata, types.UsageMetadata)
+  assert result.usage_metadata.prompt_token_count == 15
+  assert result.usage_metadata.total_token_count == 200
+  if not vertexai:
+    assert result.usage_metadata.response_token_count == 25
+    assert result.usage_metadata.response_tokens_details[0].token_count == 20
+  else:
+    # VertexAI maps candidatesTokenCount to responseTokenCount and maps
+    # candidatesTokensDetails to responseTokensDetails.
+    assert result.usage_metadata.response_token_count == 50
+    assert result.usage_metadata.response_tokens_details[0].token_count == 10
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_receive_server_content_with_turn_reason(mock_websocket, vertexai):
+  """Tests parsing of LiveServerContent with turn_complete_reason and waiting_for_input."""
+
+  raw_response_json = json.dumps({
+      "serverContent": {
+          "modelTurn": {
+              "parts": [{"text": "Please provide more details."}]
+          },
+          "turnComplete": True,
+          "turnCompleteReason": "NEED_MORE_INPUT",
+          "waitingForInput": True
+      }
+  })
+  mock_websocket.recv.return_value = raw_response_json
+
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  result = await session._receive()
+
+  # Assert the results
+  assert isinstance(result, types.LiveServerMessage)
+  assert result.server_content is not None
+
+  assert result.server_content.model_turn.parts[0].text == "Please provide more details."
+  assert result.server_content.turn_complete is True
+  assert result.server_content.turn_complete_reason == types.TurnCompleteReason.NEED_MORE_INPUT
+  assert result.server_content.waiting_for_input is True

google/genai/tests/live/test_send_client_content.py

@@ -0,0 +1,147 @@
+# 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 base64
+import json
+from unittest import mock
+
+import pytest
+from websockets import client
+
+from .. import pytest_helper
+from ... import client as gl_client
+from ... import live
+from ... import types
+
+
+def mock_api_client(vertexai=False):
+  api_client = mock.MagicMock(spec=gl_client.BaseApiClient)
+  api_client.api_key = 'TEST_API_KEY'
+  api_client._host = lambda: 'test_host'
+  api_client._http_options = {'headers': {}}  # Ensure headers exist
+  api_client.vertexai = vertexai
+  return api_client
+
+
+@pytest.fixture
+def mock_websocket():
+  websocket = mock.AsyncMock(spec=client.ClientConnection)
+  websocket.send = mock.AsyncMock()
+  websocket.recv = mock.AsyncMock(
+      return_value='{"serverContent": {"turnComplete": true}}'
+  )  # Default response
+  websocket.close = mock.AsyncMock()
+  return websocket
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_send_content_dict(mock_websocket, vertexai):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  content = [{'parts': [{'text': 'test'}]}]
+
+  await session.send_client_content(turns=content)
+  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']['turns'][0]['parts'][0]['text'] == 'test'
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_send_content_dict_list(mock_websocket, vertexai):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  content = [{'parts': [{'text': 'test'}]}]
+
+  await session.send_client_content(turns=content)
+  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']['turns'][0]['parts'][0]['text'] == 'test'
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_send_content_content(mock_websocket, vertexai):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  content = types.Content.model_validate({'parts': [{'text': 'test'}]})
+
+  await session.send_client_content(turns=content)
+  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']['turns'][0]['parts'][0]['text'] == 'test'
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_send_content_with_blob(mock_websocket, vertexai):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+  content = types.Content.model_validate(
+      {'parts': [{'inline_data': {'data': b'test'}}]}
+  )
+
+  await session.send_client_content(turns=content)
+  mock_websocket.send.assert_called_once()
+  sent_data = json.loads(mock_websocket.send.call_args[0][0])
+  assert 'client_content' in sent_data
+
+  assert pytest_helper.get_value_ignore_key_case(
+      sent_data['client_content']['turns'][0]['parts'][0], 'inline_data') == {
+          'data': base64.b64encode(b'test').decode()
+      }
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_send_client_content_turn_complete_false(
+    mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+
+  await session.send_client_content(turn_complete=False)
+  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']['turnComplete'] == False
+
+
+@pytest.mark.parametrize('vertexai', [True, False])
+@pytest.mark.asyncio
+async def test_send_client_content_empty(
+    mock_websocket, vertexai
+):
+  session = live.AsyncSession(
+      api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
+  )
+
+  await session.send_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
+  assert sent_data['client_content']['turnComplete'] == True