Skip to content

Commit e3caf79

Browse files
GWealecopybara-github
authored andcommitted
feat: expose artifact URLs to the model when available
SaveFilesAsArtifactsPlugin now resolves the canonical URI for each saved artifact. When the URI is model-accessible (`gs://`, `https://`, `http://`), we add a `Part(file_data=...)` so the LLM can fetch the filedirectly while still emitting the placeholder. If no model-accessible URI exists, we retain the original inline blob alongside the placeholder to preserve access to the uploaded bytes. Close #2016 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 828786519
1 parent 52db15f commit e3caf79

File tree

2 files changed

+188
-63
lines changed

2 files changed

+188
-63
lines changed

src/google/adk/plugins/save_files_as_artifacts_plugin.py

Lines changed: 86 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import copy
1818
import logging
1919
from typing import Optional
20+
import urllib.parse
2021

2122
from google.genai import types
2223

@@ -25,6 +26,11 @@
2526

2627
logger = logging.getLogger('google_adk.' + __name__)
2728

29+
# Schemes supported by our current LLM connectors. Vertex exposes `gs://` while
30+
# hosted endpoints use HTTPS. Expand this list when BaseLlm surfaces provider
31+
# capabilities.
32+
_MODEL_ACCESSIBLE_URI_SCHEMES = {'gs', 'https', 'http'}
33+
2834

2935
class SaveFilesAsArtifactsPlugin(BasePlugin):
3036
"""A plugin that saves files embedded in user messages as artifacts.
@@ -75,8 +81,9 @@ async def on_user_message_callback(
7581
continue
7682

7783
try:
78-
# Use display_name if available; otherwise, generate a filename
79-
file_name = part.inline_data.display_name
84+
# Use display_name if available, otherwise generate a filename
85+
inline_data = part.inline_data
86+
file_name = inline_data.display_name
8087
if not file_name:
8188
file_name = f'artifact_{invocation_context.invocation_id}_{i}'
8289
logger.info(
@@ -87,18 +94,36 @@ async def on_user_message_callback(
8794
display_name = file_name
8895

8996
# Create a copy to stop mutation of the saved artifact if the original part is modified
90-
await invocation_context.artifact_service.save_artifact(
97+
version = await invocation_context.artifact_service.save_artifact(
9198
app_name=invocation_context.app_name,
9299
user_id=invocation_context.user_id,
93100
session_id=invocation_context.session.id,
94101
filename=file_name,
95102
artifact=copy.copy(part),
96103
)
97104

98-
# Replace the inline data with a placeholder text (using the clean name)
99-
new_parts.append(
100-
types.Part(text=f'[Uploaded Artifact: "{display_name}"]')
105+
placeholder_part = types.Part(
106+
text=f'[Uploaded Artifact: "{display_name}"]'
107+
)
108+
new_parts.append(placeholder_part)
109+
110+
file_part = await self._build_file_reference_part(
111+
invocation_context=invocation_context,
112+
filename=file_name,
113+
version=version,
114+
mime_type=inline_data.mime_type,
115+
display_name=display_name,
101116
)
117+
if file_part:
118+
new_parts.append(file_part)
119+
else:
120+
logger.debug(
121+
'Artifact %s is not exposed via a model-accessible URI; keeping'
122+
' inline data in user message.',
123+
file_name,
124+
)
125+
new_parts.append(part)
126+
102127
modified = True
103128
logger.info(f'Successfully saved artifact: {file_name}')
104129

@@ -112,3 +137,58 @@ async def on_user_message_callback(
112137
return types.Content(role=user_message.role, parts=new_parts)
113138
else:
114139
return None
140+
141+
async def _build_file_reference_part(
142+
self,
143+
*,
144+
invocation_context: InvocationContext,
145+
filename: str,
146+
version: int,
147+
mime_type: Optional[str],
148+
display_name: str,
149+
) -> Optional[types.Part]:
150+
"""Constructs a file reference part if the artifact URI is model-accessible."""
151+
152+
artifact_service = invocation_context.artifact_service
153+
if not artifact_service:
154+
return None
155+
156+
try:
157+
artifact_version = await artifact_service.get_artifact_version(
158+
app_name=invocation_context.app_name,
159+
user_id=invocation_context.user_id,
160+
session_id=invocation_context.session.id,
161+
filename=filename,
162+
version=version,
163+
)
164+
except Exception as exc: # pylint: disable=broad-except
165+
logger.warning(
166+
'Failed to resolve artifact version for %s: %s', filename, exc
167+
)
168+
return None
169+
170+
if (
171+
not artifact_version
172+
or not artifact_version.canonical_uri
173+
or not _is_model_accessible_uri(artifact_version.canonical_uri)
174+
):
175+
return None
176+
177+
file_data = types.FileData(
178+
file_uri=artifact_version.canonical_uri,
179+
mime_type=mime_type or artifact_version.mime_type,
180+
display_name=display_name,
181+
)
182+
return types.Part(file_data=file_data)
183+
184+
185+
def _is_model_accessible_uri(uri: str) -> bool:
186+
try:
187+
parsed = urllib.parse.urlparse(uri)
188+
except ValueError:
189+
return False
190+
191+
if not parsed.scheme:
192+
return False
193+
194+
return parsed.scheme.lower() in _MODEL_ACCESSIBLE_URI_SCHEMES

0 commit comments

Comments
 (0)