Skip to content

Commit 2fd4079

Browse files
speedstorm1copybara-github
authored andcommitted
fix: Base64-encode inline bytes in Live send_tool_response
fixes #3022 PiperOrigin-RevId: 998886870
1 parent 4c1abcc commit 2fd4079

2 files changed

Lines changed: 111 additions & 0 deletions

File tree

‎google/genai/live.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -445,6 +445,8 @@ async def send_tool_response(
445445
if response.get('id') is None:
446446
raise ValueError(_FUNCTION_RESPONSE_REQUIRES_ID)
447447

448+
tool_response_dict = _common.encode_unserializable_types(tool_response_dict)
449+
448450
await self._ws.send(json.dumps({'tool_response': tool_response_dict}))
449451

450452
async def send_context_update(

‎google/genai/tests/live/test_send_tool_response.py‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
"""Tests for live.py."""
1818

19+
import base64
1920
import contextlib
2021
import json
2122
from unittest import mock
@@ -220,3 +221,111 @@ async def test_missing_id(mock_websocket, vertexai):
220221
if not vertexai:
221222
with pytest.raises(ValueError, match=".*must have.*"):
222223
await session.send_tool_response(function_responses=[input1, input2])
224+
225+
226+
_IMAGE_BYTES = b'\x89PNG\r\n\x1a\n\x00\xff\xfe'
227+
228+
229+
def _assert_inline_image(part):
230+
inline_data = pytest_helper.get_value_ignore_key_case(part, 'inline_data')
231+
mime_type = pytest_helper.get_value_ignore_key_case(inline_data, 'mime_type')
232+
assert mime_type == 'image/png'
233+
assert isinstance(inline_data['data'], str)
234+
assert base64.urlsafe_b64decode(inline_data['data']) == _IMAGE_BYTES
235+
236+
237+
@pytest.mark.parametrize('vertexai', [True, False])
238+
@pytest.mark.asyncio
239+
async def test_function_response_inline_bytes(mock_websocket, vertexai):
240+
"""Inline bytes in function response parts are sent as base64."""
241+
session = live.AsyncSession(
242+
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
243+
)
244+
245+
function_response_input = types.FunctionResponse(
246+
id='some-id',
247+
name='take_photo',
248+
response={'output': 'Photo taken.'},
249+
parts=[
250+
types.FunctionResponsePart.from_bytes(
251+
data=_IMAGE_BYTES, mime_type='image/png'
252+
)
253+
],
254+
)
255+
256+
await session.send_tool_response(function_responses=function_response_input)
257+
258+
mock_websocket.send.assert_called_once()
259+
sent_data = json.loads(mock_websocket.send.call_args[0][0])
260+
function_response = sent_data['tool_response']['functionResponses'][0]
261+
assert function_response['response'] == {'output': 'Photo taken.'}
262+
_assert_inline_image(function_response['parts'][0])
263+
264+
265+
@pytest.mark.parametrize('vertexai', [True, False])
266+
@pytest.mark.asyncio
267+
async def test_function_response_inline_bytes_dict(mock_websocket, vertexai):
268+
"""Inline bytes given in a plain dict are sent as base64."""
269+
session = live.AsyncSession(
270+
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
271+
)
272+
273+
function_response_input = {
274+
'id': 'some-id',
275+
'name': 'take_photo',
276+
'response': {'output': 'Photo taken.'},
277+
'parts': [
278+
{'inline_data': {'data': _IMAGE_BYTES, 'mime_type': 'image/png'}}
279+
],
280+
}
281+
282+
await session.send_tool_response(function_responses=function_response_input)
283+
284+
mock_websocket.send.assert_called_once()
285+
sent_data = json.loads(mock_websocket.send.call_args[0][0])
286+
function_response = sent_data['tool_response']['functionResponses'][0]
287+
assert function_response['response'] == {'output': 'Photo taken.'}
288+
_assert_inline_image(function_response['parts'][0])
289+
290+
291+
@pytest.mark.parametrize('vertexai', [True, False])
292+
@pytest.mark.asyncio
293+
async def test_function_response_list_mixed_parts(mock_websocket, vertexai):
294+
"""Bytes, file URI and part-less responses survive together in one list."""
295+
session = live.AsyncSession(
296+
api_client=mock_api_client(vertexai=vertexai), websocket=mock_websocket
297+
)
298+
299+
input1 = types.FunctionResponse(
300+
id='1',
301+
name='take_photo',
302+
response={'output': 'Photo taken.'},
303+
parts=[
304+
types.FunctionResponsePart.from_bytes(
305+
data=_IMAGE_BYTES, mime_type='image/png'
306+
),
307+
types.FunctionResponsePart.from_uri(
308+
file_uri='gs://bucket/photo.png', mime_type='image/png'
309+
),
310+
],
311+
)
312+
input2 = types.FunctionResponse(
313+
id='2', name='get_time', response={'time': '12:00'}
314+
)
315+
316+
await session.send_tool_response(function_responses=[input1, input2])
317+
318+
mock_websocket.send.assert_called_once()
319+
sent_data = json.loads(mock_websocket.send.call_args[0][0])
320+
function_responses = sent_data['tool_response']['functionResponses']
321+
assert len(function_responses) == 2
322+
_assert_inline_image(function_responses[0]['parts'][0])
323+
file_data = pytest_helper.get_value_ignore_key_case(
324+
function_responses[0]['parts'][1], 'file_data'
325+
)
326+
assert (
327+
pytest_helper.get_value_ignore_key_case(file_data, 'file_uri')
328+
== 'gs://bucket/photo.png'
329+
)
330+
assert 'parts' not in function_responses[1]
331+
assert function_responses[1]['response'] == {'time': '12:00'}

0 commit comments

Comments
 (0)