first vommit
This commit is contained in:
commit
a04e68bc95
4926 files changed
+764440
No files matched your search
@@ -0,0 +1,16 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,389 @@
|
||||
# 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()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_uses_header_auth_without_query_key(mock_websocket):
|
||||
api_client = mock_api_client(vertexai=False)
|
||||
api_client._websocket_base_url = lambda: 'wss://generativelanguage.googleapis.com'
|
||||
api_client._http_options.api_version = 'v1beta'
|
||||
api_client._http_options.headers['x-goog-api-key'] = 'TEST_API_KEY'
|
||||
captured = {}
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def mock_connect(uri, additional_headers=None):
|
||||
captured['uri'] = uri
|
||||
captured['headers'] = additional_headers
|
||||
yield mock_websocket
|
||||
|
||||
@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='test_model'):
|
||||
pass
|
||||
|
||||
await _test_connect()
|
||||
|
||||
assert 'TEST_API_KEY' not in captured['uri']
|
||||
assert '?key=' not in captured['uri']
|
||||
assert captured['headers']['x-goog-api-key'] == 'TEST_API_KEY'
|
||||
|
||||
|
||||
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
|
||||
@@ -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
|
||||
+147
@@ -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
|
||||
+268
@@ -0,0 +1,268 @@
|
||||
# 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 json
|
||||
import os
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from websockets import client
|
||||
|
||||
from ... import client as gl_client
|
||||
from ... import live
|
||||
from ... import types
|
||||
from .. import pytest_helper
|
||||
|
||||
|
||||
IMAGE_FILE_PATH = os.path.abspath(
|
||||
os.path.join(os.path.dirname(__file__), '../data/google.jpg')
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
api_client._api_client = api_client
|
||||
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_media_blob_dict(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
content = {'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'audio/pcm'}
|
||||
|
||||
await session.send_realtime_input(media=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['mediaChunks'][0]['data'] == 'AAAAAAAA'
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
|
||||
) == 'audio/pcm'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_media_blob(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm')
|
||||
|
||||
await session.send_realtime_input(media=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['mediaChunks'][0]['data'] == 'AAAAAAAA'
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
|
||||
) == 'audio/pcm'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_media_image(mock_websocket, vertexai, image_jpeg):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
await session.send_realtime_input(media=image_jpeg)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['mediaChunks'][0], 'mime_type'
|
||||
) == 'image/jpeg'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
'content',
|
||||
[
|
||||
{'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'audio/pcm'},
|
||||
types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm'),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_audio(mock_websocket, vertexai, content):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
|
||||
await session.send_realtime_input(audio=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['audio']['data'] == 'AAAAAAAA'
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['audio'], 'mime_type'
|
||||
) == 'audio/pcm'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_bad_audio_blob(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='image/png')
|
||||
|
||||
with pytest.raises(ValueError, match='.*Unsupported mime type.*'):
|
||||
await session.send_realtime_input(audio=content)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_bad_video_blob(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
content = types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='audio/pcm')
|
||||
|
||||
with pytest.raises(ValueError, match='.*Unsupported mime type.*'):
|
||||
await session.send_realtime_input(video=content)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_audio_stream_end(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
await session.send_realtime_input(audio_stream_end=True)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['audioStreamEnd'] == True
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
'content',
|
||||
[
|
||||
{'data': bytes([0, 0, 0, 0, 0, 0]), 'mime_type': 'image/png'},
|
||||
types.Blob(data=bytes([0, 0, 0, 0, 0, 0]), mime_type='image/png'),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_video(mock_websocket, vertexai, content):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
|
||||
await session.send_realtime_input(video=content)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['video']['data'] == 'AAAAAAAA'
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['video'], 'mime_type'
|
||||
) == 'image/png'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_video_image(mock_websocket, vertexai, image_jpeg):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
|
||||
await session.send_realtime_input(video=image_jpeg)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['realtime_input']['video'], 'mime_type'
|
||||
) == 'image/jpeg'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_text(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
|
||||
await session.send_realtime_input(text='Hello?')
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['text'] == 'Hello?'
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.parametrize('activity', [{}, types.ActivityStart()])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_activity_start(mock_websocket, vertexai, activity):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
await session.send_realtime_input(activity_start=activity)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['activityStart'] == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.parametrize('activity', [{}, types.ActivityEnd()])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_activity_end(mock_websocket, vertexai, activity):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
await session.send_realtime_input(activity_end=activity)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'realtime_input' in sent_data
|
||||
|
||||
assert sent_data['realtime_input']['activityEnd'] == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_multiple_args(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
with pytest.raises(ValueError, match='.*one argument.*'):
|
||||
await session.send_realtime_input(
|
||||
text='Hello?', activity_start=types.ActivityStart()
|
||||
)
|
||||
+222
@@ -0,0 +1,222 @@
|
||||
# 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
|
||||
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 exception_if_mldev(vertexai, exception_type: type[Exception]):
|
||||
if vertexai:
|
||||
return contextlib.nullcontext()
|
||||
else:
|
||||
return pytest.raises(exception_type)
|
||||
|
||||
|
||||
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
|
||||
api_client._api_client = api_client
|
||||
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_function_response_dict(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(
|
||||
api_client=api_client, websocket=mock_websocket
|
||||
)
|
||||
|
||||
input = {
|
||||
'name': 'get_current_weather',
|
||||
'response': {'temperature': 14.5, 'unit': 'C'},
|
||||
}
|
||||
|
||||
if not vertexai:
|
||||
input['id'] = 'some-id'
|
||||
|
||||
await session.send_tool_response(function_responses=input)
|
||||
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'tool_response' in sent_data
|
||||
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['name']
|
||||
== 'get_current_weather'
|
||||
)
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['response'][
|
||||
'temperature'
|
||||
]
|
||||
== 14.5
|
||||
)
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['response']['unit']
|
||||
== 'C'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_function_response(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
input = types.FunctionResponse(
|
||||
name='get_current_weather',
|
||||
response={
|
||||
'temperature': 14.5,
|
||||
'unit': 'C',
|
||||
'user_name': 'test_user_name',
|
||||
'userEmail': 'test_user_email',
|
||||
},
|
||||
)
|
||||
if not vertexai:
|
||||
input.id = 'some-id'
|
||||
|
||||
await session.send_tool_response(function_responses=input)
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'tool_response' in sent_data
|
||||
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['name']
|
||||
== 'get_current_weather'
|
||||
)
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['response']
|
||||
== input.response
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_function_response_scheduling(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(api_client=api_client, websocket=mock_websocket)
|
||||
|
||||
input = types.FunctionResponse(
|
||||
name='get_current_weather',
|
||||
response={'temperature': 14.5, 'unit': 'C'},
|
||||
will_continue=True,
|
||||
scheduling=types.FunctionResponseScheduling.SILENT,
|
||||
)
|
||||
if not vertexai:
|
||||
input.id = 'some-id'
|
||||
|
||||
await session.send_tool_response(function_responses=input)
|
||||
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'tool_response' in sent_data
|
||||
|
||||
assert pytest_helper.get_value_ignore_key_case(
|
||||
sent_data['tool_response']['functionResponses'][0], 'will_continue'
|
||||
)
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['scheduling']
|
||||
== 'SILENT'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_function_response_list(mock_websocket, vertexai):
|
||||
session = live.AsyncSession(
|
||||
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
|
||||
)
|
||||
|
||||
input1 = {
|
||||
'name': 'get_current_weather',
|
||||
'response': {'temperature': 14.5, 'unit': 'C'},
|
||||
}
|
||||
input2 = {
|
||||
'name': 'get_current_weather',
|
||||
'response': {'temperature': 99.9, 'unit': 'C'},
|
||||
}
|
||||
|
||||
if not vertexai:
|
||||
input1['id'] = '1'
|
||||
input2['id'] = '2'
|
||||
|
||||
await session.send_tool_response(function_responses=[input1, input2])
|
||||
mock_websocket.send.assert_called_once()
|
||||
sent_data = json.loads(mock_websocket.send.call_args[0][0])
|
||||
assert 'tool_response' in sent_data
|
||||
|
||||
assert len(sent_data['tool_response']['functionResponses']) == 2
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][0]['response'][
|
||||
'temperature'
|
||||
]
|
||||
== 14.5
|
||||
)
|
||||
assert (
|
||||
sent_data['tool_response']['functionResponses'][1]['response'][
|
||||
'temperature'
|
||||
]
|
||||
== 99.9
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('vertexai', [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_id(mock_websocket, vertexai):
|
||||
api_client = mock_api_client(vertexai=vertexai)
|
||||
session = live.AsyncSession(
|
||||
api_client=api_client, websocket=mock_websocket
|
||||
)
|
||||
|
||||
input1 = {
|
||||
'name': 'get_current_weather',
|
||||
'response': {'temperature': 14.5, 'unit': 'C'},
|
||||
'id': '1',
|
||||
}
|
||||
input2 = {
|
||||
'name': 'get_current_weather',
|
||||
'response': {'temperature': 99.9, 'unit': 'C'},
|
||||
}
|
||||
|
||||
if not vertexai:
|
||||
with pytest.raises(ValueError, match=".*must have.*"):
|
||||
await session.send_tool_response(function_responses=[input1, input2])
|
||||
Reference in new issue
Block a user