Skip to content

Commit 657369c

Browse files
GWealecopybara-github
authored andcommitted
fix: Adds plugin to save artifacts for issue #2176
PiperOrigin-RevId: 810522939
1 parent c944a12 commit 657369c

6 files changed

Lines changed: 419 additions & 3 deletions

File tree

contributing/samples/hello_world_app/agent.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
from google.adk.models.llm_request import LlmRequest
2222
from google.adk.plugins.base_plugin import BasePlugin
2323
from google.adk.plugins.context_filter_plugin import ContextFilterPlugin
24+
from google.adk.plugins.save_files_as_artifacts_plugin import SaveFilesAsArtifactsPlugin
25+
from google.adk.tools import load_artifacts
2426
from google.adk.tools.tool_context import ToolContext
2527
from google.genai import types
2628

@@ -97,6 +99,7 @@ async def check_prime(nums: list[int]) -> str:
9799
tools=[
98100
roll_die,
99101
check_prime,
102+
load_artifacts,
100103
],
101104
# planner=BuiltInPlanner(
102105
# thinking_config=types.ThinkingConfig(
@@ -145,5 +148,6 @@ async def before_model_callback(
145148
plugins=[
146149
CountInvocationPlugin(),
147150
ContextFilterPlugin(num_invocations_to_keep=3),
151+
SaveFilesAsArtifactsPlugin(),
148152
],
149153
)

contributing/samples/hello_world_app/main.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ async def run_prompt_bytes(session: Session, new_message: str):
6565
user_id=user_id_1,
6666
session_id=session.id,
6767
new_message=content,
68-
run_config=RunConfig(save_input_blobs_as_artifacts=True),
68+
run_config=RunConfig(save_input_blobs_as_artifacts=False),
6969
):
7070
if event.content.parts and event.content.parts[0].text:
7171
print(f'** {event.author}: {event.content.parts[0].text}')

src/google/adk/agents/run_config.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,15 @@ class RunConfig(BaseModel):
4848
response_modalities: Optional[list[str]] = None
4949
"""The output modalities. If not set, it's default to AUDIO."""
5050

51-
save_input_blobs_as_artifacts: bool = False
52-
"""Whether or not to save the input blobs as artifacts."""
51+
save_input_blobs_as_artifacts: bool = Field(
52+
default=False,
53+
deprecated=True,
54+
description=(
55+
'Whether or not to save the input blobs as artifacts. DEPRECATED: Use'
56+
' SaveFilesAsArtifactsPlugin instead for better control and'
57+
' flexibility. See google.adk.plugins.SaveFilesAsArtifactsPlugin.'
58+
),
59+
)
5360

5461
support_cfc: bool = False
5562
"""
Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
import logging
18+
from typing import Optional
19+
20+
from google.genai import types
21+
22+
from ..agents.invocation_context import InvocationContext
23+
from .base_plugin import BasePlugin
24+
25+
logger = logging.getLogger('google_adk.' + __name__)
26+
27+
28+
class SaveFilesAsArtifactsPlugin(BasePlugin):
29+
"""A plugin that saves files embedded in user messages as artifacts.
30+
31+
This is useful to allow users to upload files in the chat experience and have
32+
those files available to the agent.
33+
34+
We use Blob.display_name to determine
35+
the file name. Artifacts with the same name will be overwritten. A placeholder
36+
with the artifact name will be put in place of the embedded file in the user
37+
message so the model knows where to find the file. You may want to add
38+
load_artifacts tool to the agent, or load the artifacts in your own tool to
39+
use the files.
40+
"""
41+
42+
def __init__(self, name: str = 'save_files_as_artifacts_plugin'):
43+
"""Initialize the save files as artifacts plugin.
44+
45+
Args:
46+
name: The name of the plugin instance.
47+
"""
48+
super().__init__(name)
49+
50+
async def on_user_message_callback(
51+
self,
52+
*,
53+
invocation_context: InvocationContext,
54+
user_message: types.Content,
55+
) -> Optional[types.Content]:
56+
"""Process user message and save any attached files as artifacts."""
57+
if not invocation_context.artifact_service:
58+
logger.warning(
59+
'Artifact service is not set. SaveFilesAsArtifactsPlugin'
60+
' will not be enabled.'
61+
)
62+
return user_message
63+
64+
if not user_message.parts:
65+
return user_message
66+
67+
for i, part in enumerate(user_message.parts):
68+
if part.inline_data is None:
69+
continue
70+
71+
try:
72+
# Use display_name if available, otherwise generate a filename
73+
file_name = part.inline_data.display_name
74+
if not file_name:
75+
file_name = f'artifact_{invocation_context.invocation_id}_{i}'
76+
logger.info(
77+
f'No display_name found, using generated filename: {file_name}'
78+
)
79+
80+
await invocation_context.artifact_service.save_artifact(
81+
app_name=invocation_context.app_name,
82+
user_id=invocation_context.user_id,
83+
session_id=invocation_context.session.id,
84+
filename=file_name,
85+
artifact=part,
86+
)
87+
88+
# Replace the inline data with a placeholder text
89+
user_message.parts[i] = types.Part(
90+
text=f'[Uploaded Artifact: "{file_name}"]'
91+
)
92+
logger.info(f'Successfully saved artifact: {file_name}')
93+
94+
except Exception as e:
95+
logger.error(f'Failed to save artifact for part {i}: {e}')
96+
# Keep the original part if saving fails
97+
continue
98+
99+
return user_message

src/google/adk/runners.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -419,6 +419,15 @@ async def _append_new_message_to_session(
419419
raise ValueError('No parts in the new_message.')
420420

421421
if self.artifact_service and save_input_blobs_as_artifacts:
422+
# Issue deprecation warning
423+
warnings.warn(
424+
"The 'save_input_blobs_as_artifacts' parameter is deprecated. Use"
425+
' SaveFilesAsArtifactsPlugin instead for better control and'
426+
' flexibility. See google.adk.plugins.SaveFilesAsArtifactsPlugin for'
427+
' migration guidance.',
428+
DeprecationWarning,
429+
stacklevel=3,
430+
)
422431
# The runner directly saves the artifacts (if applicable) in the
423432
# user message and replaces the artifact data with a file name
424433
# placeholder.

0 commit comments

Comments
 (0)