1717import copy
1818import logging
1919from typing import Optional
20+ import urllib .parse
2021
2122from google .genai import types
2223
2526
2627logger = 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
2935class 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