fix: Convert unsupported inline artifact MIME types to text in LoadArtifactsTool

The LoadArtifactsTool now checks if an artifact's inline data MIME type is supported by Gemini. If not, it attempts to convert the artifact content into a text Part

Close #4028

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 856404510
This commit is contained in:
George Weale
2026-01-14 16:01:06 -08:00
committed by Copybara-Service
parent 7b035aa9fc
commit fdc98d5c92
2 changed files with 265 additions and 2 deletions
+103 -2
View File
@@ -14,6 +14,8 @@
from __future__ import annotations
import base64
import binascii
import json
import logging
from typing import Any
@@ -24,6 +26,19 @@ from typing_extensions import override
from .base_tool import BaseTool
# MIME types Gemini accepts for inline data in requests.
_GEMINI_SUPPORTED_INLINE_MIME_PREFIXES = (
'image/',
'audio/',
'video/',
)
_GEMINI_SUPPORTED_INLINE_MIME_TYPES = frozenset({'application/pdf'})
_TEXT_LIKE_MIME_TYPES = frozenset({
'application/csv',
'application/json',
'application/xml',
})
if TYPE_CHECKING:
from ..models.llm_request import LlmRequest
from .tool_context import ToolContext
@@ -31,6 +46,79 @@ if TYPE_CHECKING:
logger = logging.getLogger('google_adk.' + __name__)
def _normalize_mime_type(mime_type: str | None) -> str | None:
"""Returns the normalized MIME type, without parameters like charset."""
if not mime_type:
return None
return mime_type.split(';', 1)[0].strip()
def _is_inline_mime_type_supported(mime_type: str | None) -> bool:
"""Returns True if Gemini accepts this MIME type as inline data."""
normalized = _normalize_mime_type(mime_type)
if not normalized:
return False
return normalized.startswith(_GEMINI_SUPPORTED_INLINE_MIME_PREFIXES) or (
normalized in _GEMINI_SUPPORTED_INLINE_MIME_TYPES
)
def _maybe_base64_to_bytes(data: str) -> bytes | None:
"""Best-effort base64 decode for both std and urlsafe formats."""
try:
return base64.b64decode(data, validate=True)
except (binascii.Error, ValueError):
try:
return base64.urlsafe_b64decode(data)
except (binascii.Error, ValueError):
return None
def _as_safe_part_for_llm(
artifact: types.Part, artifact_name: str
) -> types.Part:
"""Returns a Part that is safe to send to Gemini."""
inline_data = artifact.inline_data
if inline_data is None:
return artifact
if _is_inline_mime_type_supported(inline_data.mime_type):
return artifact
mime_type = _normalize_mime_type(inline_data.mime_type) or (
'application/octet-stream'
)
data = inline_data.data
if data is None:
return types.Part.from_text(
text=(
f'[Artifact: {artifact_name}, type: {mime_type}. '
'No inline data was provided.]'
)
)
if isinstance(data, str):
decoded = _maybe_base64_to_bytes(data)
if decoded is None:
return types.Part.from_text(text=data)
data = decoded
if mime_type.startswith('text/') or mime_type in _TEXT_LIKE_MIME_TYPES:
try:
return types.Part.from_text(text=data.decode('utf-8'))
except UnicodeDecodeError:
return types.Part.from_text(text=data.decode('utf-8', errors='replace'))
size_kb = len(data) / 1024
return types.Part.from_text(
text=(
f'[Binary artifact: {artifact_name}, '
f'type: {mime_type}, size: {size_kb:.1f} KB. '
'Content cannot be displayed inline.]'
)
)
class LoadArtifactsTool(BaseTool):
"""A tool that loads the artifacts and adds them to the session."""
@@ -108,7 +196,8 @@ web UI)."""),
if llm_request.contents and llm_request.contents[-1].parts:
function_response = llm_request.contents[-1].parts[0].function_response
if function_response and function_response.name == 'load_artifacts':
artifact_names = function_response.response['artifact_names']
response = function_response.response or {}
artifact_names = response.get('artifact_names', [])
for artifact_name in artifact_names:
# Try session-scoped first (default behavior)
artifact = await tool_context.load_artifact(artifact_name)
@@ -122,6 +211,18 @@ web UI)."""),
if artifact is None:
logger.warning('Artifact "%s" not found, skipping', artifact_name)
continue
artifact_part = _as_safe_part_for_llm(artifact, artifact_name)
if artifact_part is not artifact:
mime_type = (
artifact.inline_data.mime_type if artifact.inline_data else None
)
logger.debug(
'Converted artifact "%s" (mime_type=%s) to text Part',
artifact_name,
mime_type,
)
llm_request.contents.append(
types.Content(
role='user',
@@ -129,7 +230,7 @@ web UI)."""),
types.Part.from_text(
text=f'Artifact {artifact_name} is:'
),
artifact,
artifact_part,
],
)
)
@@ -0,0 +1,162 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import base64
from google.adk.models.llm_request import LlmRequest
from google.adk.tools.load_artifacts_tool import _maybe_base64_to_bytes
from google.adk.tools.load_artifacts_tool import load_artifacts_tool
from google.genai import types
from pytest import mark
class _StubToolContext:
"""Minimal ToolContext stub for LoadArtifactsTool tests."""
def __init__(self, artifacts_by_name: dict[str, types.Part]):
self._artifacts_by_name = artifacts_by_name
async def list_artifacts(self) -> list[str]:
return list(self._artifacts_by_name.keys())
async def load_artifact(self, name: str) -> types.Part | None:
return self._artifacts_by_name.get(name)
@mark.asyncio
async def test_load_artifacts_converts_unsupported_mime_to_text():
"""Unsupported inline MIME types are converted to text parts."""
artifact_name = 'test.csv'
csv_bytes = b'col1,col2\n1,2\n'
artifact = types.Part(
inline_data=types.Blob(data=csv_bytes, mime_type='application/csv')
)
tool_context = _StubToolContext({artifact_name: artifact})
llm_request = LlmRequest(
contents=[
types.Content(
role='user',
parts=[
types.Part(
function_response=types.FunctionResponse(
name='load_artifacts',
response={'artifact_names': [artifact_name]},
)
)
],
)
]
)
await load_artifacts_tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
assert llm_request.contents[-1].parts[0].text == (
f'Artifact {artifact_name} is:'
)
artifact_part = llm_request.contents[-1].parts[1]
assert artifact_part.inline_data is None
assert artifact_part.text == csv_bytes.decode('utf-8')
@mark.asyncio
async def test_load_artifacts_converts_base64_unsupported_mime_to_text():
"""Unsupported base64 string data is converted to text parts."""
artifact_name = 'test.csv'
csv_bytes = b'col1,col2\n1,2\n'
csv_base64 = base64.b64encode(csv_bytes).decode('ascii')
artifact = types.Part(
inline_data=types.Blob(data=csv_base64, mime_type='application/csv')
)
tool_context = _StubToolContext({artifact_name: artifact})
llm_request = LlmRequest(
contents=[
types.Content(
role='user',
parts=[
types.Part(
function_response=types.FunctionResponse(
name='load_artifacts',
response={'artifact_names': [artifact_name]},
)
)
],
)
]
)
await load_artifacts_tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
artifact_part = llm_request.contents[-1].parts[1]
assert artifact_part.inline_data is None
assert artifact_part.text == csv_bytes.decode('utf-8')
@mark.asyncio
async def test_load_artifacts_keeps_supported_mime_types():
"""Supported inline MIME types are passed through unchanged."""
artifact_name = 'test.pdf'
artifact = types.Part(
inline_data=types.Blob(data=b'%PDF-1.4', mime_type='application/pdf')
)
tool_context = _StubToolContext({artifact_name: artifact})
llm_request = LlmRequest(
contents=[
types.Content(
role='user',
parts=[
types.Part(
function_response=types.FunctionResponse(
name='load_artifacts',
response={'artifact_names': [artifact_name]},
)
)
],
)
]
)
await load_artifacts_tool.process_llm_request(
tool_context=tool_context, llm_request=llm_request
)
artifact_part = llm_request.contents[-1].parts[1]
assert artifact_part.inline_data is not None
assert artifact_part.inline_data.mime_type == 'application/pdf'
def test_maybe_base64_to_bytes_decodes_standard_base64():
"""Standard base64 encoded strings are decoded correctly."""
original = b'hello world'
encoded = base64.b64encode(original).decode('ascii')
assert _maybe_base64_to_bytes(encoded) == original
def test_maybe_base64_to_bytes_decodes_urlsafe_base64():
"""URL-safe base64 encoded strings are decoded correctly."""
original = b'\xfb\xff\xfe' # bytes that produce +/ in std but -_ in urlsafe
encoded = base64.urlsafe_b64encode(original).decode('ascii')
assert _maybe_base64_to_bytes(encoded) == original
def test_maybe_base64_to_bytes_returns_none_for_invalid():
"""Invalid base64 strings return None."""
# Single character is invalid (base64 requires length % 4 == 0 after padding)
assert _maybe_base64_to_bytes('x') is None