Skip to content

Commit 43b484f

Browse files
GWealecopybara-github
authored andcommitted
fix: Handle file URI conversion for LiteLLM based on provider and model
This change updates how `file_data.file_uri` parts are converted to LiteLLM content. For providers like OpenAI and Azure, only URIs resembling OpenAI file IDs ("file-...") are passed as file objects. Other URIs are converted to a text placeholder Close #4038 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 855277306
1 parent 94d48fc commit 43b484f

File tree

2 files changed

+195
-3
lines changed

2 files changed

+195
-3
lines changed

src/google/adk/models/lite_llm.py

Lines changed: 75 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,45 @@ def _infer_mime_type_from_uri(uri: str) -> Optional[str]:
181181
return None
182182

183183

184+
def _looks_like_openai_file_id(file_uri: str) -> bool:
185+
"""Returns True when file_uri resembles an OpenAI/Azure file id."""
186+
return file_uri.startswith("file-")
187+
188+
189+
def _redact_file_uri_for_log(
190+
file_uri: str, *, display_name: str | None = None
191+
) -> str:
192+
"""Returns a privacy-preserving identifier for logs."""
193+
if display_name:
194+
return display_name
195+
if _looks_like_openai_file_id(file_uri):
196+
return "file-<redacted>"
197+
try:
198+
parsed = urlparse(file_uri)
199+
except ValueError:
200+
return "<unparseable>"
201+
if not parsed.scheme:
202+
return "<unknown>"
203+
segments = [segment for segment in parsed.path.split("/") if segment]
204+
tail = segments[-1] if segments else ""
205+
if tail:
206+
return f"{parsed.scheme}://<redacted>/{tail}"
207+
return f"{parsed.scheme}://<redacted>"
208+
209+
210+
def _requires_file_uri_fallback(
211+
provider: str, model: str, file_uri: str
212+
) -> bool:
213+
"""Returns True when `file_uri` should not be sent as a file content block."""
214+
if provider in _FILE_ID_REQUIRED_PROVIDERS:
215+
return not _looks_like_openai_file_id(file_uri)
216+
if provider == "anthropic":
217+
return True
218+
if provider == "vertex_ai" and not _is_litellm_gemini_model(model):
219+
return True
220+
return False
221+
222+
184223
def _decode_inline_text_data(raw_bytes: bytes) -> str:
185224
"""Decodes inline file bytes that represent textual content."""
186225
try:
@@ -447,6 +486,7 @@ async def _content_to_message_param(
447486
content: types.Content,
448487
*,
449488
provider: str = "",
489+
model: str = "",
450490
) -> Union[Message, list[Message]]:
451491
"""Converts a types.Content to a litellm Message or list of Messages.
452492
@@ -456,6 +496,7 @@ async def _content_to_message_param(
456496
Args:
457497
content: The content to convert.
458498
provider: The LLM provider name (e.g., "openai", "azure").
499+
model: The LiteLLM model string, used for provider-specific behavior.
459500
460501
Returns:
461502
A litellm Message, a list of litellm Messages.
@@ -499,7 +540,9 @@ async def _content_to_message_param(
499540

500541
if role == "user":
501542
user_parts = [part for part in content.parts if not part.thought]
502-
message_content = await _get_content(user_parts, provider=provider) or None
543+
message_content = (
544+
await _get_content(user_parts, provider=provider, model=model) or None
545+
)
503546
return ChatCompletionUserMessage(role="user", content=message_content)
504547
else: # assistant/model
505548
tool_calls = []
@@ -523,7 +566,7 @@ async def _content_to_message_param(
523566
content_parts.append(part)
524567

525568
final_content = (
526-
await _get_content(content_parts, provider=provider)
569+
await _get_content(content_parts, provider=provider, model=model)
527570
if content_parts
528571
else None
529572
)
@@ -620,6 +663,7 @@ async def _get_content(
620663
parts: Iterable[types.Part],
621664
*,
622665
provider: str = "",
666+
model: str = "",
623667
) -> OpenAIMessageContent:
624668
"""Converts a list of parts to litellm content.
625669
@@ -629,6 +673,8 @@ async def _get_content(
629673
Args:
630674
parts: The parts to convert.
631675
provider: The LLM provider name (e.g., "openai", "azure").
676+
model: The LiteLLM model string (e.g., "openai/gpt-4o",
677+
"vertex_ai/gemini-2.5-flash").
632678
633679
Returns:
634680
The litellm content.
@@ -709,6 +755,32 @@ async def _get_content(
709755
f"{part.inline_data.mime_type}."
710756
)
711757
elif part.file_data and part.file_data.file_uri:
758+
if (
759+
provider in _FILE_ID_REQUIRED_PROVIDERS
760+
and _looks_like_openai_file_id(part.file_data.file_uri)
761+
):
762+
content_objects.append({
763+
"type": "file",
764+
"file": {"file_id": part.file_data.file_uri},
765+
})
766+
continue
767+
768+
if _requires_file_uri_fallback(provider, model, part.file_data.file_uri):
769+
logger.debug(
770+
"File URI %s not supported for provider %s, using text fallback",
771+
_redact_file_uri_for_log(
772+
part.file_data.file_uri,
773+
display_name=part.file_data.display_name,
774+
),
775+
provider,
776+
)
777+
identifier = part.file_data.display_name or part.file_data.file_uri
778+
content_objects.append({
779+
"type": "text",
780+
"text": f'[File reference: "{identifier}"]',
781+
})
782+
continue
783+
712784
file_object: ChatCompletionFileUrlObject = {
713785
"file_id": part.file_data.file_uri,
714786
}
@@ -1363,7 +1435,7 @@ async def _get_completion_inputs(
13631435
messages: List[Message] = []
13641436
for content in llm_request.contents or []:
13651437
message_param_or_list = await _content_to_message_param(
1366-
content, provider=provider
1438+
content, provider=provider, model=model
13671439
)
13681440
if isinstance(message_param_or_list, list):
13691441
messages.extend(message_param_or_list)

tests/unittests/models/test_litellm.py

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2304,6 +2304,126 @@ async def test_get_content_file_uri(file_uri, mime_type):
23042304
}
23052305

23062306

2307+
@pytest.mark.asyncio
2308+
@pytest.mark.parametrize(
2309+
"provider,model",
2310+
[
2311+
("openai", "openai/gpt-4o"),
2312+
("azure", "azure/gpt-4"),
2313+
],
2314+
)
2315+
async def test_get_content_file_uri_file_id_required_falls_back_to_text(
2316+
provider, model
2317+
):
2318+
parts = [
2319+
types.Part(
2320+
file_data=types.FileData(
2321+
file_uri="gs://bucket/path/to/document.pdf",
2322+
mime_type="application/pdf",
2323+
display_name="document.pdf",
2324+
)
2325+
)
2326+
]
2327+
content = await _get_content(parts, provider=provider, model=model)
2328+
assert content == [
2329+
{"type": "text", "text": '[File reference: "document.pdf"]'}
2330+
]
2331+
2332+
2333+
@pytest.mark.asyncio
2334+
@pytest.mark.parametrize(
2335+
"provider,model",
2336+
[
2337+
("openai", "openai/gpt-4o"),
2338+
("azure", "azure/gpt-4"),
2339+
],
2340+
)
2341+
async def test_get_content_file_uri_file_id_required_preserves_file_id(
2342+
provider, model
2343+
):
2344+
parts = [
2345+
types.Part(
2346+
file_data=types.FileData(
2347+
file_uri="file-abc123",
2348+
mime_type="application/pdf",
2349+
)
2350+
)
2351+
]
2352+
content = await _get_content(parts, provider=provider, model=model)
2353+
assert content == [{"type": "file", "file": {"file_id": "file-abc123"}}]
2354+
2355+
2356+
@pytest.mark.asyncio
2357+
async def test_get_content_file_uri_anthropic_falls_back_to_text():
2358+
parts = [
2359+
types.Part(
2360+
file_data=types.FileData(
2361+
file_uri="gs://bucket/path/to/document.pdf",
2362+
mime_type="application/pdf",
2363+
display_name="document.pdf",
2364+
)
2365+
)
2366+
]
2367+
content = await _get_content(
2368+
parts, provider="anthropic", model="anthropic/claude-3-5"
2369+
)
2370+
assert content == [
2371+
{"type": "text", "text": '[File reference: "document.pdf"]'}
2372+
]
2373+
2374+
2375+
@pytest.mark.asyncio
2376+
async def test_get_content_file_uri_anthropic_openai_file_id_falls_back_to_text():
2377+
parts = [types.Part(file_data=types.FileData(file_uri="file-abc123"))]
2378+
content = await _get_content(
2379+
parts, provider="anthropic", model="anthropic/claude-3-5"
2380+
)
2381+
assert content == [
2382+
{"type": "text", "text": '[File reference: "file-abc123"]'}
2383+
]
2384+
2385+
2386+
@pytest.mark.asyncio
2387+
async def test_get_content_file_uri_vertex_ai_non_gemini_falls_back_to_text():
2388+
parts = [
2389+
types.Part(
2390+
file_data=types.FileData(
2391+
file_uri="gs://bucket/path/to/document.pdf",
2392+
mime_type="application/pdf",
2393+
display_name="document.pdf",
2394+
)
2395+
)
2396+
]
2397+
content = await _get_content(
2398+
parts, provider="vertex_ai", model="vertex_ai/claude-3-5"
2399+
)
2400+
assert content == [
2401+
{"type": "text", "text": '[File reference: "document.pdf"]'}
2402+
]
2403+
2404+
2405+
@pytest.mark.asyncio
2406+
async def test_get_content_file_uri_vertex_ai_gemini_keeps_file_block():
2407+
parts = [
2408+
types.Part(
2409+
file_data=types.FileData(
2410+
file_uri="gs://bucket/path/to/document.pdf",
2411+
mime_type="application/pdf",
2412+
)
2413+
)
2414+
]
2415+
content = await _get_content(
2416+
parts, provider="vertex_ai", model="vertex_ai/gemini-2.5-flash"
2417+
)
2418+
assert content == [{
2419+
"type": "file",
2420+
"file": {
2421+
"file_id": "gs://bucket/path/to/document.pdf",
2422+
"format": "application/pdf",
2423+
},
2424+
}]
2425+
2426+
23072427
@pytest.mark.asyncio
23082428
async def test_get_content_file_uri_infer_mime_type():
23092429
"""Test MIME type inference from file_uri extension.

0 commit comments

Comments
 (0)