mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: Add MIME type inference and default for file URIs in LiteLLM
This change ensures that file URI parts passed to LiteLLM always include a "format" field. If `mime_type` is not explicitly provided in `FileData`, the system attempts to infer it from the URI's file extension. If inference fails, a default "application/octet-stream" is used. This is necessary because LiteLLM's Vertex AI backend requires the "format" field for GCS URIs. Close #3787 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 843753810
This commit is contained in:
George Weale
committed by
Copybara-Service
Copybara-Service
parent
8782a69503
commit
5c4bae7ff2
@@ -18,6 +18,7 @@ import base64
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import mimetypes
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
@@ -33,6 +34,7 @@ from typing import Optional
|
||||
from typing import Tuple
|
||||
from typing import TypedDict
|
||||
from typing import Union
|
||||
from urllib.parse import urlparse
|
||||
import uuid
|
||||
import warnings
|
||||
|
||||
@@ -129,6 +131,51 @@ def _get_provider_from_model(model: str) -> str:
|
||||
return ""
|
||||
|
||||
|
||||
# Default MIME type when none can be inferred
|
||||
_DEFAULT_MIME_TYPE = "application/octet-stream"
|
||||
|
||||
|
||||
def _infer_mime_type_from_uri(uri: str) -> Optional[str]:
|
||||
"""Attempts to infer MIME type from a URI's path extension.
|
||||
|
||||
Args:
|
||||
uri: A URI string (e.g., 'gs://bucket/file.pdf' or
|
||||
'https://example.com/doc.json')
|
||||
|
||||
Returns:
|
||||
The inferred MIME type, or None if it cannot be determined.
|
||||
"""
|
||||
try:
|
||||
parsed = urlparse(uri)
|
||||
# Get the path component and extract filename
|
||||
path = parsed.path
|
||||
if not path:
|
||||
return None
|
||||
|
||||
# Many artifact URIs are versioned (for example, ".../filename/0" or
|
||||
# ".../filename/versions/0"). If the last path segment looks like a numeric
|
||||
# version, infer from the preceding filename instead.
|
||||
segments = [segment for segment in path.split("/") if segment]
|
||||
if not segments:
|
||||
return None
|
||||
|
||||
candidate = segments[-1]
|
||||
if candidate.isdigit():
|
||||
segments = segments[:-1]
|
||||
if segments and segments[-1].lower() in ("versions", "version"):
|
||||
segments = segments[:-1]
|
||||
|
||||
if not segments:
|
||||
return None
|
||||
|
||||
candidate = segments[-1]
|
||||
mime_type, _ = mimetypes.guess_type(candidate)
|
||||
return mime_type
|
||||
except (ValueError, AttributeError) as e:
|
||||
logger.debug("Could not infer MIME type from URI %s: %s", uri, e)
|
||||
return None
|
||||
|
||||
|
||||
def _decode_inline_text_data(raw_bytes: bytes) -> str:
|
||||
"""Decodes inline file bytes that represent textual content."""
|
||||
try:
|
||||
@@ -553,6 +600,22 @@ async def _get_content(
|
||||
file_object: ChatCompletionFileUrlObject = {
|
||||
"file_id": part.file_data.file_uri,
|
||||
}
|
||||
# Determine MIME type: use explicit value, infer from URI, or use default
|
||||
mime_type = part.file_data.mime_type
|
||||
if not mime_type:
|
||||
mime_type = _infer_mime_type_from_uri(part.file_data.file_uri)
|
||||
if not mime_type and part.file_data.display_name:
|
||||
guessed_mime_type, _ = mimetypes.guess_type(part.file_data.display_name)
|
||||
mime_type = guessed_mime_type
|
||||
if not mime_type:
|
||||
# LiteLLM's Vertex AI backend requires format for GCS URIs
|
||||
mime_type = _DEFAULT_MIME_TYPE
|
||||
logger.debug(
|
||||
"Could not determine MIME type for file_uri %s, using default: %s",
|
||||
part.file_data.file_uri,
|
||||
mime_type,
|
||||
)
|
||||
file_object["format"] = mime_type
|
||||
content_objects.append({
|
||||
"type": "file",
|
||||
"file": file_object,
|
||||
|
||||
@@ -1574,13 +1574,13 @@ async def test_content_to_message_param_user_message_with_file_uri(
|
||||
)
|
||||
|
||||
message = await _content_to_message_param(content)
|
||||
assert message["role"] == "user"
|
||||
assert isinstance(message["content"], list)
|
||||
assert message["content"][0]["type"] == "text"
|
||||
assert message["content"][0]["text"] == "Summarize this file."
|
||||
assert message["content"][1]["type"] == "file"
|
||||
assert message["content"][1]["file"]["file_id"] == file_uri
|
||||
assert "format" not in message["content"][1]["file"]
|
||||
assert message == {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Summarize this file."},
|
||||
{"type": "file", "file": {"file_id": file_uri, "format": mime_type}},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -1597,11 +1597,88 @@ async def test_content_to_message_param_user_message_file_uri_only(
|
||||
)
|
||||
|
||||
message = await _content_to_message_param(content)
|
||||
assert message["role"] == "user"
|
||||
assert isinstance(message["content"], list)
|
||||
assert message["content"][0]["type"] == "file"
|
||||
assert message["content"][0]["file"]["file_id"] == file_uri
|
||||
assert "format" not in message["content"][0]["file"]
|
||||
assert message == {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "file", "file": {"file_id": file_uri, "format": mime_type}},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_content_to_message_param_user_message_file_uri_without_mime_type():
|
||||
"""Test handling of file_data without mime_type (GcsArtifactService scenario).
|
||||
|
||||
When using GcsArtifactService, artifacts may have file_uri (gs://...) but
|
||||
without mime_type set. LiteLLM's Vertex AI backend requires the format
|
||||
field to be present, so we infer MIME type from the URI extension or use
|
||||
a default fallback to ensure compatibility.
|
||||
|
||||
See: https://github.com/google/adk-python/issues/3787
|
||||
"""
|
||||
file_part = types.Part(
|
||||
file_data=types.FileData(
|
||||
file_uri="gs://agent-artifact-bucket/app/user/session/artifact/0"
|
||||
)
|
||||
)
|
||||
content = types.Content(
|
||||
role="user",
|
||||
parts=[
|
||||
types.Part.from_text(text="Analyze this file."),
|
||||
file_part,
|
||||
],
|
||||
)
|
||||
|
||||
message = await _content_to_message_param(content)
|
||||
assert message == {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Analyze this file."},
|
||||
{
|
||||
"type": "file",
|
||||
"file": {
|
||||
"file_id": (
|
||||
"gs://agent-artifact-bucket/app/user/session/artifact/0"
|
||||
),
|
||||
"format": "application/octet-stream",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_content_to_message_param_user_message_file_uri_infer_mime_type():
|
||||
"""Test MIME type inference from file_uri extension.
|
||||
|
||||
When file_data has a file_uri with a recognizable extension but no explicit
|
||||
mime_type, the MIME type should be inferred from the extension.
|
||||
|
||||
See: https://github.com/google/adk-python/issues/3787
|
||||
"""
|
||||
file_part = types.Part(
|
||||
file_data=types.FileData(
|
||||
file_uri="gs://bucket/path/to/document.pdf",
|
||||
)
|
||||
)
|
||||
content = types.Content(
|
||||
role="user",
|
||||
parts=[file_part],
|
||||
)
|
||||
|
||||
message = await _content_to_message_param(content)
|
||||
assert message == {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "file",
|
||||
"file": {
|
||||
"file_id": "gs://bucket/path/to/document.pdf",
|
||||
"format": "application/pdf",
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -1995,9 +2072,112 @@ async def test_get_content_file_bytes(file_data, mime_type, expected_base64):
|
||||
async def test_get_content_file_uri(file_uri, mime_type):
|
||||
parts = [types.Part.from_uri(file_uri=file_uri, mime_type=mime_type)]
|
||||
content = await _get_content(parts)
|
||||
assert content[0]["type"] == "file"
|
||||
assert content[0]["file"]["file_id"] == file_uri
|
||||
assert "format" not in content[0]["file"]
|
||||
assert content[0] == {
|
||||
"type": "file",
|
||||
"file": {"file_id": file_uri, "format": mime_type},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_content_file_uri_infer_mime_type():
|
||||
"""Test MIME type inference from file_uri extension.
|
||||
|
||||
When file_data has a file_uri with a recognizable extension but no explicit
|
||||
mime_type, the MIME type should be inferred from the extension.
|
||||
|
||||
See: https://github.com/google/adk-python/issues/3787
|
||||
"""
|
||||
# Use Part constructor directly to test MIME type inference in _get_content
|
||||
# (types.Part.from_uri does its own inference, so we bypass it)
|
||||
parts = [
|
||||
types.Part(
|
||||
file_data=types.FileData(file_uri="gs://bucket/path/to/document.pdf")
|
||||
)
|
||||
]
|
||||
content = await _get_content(parts)
|
||||
assert content[0] == {
|
||||
"type": "file",
|
||||
"file": {
|
||||
"file_id": "gs://bucket/path/to/document.pdf",
|
||||
"format": "application/pdf",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_content_file_uri_versioned_infer_mime_type():
|
||||
"""Test MIME type inference from versioned artifact URIs."""
|
||||
parts = [
|
||||
types.Part(
|
||||
file_data=types.FileData(
|
||||
file_uri="gs://bucket/path/to/document.pdf/0"
|
||||
)
|
||||
)
|
||||
]
|
||||
content = await _get_content(parts)
|
||||
assert content[0]["file"]["format"] == "application/pdf"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_content_file_uri_infers_from_display_name():
|
||||
"""Test MIME type inference from display_name when URI lacks extension."""
|
||||
parts = [
|
||||
types.Part(
|
||||
file_data=types.FileData(
|
||||
file_uri="gs://bucket/artifact/0",
|
||||
display_name="document.pdf",
|
||||
)
|
||||
)
|
||||
]
|
||||
content = await _get_content(parts)
|
||||
assert content[0]["file"]["format"] == "application/pdf"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_content_file_uri_default_mime_type():
|
||||
"""Test that file_uri without extension uses default MIME type.
|
||||
|
||||
When file_data has a file_uri without a recognizable extension and no explicit
|
||||
mime_type, a default MIME type should be used to ensure compatibility with
|
||||
LiteLLM backends.
|
||||
|
||||
See: https://github.com/google/adk-python/issues/3787
|
||||
"""
|
||||
# Use Part constructor directly to create file_data without mime_type
|
||||
# (types.Part.from_uri requires a valid mime_type when it can't infer)
|
||||
parts = [
|
||||
types.Part(file_data=types.FileData(file_uri="gs://bucket/artifact/0"))
|
||||
]
|
||||
content = await _get_content(parts)
|
||||
assert content[0] == {
|
||||
"type": "file",
|
||||
"file": {
|
||||
"file_id": "gs://bucket/artifact/0",
|
||||
"format": "application/octet-stream",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"uri,expected_mime_type",
|
||||
[
|
||||
("gs://bucket/file.pdf", "application/pdf"),
|
||||
("gs://bucket/path/to/document.json", "application/json"),
|
||||
("gs://bucket/image.png", "image/png"),
|
||||
("gs://bucket/image.jpg", "image/jpeg"),
|
||||
("gs://bucket/audio.mp3", "audio/mpeg"),
|
||||
("gs://bucket/video.mp4", "video/mp4"),
|
||||
],
|
||||
)
|
||||
async def test_get_content_file_uri_mime_type_inference(
|
||||
uri, expected_mime_type
|
||||
):
|
||||
"""Test MIME type inference from various file extensions."""
|
||||
# Use Part constructor directly to test MIME type inference in _get_content
|
||||
parts = [types.Part(file_data=types.FileData(file_uri=uri))]
|
||||
content = await _get_content(parts)
|
||||
assert content[0]["file"]["format"] == expected_mime_type
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
Reference in New Issue
Block a user