diff --git a/src/google/adk/flows/llm_flows/_code_execution.py b/src/google/adk/flows/llm_flows/_code_execution.py index 5c0a5777..d37e773c 100644 --- a/src/google/adk/flows/llm_flows/_code_execution.py +++ b/src/google/adk/flows/llm_flows/_code_execution.py @@ -19,6 +19,7 @@ from __future__ import annotations import base64 import copy import dataclasses +import datetime import os import re from typing import AsyncGenerator @@ -275,6 +276,43 @@ async def _run_post_processor( return if isinstance(code_executor, BuiltInCodeExecutor): + event_actions = EventActions() + + # If an image is generated, save it to the artifact service and add it to + # the event actions. + for part in llm_response.content.parts: + if part.inline_data and part.inline_data.mime_type.startswith('image/'): + if invocation_context.artifact_service is None: + raise ValueError('Artifact service is not initialized.') + + if part.inline_data.display_name: + file_name = part.inline_data.display_name + else: + now = datetime.datetime.now().astimezone() + timestamp = now.strftime('%Y%m%d_%H%M%S') + file_extension = part.inline_data.mime_type.split('/')[-1] + file_name = f'{timestamp}.{file_extension}' + + version = await invocation_context.artifact_service.save_artifact( + app_name=invocation_context.app_name, + user_id=invocation_context.user_id, + session_id=invocation_context.session.id, + filename=file_name, + artifact=types.Part.from_bytes( + data=part.inline_data.data, + mime_type=part.inline_data.mime_type, + ), + ) + event_actions.artifact_delta[file_name] = version + part.inline_data = None + part.text = f'artifact: {file_name}' + + yield Event( + invocation_id=invocation_context.invocation_id, + author=agent.name, + branch=invocation_context.branch, + actions=event_actions, + ) return code_executor_context = CodeExecutorContext(invocation_context.session.state) diff --git a/tests/unittests/flows/llm_flows/test_code_execution.py b/tests/unittests/flows/llm_flows/test_code_execution.py new file mode 100644 index 00000000..4212e19a --- /dev/null +++ b/tests/unittests/flows/llm_flows/test_code_execution.py @@ -0,0 +1,102 @@ +# 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 == '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'artifact: {expected_filename2}' + assert not llm_response.content.parts[2].inline_data