mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
Kathy WuandCopybara-Service
d8c66fe688
Since it appears in the same bubble as the rest of the LLM's response text, make it more human readable so it doesn't look out of place. PiperOrigin-RevId: 822729061
108 lines
3.8 KiB
Python
108 lines
3.8 KiB
Python
# 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.
|
|
|
|
"""Unit tests for Code Execution logic."""
|
|
|
|
import datetime
|
|
from unittest.mock import AsyncMock
|
|
from unittest.mock import MagicMock
|
|
from unittest.mock import patch
|
|
|
|
from google.adk.agents.llm_agent import Agent
|
|
from google.adk.code_executors.built_in_code_executor import BuiltInCodeExecutor
|
|
from google.adk.flows.llm_flows._code_execution import response_processor
|
|
from google.adk.models.llm_response import LlmResponse
|
|
from google.genai import types
|
|
import pytest
|
|
|
|
from ... import testing_utils
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch('google.adk.flows.llm_flows._code_execution.datetime')
|
|
async def test_builtin_code_executor_image_artifact_creation(mock_datetime):
|
|
"""Test BuiltInCodeExecutor creates artifacts for images in response."""
|
|
mock_now = datetime.datetime(2025, 1, 1, 12, 0, 0)
|
|
mock_datetime.datetime.now.return_value.astimezone.return_value = mock_now
|
|
code_executor = BuiltInCodeExecutor()
|
|
agent = Agent(name='test_agent', code_executor=code_executor)
|
|
invocation_context = await testing_utils.create_invocation_context(
|
|
agent=agent, user_content='test message'
|
|
)
|
|
invocation_context.artifact_service = MagicMock()
|
|
invocation_context.artifact_service.save_artifact = AsyncMock(
|
|
return_value='v1'
|
|
)
|
|
llm_response = LlmResponse(
|
|
content=types.Content(
|
|
parts=[
|
|
types.Part(
|
|
inline_data=types.Blob(
|
|
mime_type='image/png',
|
|
data=b'image1',
|
|
display_name='image_1.png',
|
|
)
|
|
),
|
|
types.Part(text='this is text'),
|
|
types.Part(
|
|
inline_data=types.Blob(mime_type='image/jpeg', data=b'image2')
|
|
),
|
|
]
|
|
)
|
|
)
|
|
|
|
events = []
|
|
async for event in response_processor.run_async(
|
|
invocation_context, llm_response
|
|
):
|
|
events.append(event)
|
|
|
|
expected_timestamp = mock_now.strftime('%Y%m%d_%H%M%S')
|
|
expected_filename2 = f'{expected_timestamp}.jpeg'
|
|
|
|
assert invocation_context.artifact_service.save_artifact.call_count == 2
|
|
invocation_context.artifact_service.save_artifact.assert_any_call(
|
|
app_name=invocation_context.app_name,
|
|
user_id=invocation_context.user_id,
|
|
session_id=invocation_context.session.id,
|
|
filename='image_1.png',
|
|
artifact=types.Part.from_bytes(data=b'image1', mime_type='image/png'),
|
|
)
|
|
invocation_context.artifact_service.save_artifact.assert_any_call(
|
|
app_name=invocation_context.app_name,
|
|
user_id=invocation_context.user_id,
|
|
session_id=invocation_context.session.id,
|
|
filename=expected_filename2,
|
|
artifact=types.Part.from_bytes(data=b'image2', mime_type='image/jpeg'),
|
|
)
|
|
|
|
assert len(events) == 1
|
|
assert events[0].actions.artifact_delta == {
|
|
'image_1.png': 'v1',
|
|
expected_filename2: 'v1',
|
|
}
|
|
assert not events[0].content
|
|
assert llm_response.content is not None
|
|
assert len(llm_response.content.parts) == 3
|
|
assert (
|
|
llm_response.content.parts[0].text == 'Saved as artifact: image_1.png. '
|
|
)
|
|
assert not llm_response.content.parts[0].inline_data
|
|
assert llm_response.content.parts[1].text == 'this is text'
|
|
assert (
|
|
llm_response.content.parts[2].text
|
|
== f'Saved as artifact: {expected_filename2}. '
|
|
)
|
|
assert not llm_response.content.parts[2].inline_data
|