Skip to content

Commit ccbedd8

Browse files
xuanyang15copybara-github
authored andcommitted
feat: Support LoadArtifacts parsing for binary text documents (e.g., DOCX)
Co-authored-by: Xuan Yang <xygoogle@google.com> PiperOrigin-RevId: 948373781
1 parent faa1744 commit ccbedd8

4 files changed

Lines changed: 339 additions & 1 deletion

File tree

src/google/adk/models/google_llm.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -501,6 +501,7 @@ async def wait_5_seconds(tool_context: Any = None) -> Any:
501501
)
502502

503503
async def _preprocess_request(self, llm_request: LlmRequest) -> None:
504+
from ..tools import load_artifacts_tool # pylint: disable=import-outside-toplevel
504505

505506
if self._api_backend == GoogleLLMVariant.GEMINI_API:
506507
# Using API key from Google AI Studio to call model doesn't support labels.
@@ -528,6 +529,24 @@ async def _preprocess_request(self, llm_request: LlmRequest) -> None:
528529
llm_request.config.system_instruction = None
529530
await self._adapt_computer_use_tool(llm_request)
530531

532+
# Sanitize inputs by ensuring unsupported inline types (e.g. DOCX from UI)
533+
# are converted to plain text using load_artifacts_tool._as_safe_part_for_llm.
534+
if llm_request.contents:
535+
for content in llm_request.contents:
536+
if not content.parts:
537+
continue
538+
new_parts = []
539+
for part in content.parts:
540+
if part.inline_data:
541+
# GE inline_data does not preserve filenames, so we pass a dummy
542+
# 'inline-file' name as a placeholder for
543+
# _as_safe_part_for_llm's required artifact_name argument.
544+
part = load_artifacts_tool._as_safe_part_for_llm( # pylint: disable=protected-access
545+
part, 'inline-file'
546+
)
547+
new_parts.append(part)
548+
content.parts = new_parts
549+
531550
def _merge_tracking_headers(self, headers: dict[str, str]) -> dict[str, str]:
532551
"""Merge tracking headers to the given headers."""
533552
return merge_tracking_headers(headers)

src/google/adk/tools/load_artifacts_tool.py

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,14 @@
1616

1717
import base64
1818
import binascii
19+
import io
1920
import json
2021
import logging
22+
import re
23+
import struct
2124
from typing import Any
2225
from typing import TYPE_CHECKING
26+
import zipfile
2327

2428
from google.genai import types
2529
from typing_extensions import override
@@ -96,6 +100,42 @@ def _maybe_base64_to_bytes(data: str) -> bytes | None:
96100
return None
97101

98102

