Skip to content

Commit 64a1208

Browse files
Aakash SAakash S
authored andcommitted
fix(plugins): Add 20MB file size validation to SaveFilesAsArtifactsPlugin with clear error messages and Files API guidance. Fixes #3751
1 parent 8da61be commit 64a1208

File tree

2 files changed

+149
-1
lines changed

2 files changed

+149
-1
lines changed

src/google/adk/plugins/save_files_as_artifacts_plugin.py

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,10 @@
3131
# capabilities.
3232
_MODEL_ACCESSIBLE_URI_SCHEMES = {'gs', 'https', 'http'}
3333

34+
# Maximum file size for inline_data (20MB as per Gemini API documentation)
35+
# https://ai.google.dev/gemini-api/docs/files
36+
_MAX_INLINE_DATA_SIZE_BYTES = 20 * 1024 * 1024 # 20 MB
37+
3438

3539
class SaveFilesAsArtifactsPlugin(BasePlugin):
3640
"""A plugin that saves files embedded in user messages as artifacts.
@@ -81,8 +85,28 @@ async def on_user_message_callback(
8185
continue
8286

8387
try:
84-
# Use display_name if available, otherwise generate a filename
88+
# Check file size before processing
8589
inline_data = part.inline_data
90+
file_size = len(inline_data.data) if inline_data.data else 0
91+
92+
if file_size > _MAX_INLINE_DATA_SIZE_BYTES:
93+
# File exceeds the inline_data limit
94+
file_size_mb = file_size / (1024 * 1024)
95+
limit_mb = _MAX_INLINE_DATA_SIZE_BYTES / (1024 * 1024)
96+
error_message = (
97+
f'File size ({file_size_mb:.2f} MB) exceeds the maximum allowed'
98+
f' size for inline uploads ({limit_mb:.0f} MB). Please use the'
99+
' Files API to upload files larger than 20 MB. See'
100+
' https://ai.google.dev/gemini-api/docs/files for more'
101+
' information.'
102+
)
103+
logger.error(error_message)
104+
# Replace with error message part
105+
new_parts.append(types.Part(text=f'[Upload Error: {error_message}]'))
106+
modified = True
107+
continue
108+
109+
# Use display_name if available, otherwise generate a filename
86110
file_name = inline_data.display_name
87111
if not file_name:
88112
file_name = f'artifact_{invocation_context.invocation_id}_{i}'

tests/unittests/plugins/test_save_files_as_artifacts.py

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -303,3 +303,127 @@ def test_plugin_name_default(self):
303303
"""Test that plugin has correct default name."""
304304
plugin = SaveFilesAsArtifactsPlugin()
305305
assert plugin.name == "save_files_as_artifacts_plugin"
306+
307+
@pytest.mark.asyncio
308+
async def test_file_size_exceeds_limit(self):
309+
"""Test that files exceeding 20MB limit show error message."""
310+
# Create a file larger than 20MB (20 * 1024 * 1024 bytes)
311+
large_file_data = b"x" * (21 * 1024 * 1024) # 21 MB
312+
inline_data = types.Blob(
313+
display_name="large_file.pdf",
314+
data=large_file_data,
315+
mime_type="application/pdf",
316+
)
317+
318+
user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
319+
320+
result = await self.plugin.on_user_message_callback(
321+
invocation_context=self.mock_context, user_message=user_message
322+
)
323+
324+
# Should not try to save the artifact
325+
self.mock_context.artifact_service.save_artifact.assert_not_called()
326+
327+
# Should return modified content with error message
328+
assert result is not None
329+
assert len(result.parts) == 1
330+
assert result.parts[0].text is not None
331+
assert "[Upload Error:" in result.parts[0].text
332+
assert "21.00 MB" in result.parts[0].text
333+
assert "20 MB" in result.parts[0].text
334+
assert "Files API" in result.parts[0].text
335+
336+
@pytest.mark.asyncio
337+
async def test_file_size_at_limit(self):
338+
"""Test that files exactly at 20MB limit are processed successfully."""
339+
# Create a file exactly 20MB (20 * 1024 * 1024 bytes)
340+
file_data = b"x" * (20 * 1024 * 1024) # Exactly 20 MB
341+
inline_data = types.Blob(
342+
display_name="max_size_file.pdf",
343+
data=file_data,
344+
mime_type="application/pdf",
345+
)
346+
347+
user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
348+
349+
result = await self.plugin.on_user_message_callback(
350+
invocation_context=self.mock_context, user_message=user_message
351+
)
352+
353+
# Should save the artifact since it's at the limit
354+
self.mock_context.artifact_service.save_artifact.assert_called_once()
355+
assert result is not None
356+
assert len(result.parts) == 2
357+
assert result.parts[0].text == '[Uploaded Artifact: "max_size_file.pdf"]'
358+
assert result.parts[1].file_data is not None
359+
360+
@pytest.mark.asyncio
361+
async def test_file_size_just_over_limit(self):
362+
"""Test that files just over 20MB limit show error message."""
363+
# Create a file just over 20MB
364+
large_file_data = b"x" * (20 * 1024 * 1024 + 1) # 20 MB + 1 byte
365+
inline_data = types.Blob(
366+
display_name="slightly_too_large.pdf",
367+
data=large_file_data,
368+
mime_type="application/pdf",
369+
)
370+
371+
user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
372+
373+
result = await self.plugin.on_user_message_callback(
374+
invocation_context=self.mock_context, user_message=user_message
375+
)
376+
377+
# Should not try to save the artifact
378+
self.mock_context.artifact_service.save_artifact.assert_not_called()
379+
380+
# Should return error message
381+
assert result is not None
382+
assert len(result.parts) == 1
383+
assert "[Upload Error:" in result.parts[0].text
384+
385+
@pytest.mark.asyncio
386+
async def test_mixed_file_sizes(self):
387+
"""Test processing multiple files with mixed sizes."""
388+
# Small file (should succeed)
389+
small_file_data = b"x" * (5 * 1024 * 1024) # 5 MB
390+
small_inline_data = types.Blob(
391+
display_name="small.pdf",
392+
data=small_file_data,
393+
mime_type="application/pdf",
394+
)
395+
396+
# Large file (should fail)
397+
large_file_data = b"x" * (25 * 1024 * 1024) # 25 MB
398+
large_inline_data = types.Blob(
399+
display_name="large.pdf",
400+
data=large_file_data,
401+
mime_type="application/pdf",
402+
)
403+
404+
user_message = types.Content(
405+
parts=[
406+
types.Part(inline_data=small_inline_data),
407+
types.Part(inline_data=large_inline_data),
408+
]
409+
)
410+
411+
result = await self.plugin.on_user_message_callback(
412+
invocation_context=self.mock_context, user_message=user_message
413+
)
414+
415+
# Should only save the small file
416+
self.mock_context.artifact_service.save_artifact.assert_called_once_with(
417+
app_name="test_app",
418+
user_id="test_user",
419+
session_id="test_session",
420+
filename="small.pdf",
421+
artifact=user_message.parts[0],
422+
)
423+
424+
# Should return both success and error messages
425+
assert result is not None
426+
assert len(result.parts) == 3 # [success placeholder, file_data, error]
427+
assert '[Uploaded Artifact: "small.pdf"]' in result.parts[0].text
428+
assert result.parts[1].file_data is not None
429+
assert "[Upload Error:" in result.parts[2].text

0 commit comments

Comments
 (0)