Skip to content

Commit 6883209

Browse files
author
Lin-Nikaido
committed
fix
1 parent 372a61d commit 6883209

2 files changed

Lines changed: 79 additions & 77 deletions

File tree

src/google/adk/models/google_llm.py

Lines changed: 45 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -60,45 +60,42 @@
6060

6161
_SUPPORTED_FILE_CONTENT_MIME_TYPES = frozenset({
6262
# Images
63-
"image/png",
64-
"image/jpeg",
65-
"image/webp",
66-
"image/heic",
67-
"image/heif",
68-
63+
'image/png',
64+
'image/jpeg',
65+
'image/webp',
66+
'image/heic',
67+
'image/heif',
6968
# Documents & Text
70-
"application/pdf",
71-
"text/plain",
72-
"text/csv",
73-
"text/html",
74-
"text/md",
75-
"text/x-python",
76-
"text/javascript",
77-
69+
'application/pdf',
70+
'text/plain',
71+
'text/csv',
72+
'text/html',
73+
'text/md',
74+
'text/x-python',
75+
'text/javascript',
7876
# Audio
79-
"audio/wav",
80-
"audio/mp3",
81-
"audio/aiff",
82-
"audio/aac",
83-
"audio/ogg",
84-
"audio/flac",
85-
"audio/mpeg",
86-
"audio/mpga",
87-
"audio/m4a",
88-
"audio/pcm",
89-
"audio/webm",
90-
77+
'audio/wav',
78+
'audio/mp3',
79+
'audio/aiff',
80+
'audio/aac',
81+
'audio/ogg',
82+
'audio/flac',
83+
'audio/mpeg',
84+
'audio/mpga',
85+
'audio/m4a',
86+
'audio/pcm',
87+
'audio/webm',
9188
# Video
92-
"video/mp4",
93-
"video/mpeg",
94-
"video/mov",
95-
"video/quicktime",
96-
"video/avi",
97-
"video/x-flv",
98-
"video/mpg",
99-
"video/webm",
100-
"video/wmv",
101-
"video/3gpp"
89+
'video/mp4',
90+
'video/mpeg',
91+
'video/mov',
92+
'video/quicktime',
93+
'video/avi',
94+
'video/x-flv',
95+
'video/mpg',
96+
'video/webm',
97+
'video/wmv',
98+
'video/3gpp',
10299
})
103100

104101

@@ -499,19 +496,25 @@ async def _preprocess_request(self, llm_request: LlmRequest) -> None:
499496
for part in content.parts:
500497
# Create copies to avoid mutating the original objects
501498
if part.inline_data:
502-
mime_type = (part.inline_data.mime_type or "").lower()
499+
mime_type = (part.inline_data.mime_type or '').lower()
503500
if mime_type not in _SUPPORTED_FILE_CONTENT_MIME_TYPES:
504-
identifier = part.inline_data.display_name or "inline_file"
505-
part.text = (part.text or "") + f'\n[File reference: "{identifier}"]'
501+
identifier = part.inline_data.display_name or 'inline_file'
502+
part.text = (
503+
part.text or ''
504+
) + f'\n[File reference: "{identifier}"]'
506505
part.inline_data = None
507506
part.inline_data = copy.copy(part.inline_data)
508507
_remove_display_name_if_present(part.inline_data)
509508

510509
if part.file_data:
511-
mime_type = (part.file_data.mime_type or "").lower()
512-
identifier = part.file_data.display_name or part.file_data.file_uri
510+
mime_type = (part.file_data.mime_type or '').lower()
511+
identifier = (
512+
part.file_data.display_name or part.file_data.file_uri
513+
)
513514
if mime_type not in _SUPPORTED_FILE_CONTENT_MIME_TYPES:
514-
part.text = (part.text or "") + f'\n[File reference: "{identifier}"]'
515+
part.text = (
516+
part.text or ''
517+
) + f'\n[File reference: "{identifier}"]'
515518
part.file_data = None
516519
part.file_data = copy.copy(part.file_data)
517520
_remove_display_name_if_present(part.file_data)

tests/unittests/models/test_google_llm.py

Lines changed: 34 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -2171,46 +2171,45 @@ async def __aexit__(self, *args):
21712171

21722172
@pytest.mark.asyncio
21732173
async def test_preprocess_request_unsupported_mime_type(gemini_llm):
2174-
"""Verifies that MS Office files are escaped to a text reference."""
2175-
unsupported_part = types.Part(
2176-
file_data=types.FileData(
2177-
mime_type="application/vnd.ms-excel",
2178-
file_uri="gs://bucket/data.xls",
2179-
display_name="data.xls"
2180-
)
2181-
)
2182-
req = LlmRequest(
2183-
model="gemini-2.0-flash",
2184-
contents=[types.Content(parts=[unsupported_part])]
2185-
)
2174+
"""Verifies that MS Office files are escaped to a text reference."""
2175+
unsupported_part = types.Part(
2176+
file_data=types.FileData(
2177+
mime_type="application/vnd.ms-excel",
2178+
file_uri="gs://bucket/data.xls",
2179+
display_name="data.xls",
2180+
)
2181+
)
2182+
req = LlmRequest(
2183+
model="gemini-2.0-flash",
2184+
contents=[types.Content(parts=[unsupported_part])],
2185+
)
21862186

2187-
await gemini_llm._preprocess_request(req)
2187+
await gemini_llm._preprocess_request(req)
21882188

2189-
processed_part = req.contents[0].parts[0]
2190-
# File_data should be stripped to avoid the 400 error
2191-
assert processed_part.file_data is None
2192-
# Text fallback should be present
2193-
assert '[File reference: "data.xls"]' in processed_part.text
2189+
processed_part = req.contents[0].parts[0]
2190+
# File_data should be stripped to avoid the 400 error
2191+
assert processed_part.file_data is None
2192+
# Text fallback should be present
2193+
assert '[File reference: "data.xls"]' in processed_part.text
21942194

21952195

21962196
@pytest.mark.asyncio
21972197
async def test_preprocess_request_supported_mime_type(gemini_llm):
2198-
"""Verifies that PDF files are passed through without modification."""
2199-
supported_part = types.Part(
2200-
file_data=types.FileData(
2201-
mime_type="application/pdf",
2202-
file_uri="gs://bucket/doc.pdf",
2203-
display_name="doc.pdf"
2204-
)
2205-
)
2206-
req = LlmRequest(
2207-
model="gemini-2.0-flash",
2208-
contents=[types.Content(parts=[supported_part])]
2209-
)
2198+
"""Verifies that PDF files are passed through without modification."""
2199+
supported_part = types.Part(
2200+
file_data=types.FileData(
2201+
mime_type="application/pdf",
2202+
file_uri="gs://bucket/doc.pdf",
2203+
display_name="doc.pdf",
2204+
)
2205+
)
2206+
req = LlmRequest(
2207+
model="gemini-2.0-flash", contents=[types.Content(parts=[supported_part])]
2208+
)
22102209

2211-
await gemini_llm._preprocess_request(req)
2210+
await gemini_llm._preprocess_request(req)
22122211

2213-
processed_part = req.contents[0].parts[0]
2214-
# file_data should still be intact
2215-
assert processed_part.file_data is not None
2216-
assert processed_part.file_data.mime_type == "application/pdf"
2212+
processed_part = req.contents[0].parts[0]
2213+
# file_data should still be intact
2214+
assert processed_part.file_data is not None
2215+
assert processed_part.file_data.mime_type == "application/pdf"

0 commit comments

Comments
 (0)