fix(gemini): keep Files API refs out of the async remote media walker

Gemini AI Studio file URIs under generativelanguage.googleapis.com/v1beta/files/
answer 403 when fetched and must pass through as file_data.file_uri, which the
sync transform already did. Give the shared walker a skip_url_prefixes parameter,
pass that prefix from the Gemini body builder, and skip it in the AI Studio
message transform for both image_url and file parts
This commit is contained in:
mateo-berri 2026-09-04 18:11:41 -07:00
parent 199b44a475
commit e355203014
5 changed files with 99 additions and 8 deletions

View file

@ -199,13 +199,18 @@ def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]:
return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one
def _inline_part(part: object, data_urls: Mapping[str, str]) -> object:
remote: Final = _parse_remote_part(part)
data_url: Final = data_urls.get(remote.url) if remote is not None else None
return _inline(remote, data_url) if remote is not None and data_url is not None else part
def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) -> AllMessageValues:
parts: Final = _content_parts(message)
if not parts:
return message
inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks
_inline(remote, data_urls[remote.url]) if (remote := _parse_remote_part(part)) is not None else part
for part in parts
_inline_part(part, data_urls) for part in parts
]
inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message
return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined
@ -213,13 +218,14 @@ def _inline_message(message: AllMessageValues, data_urls: Mapping[str, str]) ->
async def async_inline_remote_media(
messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues]
skip_url_prefixes: tuple[str, ...] = (),
) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues]
remote_urls: Final = tuple(
dict.fromkeys(
remote.url
for message in messages
for part in _content_parts(message)
if (remote := _parse_remote_part(part)) is not None
if (remote := _parse_remote_part(part)) is not None and not remote.url.startswith(skip_url_prefixes)
)
)
if not remote_urls:

View file

@ -12,7 +12,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionFileObject
from litellm.types.llms.vertex_ai import ContentType, PartType
from litellm.utils import supports_reasoning
from ...vertex_ai.gemini.transformation import _gemini_convert_messages_with_history
from ...vertex_ai.gemini.transformation import GEMINI_FILES_API_URI_PREFIX, _gemini_convert_messages_with_history
from ...vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
@ -127,7 +127,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
if element.get("type") == "image_url":
img_element = cast(ChatCompletionImageObject, element) # cast-ok: runtime type tag checked
_image_url, format, detail = _image_url_fields(img_element)
if _image_url and "https://" in _image_url:
if (
_image_url
and "https://" in _image_url
and not _image_url.startswith(GEMINI_FILES_API_URI_PREFIX)
):
image_obj = convert_to_anthropic_image_obj(_image_url, format=format)
converted_image_url = convert_generic_image_chunk_to_openai_image_obj(image_obj)
if detail is not None:
@ -147,7 +151,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
llm_provider="gemini",
)
file_id = _file_field.get("file_id")
if file_id and ("http://" in file_id or "https://" in file_id):
if (
file_id
and ("http://" in file_id or "https://" in file_id)
and not file_id.startswith(GEMINI_FILES_API_URI_PREFIX)
):
# Convert HTTP/HTTPS file URL to base64 data
try:
base64_data = convert_url_to_base64(file_id)

View file

