mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
f667c7445e
commit
e68006386f
@@ -16,6 +16,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import contextlib
|
import contextlib
|
||||||
|
import copy
|
||||||
from functools import cached_property
|
from functools import cached_property
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
@@ -300,8 +301,13 @@ class Gemini(BaseLlm):
|
|||||||
if not content.parts:
|
if not content.parts:
|
||||||
continue
|
continue
|
||||||
for part in content.parts:
|
for part in content.parts:
|
||||||
_remove_display_name_if_present(part.inline_data)
|
# Create copies to avoid mutating the original objects
|
||||||
_remove_display_name_if_present(part.file_data)
|
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
|
# Initialize config if needed
|
||||||
if llm_request.config and llm_request.config.tools:
|
if llm_request.config and llm_request.config.tools:
|
||||||
|
|||||||
@@ -14,6 +14,7 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import copy
|
||||||
import logging
|
import logging
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
@@ -29,14 +30,15 @@ class SaveFilesAsArtifactsPlugin(BasePlugin):
|
|||||||
"""A plugin that saves files embedded in user messages as artifacts.
|
"""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
|
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
|
We use Blob.display_name to determine the file name. By default, artifacts are
|
||||||
the file name. Artifacts with the same name will be overwritten. A placeholder
|
session-scoped. For cross-session persistence, prefix the filename with
|
||||||
with the artifact name will be put in place of the embedded file in the user
|
"user:".
|
||||||
message so the model knows where to find the file. You may want to add
|
Artifacts with the same name will be overwritten. A placeholder with the
|
||||||
load_artifacts tool to the agent, or load the artifacts in your own tool to
|
artifact name will be put in place of the embedded file in the user message
|
||||||
use the files.
|
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'):
|
def __init__(self, name: str = 'save_files_as_artifacts_plugin'):
|
||||||
@@ -62,10 +64,14 @@ class SaveFilesAsArtifactsPlugin(BasePlugin):
|
|||||||
return user_message
|
return user_message
|
||||||
|
|
||||||
if not user_message.parts:
|
if not user_message.parts:
|
||||||
return user_message
|
return None
|
||||||
|
|
||||||
|
new_parts = []
|
||||||
|
modified = False
|
||||||
|
|
||||||
for i, part in enumerate(user_message.parts):
|
for i, part in enumerate(user_message.parts):
|
||||||
if part.inline_data is None:
|
if part.inline_data is None:
|
||||||
|
new_parts.append(part)
|
||||||
continue
|
continue
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -77,23 +83,32 @@ class SaveFilesAsArtifactsPlugin(BasePlugin):
|
|||||||
f'No display_name found, using generated filename: {file_name}'
|
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(
|
await invocation_context.artifact_service.save_artifact(
|
||||||
app_name=invocation_context.app_name,
|
app_name=invocation_context.app_name,
|
||||||
user_id=invocation_context.user_id,
|
user_id=invocation_context.user_id,
|
||||||
session_id=invocation_context.session.id,
|
session_id=invocation_context.session.id,
|
||||||
filename=file_name,
|
filename=file_name,
|
||||||
artifact=part,
|
artifact=copy.copy(part),
|
||||||
)
|
)
|
||||||
|
|
||||||
# Replace the inline data with a placeholder text
|
# Replace the inline data with a placeholder text (using the clean name)
|
||||||
user_message.parts[i] = types.Part(
|
new_parts.append(
|
||||||
text=f'[Uploaded Artifact: "{file_name}"]'
|
types.Part(text=f'[Uploaded Artifact: "{display_name}"]')
|
||||||
)
|
)
|
||||||
|
modified = True
|
||||||
logger.info(f'Successfully saved artifact: {file_name}')
|
logger.info(f'Successfully saved artifact: {file_name}')
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f'Failed to save artifact for part {i}: {e}')
|
logger.error(f'Failed to save artifact for part {i}: {e}')
|
||||||
# Keep the original part if saving fails
|
# Keep the original part if saving fails
|
||||||
|
new_parts.append(part)
|
||||||
continue
|
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
|
from __future__ import annotations
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
import logging
|
||||||
from typing import Any
|
from typing import Any
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
@@ -27,6 +28,8 @@ if TYPE_CHECKING:
|
|||||||
from ..models.llm_request import LlmRequest
|
from ..models.llm_request import LlmRequest
|
||||||
from .tool_context import ToolContext
|
from .tool_context import ToolContext
|
||||||
|
|
||||||
|
logger = logging.getLogger('google_adk.' + __name__)
|
||||||
|
|
||||||
|
|
||||||
class LoadArtifactsTool(BaseTool):
|
class LoadArtifactsTool(BaseTool):
|
||||||
"""A tool that loads the artifacts and adds them to the session."""
|
"""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':
|
if function_response and function_response.name == 'load_artifacts':
|
||||||
artifact_names = function_response.response['artifact_names']
|
artifact_names = function_response.response['artifact_names']
|
||||||
for artifact_name in artifact_names:
|
for artifact_name in artifact_names:
|
||||||
|
# Try session-scoped first (default behavior)
|
||||||
artifact = await tool_context.load_artifact(artifact_name)
|
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(
|
llm_request.contents.append(
|
||||||
types.Content(
|
types.Content(
|
||||||
role='user',
|
role='user',
|
||||||
|
|||||||
@@ -277,3 +277,48 @@ async def test_list_versions(service_type):
|
|||||||
)
|
)
|
||||||
|
|
||||||
assert response_versions == list(range(4))
|
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
|
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(
|
self.mock_context.artifact_service.save_artifact.assert_called_once_with(
|
||||||
app_name="test_app",
|
app_name="test_app",
|
||||||
user_id="test_user",
|
user_id="test_user",
|
||||||
@@ -66,7 +66,7 @@ class TestSaveFilesAsArtifactsPlugin:
|
|||||||
artifact=original_part,
|
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"]'
|
assert result.parts[0].text == '[Uploaded Artifact: "test_document.pdf"]'
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@@ -85,7 +85,7 @@ class TestSaveFilesAsArtifactsPlugin:
|
|||||||
invocation_context=self.mock_context, user_message=user_message
|
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"
|
expected_filename = "artifact_test_invocation_123_0"
|
||||||
self.mock_context.artifact_service.save_artifact.assert_called_once_with(
|
self.mock_context.artifact_service.save_artifact.assert_called_once_with(
|
||||||
app_name="test_app",
|
app_name="test_app",
|
||||||
@@ -95,8 +95,12 @@ class TestSaveFilesAsArtifactsPlugin:
|
|||||||
artifact=original_part,
|
artifact=original_part,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Verify message was modified with generated filename
|
# Verify message was modified with generated filename (clean name)
|
||||||
assert result.parts[0].text == f'[Uploaded Artifact: "{expected_filename}"]'
|
generated_display_name = "artifact_test_invocation_123_0"
|
||||||
|
assert (
|
||||||
|
result.parts[0].text
|
||||||
|
== f'[Uploaded Artifact: "{generated_display_name}"]'
|
||||||
|
)
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_multiple_files_in_message(self):
|
async def test_multiple_files_in_message(self):
|
||||||
@@ -138,7 +142,7 @@ class TestSaveFilesAsArtifactsPlugin:
|
|||||||
)
|
)
|
||||||
assert second_call[1]["filename"] == "file2.jpg"
|
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[0].text == '[Uploaded Artifact: "file1.txt"]'
|
||||||
assert result.parts[1].text == "Some text between files" # Unchanged
|
assert result.parts[1].text == "Some text between files" # Unchanged
|
||||||
assert result.parts[2].text == '[Uploaded Artifact: "file2.jpg"]'
|
assert result.parts[2].text == '[Uploaded Artifact: "file2.jpg"]'
|
||||||
@@ -174,9 +178,8 @@ class TestSaveFilesAsArtifactsPlugin:
|
|||||||
invocation_context=self.mock_context, user_message=user_message
|
invocation_context=self.mock_context, user_message=user_message
|
||||||
)
|
)
|
||||||
|
|
||||||
# Should return original message unchanged
|
# Should return None to proceed with original message
|
||||||
assert result == user_message
|
assert result is None
|
||||||
assert result.parts == []
|
|
||||||
|
|
||||||
# Should not try to save any artifacts
|
# Should not try to save any artifacts
|
||||||
self.mock_context.artifact_service.save_artifact.assert_not_called()
|
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
|
invocation_context=self.mock_context, user_message=user_message
|
||||||
)
|
)
|
||||||
|
|
||||||
# Should return original message unchanged
|
# Should return None to proceed with original message
|
||||||
assert result == user_message
|
assert result is None
|
||||||
assert result.parts[0].text == "Hello world"
|
|
||||||
assert result.parts[1].text == "No files here"
|
|
||||||
|
|
||||||
# Should not try to save any artifacts
|
# Should not try to save any artifacts
|
||||||
self.mock_context.artifact_service.save_artifact.assert_not_called()
|
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
|
invocation_context=self.mock_context, user_message=user_message
|
||||||
)
|
)
|
||||||
|
|
||||||
# Should preserve original part when saving fails
|
# Should return None when saving fails (no modifications made)
|
||||||
assert result.parts[0] == original_part
|
assert result is None
|
||||||
assert result.parts[0].inline_data == inline_data
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_mixed_success_and_failure(self):
|
async def test_mixed_success_and_failure(self):
|
||||||
@@ -264,7 +264,7 @@ class TestSaveFilesAsArtifactsPlugin:
|
|||||||
invocation_context=self.mock_context, user_message=user_message
|
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"]'
|
assert result.parts[0].text == '[Uploaded Artifact: "success.pdf"]'
|
||||||
|
|
||||||
# Second file should remain unchanged due to failure
|
# Second file should remain unchanged due to failure
|
||||||
@@ -287,7 +287,7 @@ class TestSaveFilesAsArtifactsPlugin:
|
|||||||
invocation_context=self.mock_context, user_message=user_message
|
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"]'
|
expected_text = '[Uploaded Artifact: "test file with spaces.docx"]'
|
||||||
assert result.parts[0].text == expected_text
|
assert result.parts[0].text == expected_text
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user