fix: Fixes a bug that causes intermittent pydantic validation errors when uploading files

The root cause is an unsafe in-memory mutation. The `SaveFilesAsArtifactsPlugin` was saving a direct reference to the message part and then modifying the message list in-place. This created a race condition where downstream code could alter the original part *after* it had been saved as an artifact, leading to a corrupted state.
This CL saves a `copy.copy()` of the artifact, which create a snapshot of the data.
Also Changes the plugin to return a new `types.Content` object instead of modifying the original message in-place

PiperOrigin-RevId: 814308070
This commit is contained in:
George Weale
2025-10-02 13:43:34 -07:00
committed by Copybara-Service
parent f667c7445e
commit e68006386f
5 changed files with 113 additions and 33 deletions
+8 -2
View File
@@ -16,6 +16,7 @@
from __future__ import annotations
import contextlib
import copy
from functools import cached_property
import logging
import os
@@ -300,8 +301,13 @@ class Gemini(BaseLlm):
if not content.parts:
continue
for part in content.parts:
_remove_display_name_if_present(part.inline_data)
_remove_display_name_if_present(part.file_data)
# Create copies to avoid mutating the original objects
if part.inline_data:
part.inline_data = copy.copy(part.inline_data)
_remove_display_name_if_present(part.inline_data)
if part.file_data:
part.file_data = copy.copy(part.file_data)
_remove_display_name_if_present(part.file_data)
# Initialize config if needed
if llm_request.config and llm_request.config.tools:
@@ -14,6 +14,7 @@
from __future__ import annotations
import copy
import logging
from typing import Optional
@@ -29,14 +30,15 @@ class SaveFilesAsArtifactsPlugin(BasePlugin):
"""A plugin that saves files embedded in user messages as artifacts.
This is useful to allow users to upload files in the chat experience and have
those files available to the agent.
those files available to the agent within the current session.
We use Blob.display_name to determine
the file name. Artifacts with the same name will be overwritten. A placeholder
with the artifact name will be put in place of the embedded file in the user
message so the model knows where to find the file. You may want to add
load_artifacts tool to the agent, or load the artifacts in your own tool to
use the files.
We use Blob.display_name to determine the file name. By default, artifacts are
session-scoped. For cross-session persistence, prefix the filename with
"user:".
Artifacts with the same name will be overwritten. A placeholder with the
artifact name will be put in place of the embedded file in the user message
so the model knows where to find the file. You may want to add load_artifacts
tool to the agent, or load the artifacts in your own tool to use the files.
"""
def __init__(self, name: str = 'save_files_as_artifacts_plugin'):
@@ -62,10 +64,14 @@ class SaveFilesAsArtifactsPlugin(BasePlugin):
return user_message
if not user_message.parts:
return user_message
return None
new_parts = []
modified = False
for i, part in enumerate(user_message.parts):
if part.inline_data is None:
new_parts.append(part)
continue
try:
@@ -77,23 +83,32 @@ class SaveFilesAsArtifactsPlugin(BasePlugin):
f'No display_name found, using generated filename: {file_name}'
)
# Store original filename for display to user/ placeholder
display_name = file_name
# Create a copy to stop mutation of the saved artifact if the original part is modified
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=part,
artifact=copy.copy(part),
)
# Replace the inline data with a placeholder text
user_message.parts[i] = types.Part(
text=f'[Uploaded Artifact: "{file_name}"]'
# Replace the inline data with a placeholder text (using the clean name)
new_parts.append(
types.Part(text=f'[Uploaded Artifact: "{display_name}"]')
)
modified = True
logger.info(f'Successfully saved artifact: {file_name}')
except Exception as e:
logger.error(f'Failed to save artifact for part {i}: {e}')
# Keep the original part if saving fails
new_parts.append(part)
continue
return user_message
if modified:
return types.Content(role=user_message.role, parts=new_parts)
else:
return None
@@ -15,6 +15,7 @@
from __future__ import annotations
import json
import logging
from typing import Any
from typing import TYPE_CHECKING
@@ -27,6 +28,8 @@ if TYPE_CHECKING:
from ..models.llm_request import LlmRequest
from .tool_context import ToolContext
logger = logging.getLogger('google_adk.' + __name__)
class LoadArtifactsTool(BaseTool):
"""A tool that loads the artifacts and adds them to the session."""
@@ -96,7 +99,18 @@ class LoadArtifactsTool(BaseTool):
if function_response and function_response.name == 'load_artifacts':
artifact_names = function_response.response['artifact_names']
for artifact_name in artifact_names:
# Try session-scoped first (default behavior)
artifact = await tool_context.load_artifact(artifact_name)
# If not found and name doesn't already have user: prefix,
# try cross-session artifacts with user: prefix
if artifact is None and not artifact_name.startswith('user:'):
prefixed_name = f'user:{artifact_name}'
artifact = await tool_context.load_artifact(prefixed_name)
if artifact is None:
logger.warning('Artifact "%s" not found, skipping', artifact_name)
continue
llm_request.contents.append(
types.Content(
role='user',
@@ -277,3 +277,48 @@ async def test_list_versions(service_type):
)
assert response_versions == list(range(4))
@pytest.mark.asyncio
async def test_list_keys_preserves_user_prefix():
"""Tests that list_artifact_keys preserves 'user:' prefix in returned names."""
artifact_service = InMemoryArtifactService()
artifact = types.Part.from_bytes(data=b"test_data", mime_type="text/plain")
app_name = "app0"
user_id = "user0"
session_id = "123"
# Save artifacts with "user:" prefix (cross-session artifacts)
await artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename="user:document.pdf",
artifact=artifact,
)
await artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename="user:image.png",
artifact=artifact,
)
# Save session-scoped artifact without prefix
await artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename="session_file.txt",
artifact=artifact,
)
# List artifacts should return names with "user:" prefix for user-scoped artifacts
artifact_keys = await artifact_service.list_artifact_keys(
app_name=app_name, user_id=user_id, session_id=session_id
)
# Should contain prefixed names and session file
expected_keys = ["user:document.pdf", "user:image.png", "session_file.txt"]
assert sorted(artifact_keys) == sorted(expected_keys)
@@ -57,7 +57,7 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message
)
# Verify artifact was saved with correct filename
# Verify artifact was saved with correct filename (session-scoped by default)
self.mock_context.artifact_service.save_artifact.assert_called_once_with(
app_name="test_app",
user_id="test_user",
@@ -66,7 +66,7 @@ class TestSaveFilesAsArtifactsPlugin:
artifact=original_part,
)
# Verify message was modified with placeholder
# Verify message was modified with placeholder (clean name)
assert result.parts[0].text == '[Uploaded Artifact: "test_document.pdf"]'
@pytest.mark.asyncio
@@ -85,7 +85,7 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message
)
# Verify artifact was saved with generated filename
# Verify artifact was saved with generated filename (session-scoped by default)
expected_filename = "artifact_test_invocation_123_0"
self.mock_context.artifact_service.save_artifact.assert_called_once_with(
app_name="test_app",
@@ -95,8 +95,12 @@ class TestSaveFilesAsArtifactsPlugin:
artifact=original_part,
)
# Verify message was modified with generated filename
assert result.parts[0].text == f'[Uploaded Artifact: "{expected_filename}"]'
# Verify message was modified with generated filename (clean name)
generated_display_name = "artifact_test_invocation_123_0"
assert (
result.parts[0].text
== f'[Uploaded Artifact: "{generated_display_name}"]'
)
@pytest.mark.asyncio
async def test_multiple_files_in_message(self):
@@ -138,7 +142,7 @@ class TestSaveFilesAsArtifactsPlugin:
)
assert second_call[1]["filename"] == "file2.jpg"
# Verify message parts were modified correctly
# Verify message parts were modified correctly (clean names)
assert result.parts[0].text == '[Uploaded Artifact: "file1.txt"]'
assert result.parts[1].text == "Some text between files" # Unchanged
assert result.parts[2].text == '[Uploaded Artifact: "file2.jpg"]'
@@ -174,9 +178,8 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message
)
# Should return original message unchanged
assert result == user_message
assert result.parts == []
# Should return None to proceed with original message
assert result is None
# Should not try to save any artifacts
self.mock_context.artifact_service.save_artifact.assert_not_called()
@@ -193,10 +196,8 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message
)
# Should return original message unchanged
assert result == user_message
assert result.parts[0].text == "Hello world"
assert result.parts[1].text == "No files here"
# Should return None to proceed with original message
assert result is None
# Should not try to save any artifacts
self.mock_context.artifact_service.save_artifact.assert_not_called()
@@ -221,9 +222,8 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message
)
# Should preserve original part when saving fails
assert result.parts[0] == original_part
assert result.parts[0].inline_data == inline_data
# Should return None when saving fails (no modifications made)
assert result is None
@pytest.mark.asyncio
async def test_mixed_success_and_failure(self):
@@ -264,7 +264,7 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message
)
# First file should be replaced with placeholder
# First file should be replaced with placeholder (clean name)
assert result.parts[0].text == '[Uploaded Artifact: "success.pdf"]'
# Second file should remain unchanged due to failure
@@ -287,7 +287,7 @@ class TestSaveFilesAsArtifactsPlugin:
invocation_context=self.mock_context, user_message=user_message
)
# Verify exact format of placeholder text
# Verify exact format of placeholder text (clean name)
expected_text = '[Uploaded Artifact: "test file with spaces.docx"]'
assert result.parts[0].text == expected_text