Skip to content

Commit dcf9345

Browse files
author
Lin-Nikaido
committed
fix
1 parent 9ed9594 commit dcf9345

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
@@ -59,45 +59,42 @@
5959

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

103100

@@ -484,19 +481,25 @@ async def _preprocess_request(self, llm_request: LlmRequest) -> None:
484481
for part in content.parts:
485482
# Create copies to avoid mutating the original objects
486483
if part.inline_data:
487-
mime_type = (part.inline_data.mime_type or "").lower()
484+
mime_type = (part.inline_data.mime_type or '').lower()
488485
if mime_type not in _SUPPORTED_FILE_CONTENT_MIME_TYPES:
489-
identifier = part.inline_data.display_name or "inline_file"
490-
part.text = (part.text or "") + f'\n[File reference: "{identifier}"]'
486+
identifier = part.inline_data.display_name or 'inline_file'
487+
part.text = (
488+
part.text or ''
489+
) + f'\n[File reference: "{identifier}"]'
491490
part.inline_data = None
492491
part.inline_data = copy.copy(part.inline_data)
493492
_remove_display_name_if_present(part.inline_data)
494493

495494
if part.file_data:
496-
mime_type = (part.file_data.mime_type or "").lower()
497-
identifier = part.file_data.display_name or part.file_data.file_uri
495+
mime_type = (part.file_data.mime_type or '').lower()
496+
identifier = (
497+
part.file_data.display_name or part.file_data.file_uri
498+
)
498499
if mime_type not in _SUPPORTED_FILE_CONTENT_MIME_TYPES:
499-
part.text = (part.text or "") + f'\n[File reference: "{identifier}"]'
500+
part.text = (
501+
part.text or ''
502+
) + f'\n[File reference: "{identifier}"]'
500503
part.file_data = None
501504
part.file_data = copy.copy(part.file_data)
502505
_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
@@ -2143,46 +2143,45 @@ async def __aexit__(self, *args):
21432143

21442144
@pytest.mark.asyncio
21452145
async def test_preprocess_request_unsupported_mime_type(gemini_llm):
2146-
"""Verifies that MS Office files are escaped to a text reference."""
2147-
unsupported_part = types.Part(
2148-
file_data=types.FileData(
2149-
mime_type="application/vnd.ms-excel",
2150-
file_uri="gs://bucket/data.xls",
2151-
display_name="data.xls"
2152-
)
2153-
)
2154-
req = LlmRequest(
2155-
model="gemini-2.0-flash",
2156-
contents=[types.Content(parts=[unsupported_part])]
2157-
)
2146+
"""Verifies that MS Office files are escaped to a text reference."""
2147+
unsupported_part = types.Part(
2148+
file_data=types.FileData(
2149+
mime_type="application/vnd.ms-excel",
2150+
file_uri="gs://bucket/data.xls",
2151+
display_name="data.xls",
2152+
)
2153+
)
2154+
req = LlmRequest(
2155+
model="gemini-2.0-flash",
2156+
contents=[types.Content(parts=[unsupported_part])],
2157+
)
21582158

2159-
await gemini_llm._preprocess_request(req)
2159+
await gemini_llm._preprocess_request(req)
21602160

2161-
processed_part = req.contents[0].parts[0]
2162-
# File_data should be stripped to avoid the 400 error
2163-
assert processed_part.file_data is None
2164-
# Text fallback should be present
2165-
assert '[File reference: "data.xls"]' in processed_part.text
2161+
processed_part = req.contents[0].parts[0]
2162+
# File_data should be stripped to avoid the 400 error
2163+
assert processed_part.file_data is None
2164+
# Text fallback should be present
2165+
assert '[File reference: "data.xls"]' in processed_part.text
21662166

21672167

21682168
@pytest.mark.asyncio
21692169
async def test_preprocess_request_supported_mime_type(gemini_llm):
2170-
"""Verifies that PDF files are passed through without modification."""
2171-
supported_part = types.Part(
2172-
file_data=types.FileData(
2173-
mime_type="application/pdf",
2174-
file_uri="gs://bucket/doc.pdf",
2175-
display_name="doc.pdf"
2176-
)
2177-
)
2178-
req = LlmRequest(
2179-
model="gemini-2.0-flash",
2180-
contents=[types.Content(parts=[supported_part])]
2181-
)
2170+
"""Verifies that PDF files are passed through without modification."""
2171+
supported_part = types.Part(
2172+
file_data=types.FileData(
2173+
mime_type="application/pdf",
2174+
file_uri="gs://bucket/doc.pdf",
2175+
display_name="doc.pdf",
2176+
)
2177+
)
2178+
req = LlmRequest(
2179+
model="gemini-2.0-flash", contents=[types.Content(parts=[supported_part])]
2180+
)
21822181

2183-
await gemini_llm._preprocess_request(req)
2182+
await gemini_llm._preprocess_request(req)
21842183

2185-
processed_part = req.contents[0].parts[0]
2186-
# file_data should still be intact
2187-
assert processed_part.file_data is not None
2188-
assert processed_part.file_data.mime_type == "application/pdf"
2184+
processed_part = req.contents[0].parts[0]
2185+
# file_data should still be intact
2186+
assert processed_part.file_data is not None
2187+
assert processed_part.file_data.mime_type == "application/pdf"

0 commit comments

Comments
 (0)