103+
def _try_extract_docx_text(data: bytes) -> str | None:
104+
"""Extracts raw text from a DOCX binary."""
105+
# We use regex instead of standard XML parser to avoid XML bomb vulnerabilities,
106+
# and cap the zip extraction at 10 MB to prevent zip bombs.
107+
try:
108+
with zipfile.ZipFile(io.BytesIO(data)) as docx_zip:
109+
if 'word/document.xml' not in docx_zip.namelist():
110+
return None
111+
with docx_zip.open('word/document.xml') as xml_file:
112+
xml_content = xml_file.read(10 * 1024 * 1024).decode(
113+
'utf-8', errors='ignore'
114+
)
115+
116+
# Find the prefix for the WordprocessingML namespace
117+
# xmlns:w="..." or xmlns:something="..."
118+
ns_match = re.search(
119+
r'xmlns:(\w+)="http://schemas.openxmlformats.org/wordprocessingml/2006/main"',
120+
xml_content,
121+
)
122+
prefix = ns_match.group(1) if ns_match else 'w'
123+
124+
p_tag = f'{prefix}:p'
125+
t_tag = f'{prefix}:t'
126+
127+
paragraphs = []
128+
for p in re.split(rf'<{p_tag}(?:[^>]*)>', xml_content):
129+
texts = re.findall(rf'<{t_tag}(?:[^>]*)>([^<]*)</{t_tag}>', p)
130+
if texts:
131+
paragraphs.append(''.join(texts))
132+
133+
return '\n'.join(paragraphs)
134+
except (zipfile.BadZipFile, KeyError, struct.error) as e:
135+
logger.debug('Failed to parse docx layout: %s', e)
136+
return None
137+
138+
99139
def _as_safe_part_for_llm(
100140
artifact: types.Part, artifact_name: str
101141
) -> types.Part:
@@ -125,7 +165,23 @@ def _as_safe_part_for_llm(
125165
return types.Part.from_text(text=data)
126166
data = decoded
127167

128-
if mime_type.startswith('text/') or mime_type in _TEXT_LIKE_MIME_TYPES:
168+
# Attempt DOCX extraction if file seems to be a docx document.
169+
is_docx = mime_type in (
170+
'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
171+
'application/octet-stream',
172+
) or artifact_name.lower().endswith('.docx')
173+
if is_docx:
174+
extracted_text = _try_extract_docx_text(data)
175+
if extracted_text is not None:
176+
return types.Part.from_text(text=extracted_text)
177+
178+
# Fallback to general text extraction
179+
is_text_like = (
180+
mime_type.startswith('text/')
181+
or mime_type in _TEXT_LIKE_MIME_TYPES
182+
or artifact_name.lower().endswith(('.csv', '.txt', '.json', '.xml'))
183+
)
184+
if is_text_like:
129185
try:
130186
return types.Part.from_text(text=data.decode('utf-8'))
131187
except UnicodeDecodeError:

tests/unittests/models/test_google_llm.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
from google.adk.models.google_llm import Gemini
3030
from google.adk.models.llm_request import LlmRequest
3131
from google.adk.models.llm_response import LlmResponse
32+
from google.adk.tools import load_artifacts_tool
3233
from google.adk.utils._client_labels_utils import _AGENT_ENGINE_TELEMETRY_ENV_VARIABLE_NAME
3334
from google.adk.utils._client_labels_utils import _AGENT_ENGINE_TELEMETRY_TAG
3435
from google.adk.utils._google_client_headers import get_tracking_headers
@@ -1037,6 +1038,42 @@ async def test_preprocess_request_handles_backend_specific_fields(
10371038
assert llm_request_with_files.config.labels == expected_labels
10381039

10391040

1041+
@pytest.mark.asyncio
1042+
async def test_preprocess_request_converts_inline_data_safely():
1043+
"""Tests that _preprocess_request uses _as_safe_part_for_llm to sanitize inline data."""
1044+
with mock.patch.object(
1045+
load_artifacts_tool, "_as_safe_part_for_llm", autospec=True
1046+
) as mock_safe_part:
1047+
# Arrange
1048+
mock_safe_part.return_value = Part.from_text(text="safe_text")
1049+
my_gemini_llm = Gemini(model="gemini-2.5-flash")
1050+
1051+
my_llm_request = LlmRequest(
1052+
model="gemini-2.5-flash",
1053+
contents=[
1054+
Content(
1055+
role="user",
1056+
parts=[
1057+
Part(
1058+
inline_data=types.Blob(
1059+
data=b"some bytes",
1060+
mime_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
1061+
)
1062+
)
1063+
],
1064+
)
1065+
],
1066+
)
1067+
1068+
# Act
1069+
await my_gemini_llm._preprocess_request(my_llm_request) # pylint: disable=protected-access
1070+
1071+
# Assert
1072+
mock_safe_part.assert_called_once()
1073+
assert mock_safe_part.call_args[0][1] == "inline-file"
1074+
assert my_llm_request.contents[0].parts[0].text == "safe_text"
1075+
1076+
10401077
@pytest.mark.asyncio
10411078
async def test_generate_content_async_stream_aggregated_content_regardless_of_finish_reason():
10421079
"""Test that aggregated content is generated regardless of finish_reason."""

0 commit comments

Comments
 (0)