|
16 | 16 |
|
17 | 17 | """Tests for live.py.""" |
18 | 18 |
|
| 19 | +import base64 |
19 | 20 | import contextlib |
20 | 21 | import json |
21 | 22 | from unittest import mock |
@@ -220,3 +221,111 @@ async def test_missing_id(mock_websocket, vertexai): |
220 | 221 | if not vertexai: |
221 | 222 | with pytest.raises(ValueError, match=".*must have.*"): |
222 | 223 | 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