fix: Add support for file URIs in LiteLLM content conversion to fix issue #3131

changed the LiteLLM content conversion so Part.file_data.file_uri (like the gs://…) becomes a file object with file_id, making sure GCS-backed files reach LiteLLM proxies instead of being dropped add unit tests covering both _get_content and _content_to_message_param paths for file URIs

PiperOrigin-RevId: 817658432
This commit is contained in:
George Weale
2025-10-10 08:38:47 -07:00
committed by Copybara-Service
parent 998264a5b1
commit 85ed500871
2 changed files with 66 additions and 1 deletions
+12 -1
View File
@@ -65,8 +65,9 @@ _NEW_LINE = "\n"
_EXCLUDED_PART_FIELD = {"inline_data": {"data"}}
class ChatCompletionFileUrlObject(TypedDict):
class ChatCompletionFileUrlObject(TypedDict, total=False):
file_data: str
file_id: str
format: str
@@ -281,6 +282,16 @@ def _get_content(
})
else:
raise ValueError("LiteLlm(BaseLlm) does not support this content part.")
elif part.file_data and part.file_data.file_uri:
file_object: ChatCompletionFileUrlObject = {
"file_id": part.file_data.file_uri,
}
if part.file_data.mime_type:
file_object["format"] = part.file_data.mime_type
content_objects.append({
"type": "file",
"file": file_object,
})
return content_objects
+54
View File
@@ -1005,6 +1005,47 @@ def test_content_to_message_param_user_message():
assert message["content"] == "Test prompt"
def test_content_to_message_param_user_message_with_file_uri():
file_part = types.Part.from_uri(
file_uri="gs://bucket/document.pdf", mime_type="application/pdf"
)
content = types.Content(
role="user",
parts=[
types.Part.from_text(text="Summarize this file."),
file_part,
],
)
message = _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"] == "gs://bucket/document.pdf"
assert message["content"][1]["file"]["format"] == "application/pdf"
def test_content_to_message_param_user_message_file_uri_only():
file_part = types.Part.from_uri(
file_uri="gs://bucket/only.pdf", mime_type="application/pdf"
)
content = types.Content(
role="user",
parts=[
file_part,
],
)
message = _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"] == "gs://bucket/only.pdf"
assert message["content"][0]["file"]["format"] == "application/pdf"
def test_content_to_message_param_multi_part_function_response():
part1 = types.Part.from_function_response(
name="function_one",
@@ -1183,6 +1224,19 @@ def test_get_content_pdf():
assert content[0]["file"]["format"] == "application/pdf"
def test_get_content_file_uri():
parts = [
types.Part.from_uri(
file_uri="gs://bucket/document.pdf",
mime_type="application/pdf",
)
]
content = _get_content(parts)
assert content[0]["type"] == "file"
assert content[0]["file"]["file_id"] == "gs://bucket/document.pdf"
assert content[0]["file"]["format"] == "application/pdf"
def test_get_content_audio():
parts = [
types.Part.from_bytes(data=b"test_audio_data", mime_type="audio/mpeg")