@ -69,6 +69,7 @@ _GCS_METADATA_VERTEX_BASE: Any | None = None
# Shared sync client for GCS JSON API metadata reads so proxy/SSL settings
# from litellm's HTTP stack apply (see Greptile review on PR #27278).
_GCS_METADATA_HTTP_HANDLER: HTTPHandler | None = None
GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/"
_GEMINI_MIME_TYPE_ALIASES: Final[dict[str, str]] = {
"image/jpg": "image/jpeg",
}
@ -557,7 +558,7 @@ def _process_gemini_media(
file_data = FileDataType(mime_type=mime_type, file_uri=image_url)
part: PartType = {"file_data": file_data}
return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata)
elif image_url.startswith("https://generativelanguage.googleapis.com/v1beta/files/"):
elif image_url.startswith(GEMINI_FILES_API_URI_PREFIX):
# Gemini Files API URIs — the file is already uploaded to Google's
# servers; pass the URI through as file_data without fetching it.
# These URLs return 403 when accessed directly, so we must not try
@ -1349,7 +1350,11 @@ async def async_transform_request_body(
vertex_auth_header=vertex_auth_header,
)
inlined_messages: Final = await async_inline_remote_media(messages) if custom_llm_provider == "gemini" else messages
inlined_messages: Final = (
await async_inline_remote_media(messages, skip_url_prefixes=(GEMINI_FILES_API_URI_PREFIX,))
if custom_llm_provider == "gemini"
else messages
)
if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages):
# _transform_request_body may issue a sync httpx.get (up to 5s timeout)

View file

@ -308,6 +308,34 @@ async def test_async_inline_remote_media_inlines_every_remote_part_shape(async_o
assert messages == snapshot
async def test_async_inline_remote_media_leaves_skipped_url_prefixes_untouched(async_only_image_fetch):
skipped_prefix = "https://generativelanguage.googleapis.com/v1beta/files/"
skipped_file = f"{skipped_prefix}{uuid.uuid4().hex}"
skipped_image = f"{skipped_prefix}{uuid.uuid4().hex}"
fetched_image = f"https://img.example/{uuid.uuid4()}.png"
messages = [
{
"role": "user",
"content": [
{"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}},
{"type": "image_url", "image_url": {"url": skipped_image}},
{"type": "image_url", "image_url": fetched_image},
],
}
]
snapshot = copy.deepcopy(messages)
inlined = await async_inline_remote_media(messages, skip_url_prefixes=(skipped_prefix,))
assert inlined[0]["content"] == [
{"type": "file", "file": {"file_id": skipped_file, "format": "application/pdf"}},
{"type": "image_url", "image_url": {"url": skipped_image}},
{"type": "image_url", "image_url": async_only_image_fetch.data_url},
]
assert async_only_image_fetch.fetched == [fetched_image]
assert messages == snapshot
async def test_async_inline_remote_media_leaves_messages_without_remote_parts_alone(async_only_image_fetch):
messages = [
{"role": "user", "content": [{"type": "text", "text": "hi"}]},

View file

@ -1,4 +1,5 @@
import json
import uuid
import httpx
import pytest
@ -380,3 +381,46 @@ async def test_gemini_ai_studio_async_completion_inlines_remote_images_off_the_e
assert async_only_image_fetch.fetched == [image_url]
assert image_url not in captured["body"]
assert async_only_image_fetch.base64_png in captured["body"]
async def test_gemini_ai_studio_async_completion_passes_files_api_uris_through_unfetched(async_only_image_fetch):
files_api_pdf = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}"
files_api_image = f"https://generativelanguage.googleapis.com/v1beta/files/{uuid.uuid4().hex}"
captured = {}
def handle(request):
captured["body"] = json.loads(request.content)
return httpx.Response(
200,
json={
"candidates": [{"content": {"parts": [{"text": "A report"}], "role": "model"}, "finishReason": "STOP"}],
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
},
)
client = AsyncHTTPHandler()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(handle))
response = await litellm.acompletion(
model="gemini/gemini-3.8-flash",
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "Summarize these"},
{"type": "file", "file": {"file_id": files_api_pdf, "format": "application/pdf"}},
{"type": "image_url", "image_url": {"url": files_api_image, "format": "image/png"}},
],
}
],
api_key="fake-gemini-key",
client=client,
)
assert response.choices[0].message.content == "A report"
assert async_only_image_fetch.fetched == []
file_parts = [part["file_data"] for part in captured["body"]["contents"][0]["parts"] if "file_data" in part]
assert file_parts == [
{"mime_type": "application/pdf", "file_uri": files_api_pdf},
{"mime_type": "image/png", "file_uri": files_api_image},
]