Merge pull request #24662 from Sameerlite/litellm_gemini-retrieve-file-url-normalize

feat(gemini): normalize AI Studio file retrieve URL
2026-10-10 03:28:53 +00:00 · 2026-03-27 09:59:46 -07:00 · 2026-03-27 09:59:46 -07:00 · 695304d758
commit 695304d758
parent 210e30f055 b212b340ab
2 changed files with 98 additions and 32 deletions
--- a/litellm/llms/gemini/files/transformation.py
+++ b/litellm/llms/gemini/files/transformation.py
@ -5,6 +5,7 @@ For vertex ai, check out the vertex_ai/files/handler.py file.
 """
 import time
 from typing import Any, List, Literal, Optional
+from urllib.parse import urlparse

 import httpx
 from openai.types.file_deleted import FileDeleted
@ -209,27 +210,58 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
        """
        Get the URL to retrieve a file from Google AI Studio.

-        We expect file_id to be the URI (e.g. https://generativelanguage.googleapis.com/v1beta/files/...)
-        as returned by the upload response.
+        Endpoint:
+        GET https://generativelanguage.googleapis.com/v1beta/{name=files/*}
+
+        The URL should look like:
+        https://generativelanguage.googleapis.com/v1beta/files/{file_id}?key=API_KEY
+
+        We expect file_id to be just the file identifier (e.g., files/abc123 or abc123)
+        as returned by the upload response. (If it's a full URL, extract the file name.)
        """
        api_key = litellm_params.get("api_key") or self.get_api_key()
        if not api_key:
            raise ValueError("api_key is required")

-        if file_id.startswith("http"):
-            url = "{}?key={}".format(file_id, api_key)
-        else:
-            # Fallback for just file name (files/...)
-            api_base = (
-                self.get_api_base(litellm_params.get("api_base"))
-                or "https://generativelanguage.googleapis.com"
-            )
-            api_base = api_base.rstrip("/")
-            url = "{}/v1beta/{}?key={}".format(api_base, file_id, api_key)
+        file_part = self._normalize_gemini_file_id(file_id)
+
+        api_base = (
+            self.get_api_base(litellm_params.get("api_base"))
+            or "https://generativelanguage.googleapis.com"
+        )
+        api_base = api_base.rstrip("/")
+
+        url = f"{api_base}/v1beta/{file_part}?key={api_key}"

        # Return empty params dict - API key is already in URL, no query params needed
        return url, {}

+    def _normalize_gemini_file_id(self, file_id: str) -> str:
+        """
+        Normalize file identifier into `files/{id}` form.
+
+        Supports:
+        - `abc123`
+        - `files/abc123`
+        - `https://generativelanguage.googleapis.com/v1beta/files/abc123`
+        """
+        if file_id.startswith(("http://", "https://")):
+            parsed = urlparse(file_id)
+            path = parsed.path.lstrip("/")
+            files_index = path.find("files/")
+            if files_index != -1:
+                normalized_file_id = path[files_index:]
+            else:
+                normalized_file_id = path
+        else:
+            normalized_file_id = file_id
+
+        normalized_file_id = normalized_file_id.strip("/")
+        if not normalized_file_id.startswith("files/"):
+            normalized_file_id = f"files/{normalized_file_id}"
+
+        return normalized_file_id
+
    def transform_retrieve_file_response(
        self,
        raw_response: httpx.Response,
@ -240,8 +272,9 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
        Transform Gemini's file retrieval response into OpenAI-style FileObject
        """
        try:
+            verbose_logger.debug(f"Retrieve file response: {raw_response.text}")
            response_json = raw_response.json()
-
+            verbose_logger.debug(f"Response JSON: {response_json}")
            # Map Gemini state to OpenAI status
            gemini_state = response_json.get("state", "STATE_UNSPECIFIED")
            # Explicitly type status as the Literal union
--- a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py
+++ b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py
@ -3,10 +3,10 @@ Test Google AI Studio (Gemini) files transformation functionality
 """

 import os
-import pytest
 from unittest.mock import Mock, patch

 import httpx
+import pytest

 from litellm.llms.gemini.files.transformation import GoogleAIStudioFilesHandler
 from litellm.types.llms.openai import OpenAIFileObject
@ -23,7 +23,7 @@ class TestGoogleAIStudioFilesTransformation:
        """
        Test that transform_retrieve_file_request returns empty params dict
        to avoid 'Content-Type' query parameter error
-        
+
        Regression test for: https://github.com/BerriAI/litellm/issues/XXX
        When retrieving a file, the API was incorrectly trying to pass Content-Type
        as a query parameter, which Gemini API rejected.
@ -37,14 +37,19 @@ class TestGoogleAIStudioFilesTransformation:
            litellm_params=litellm_params,
        )

-        # Verify URL is constructed correctly with API key
-        assert "key=test-api-key" in url
-        assert file_id in url
+        # Verify URL is constructed exactly as required:
+        # https://generativelanguage.googleapis.com/v1beta/files/{file_id}?key=API_KEY
+        assert (
+            url
+            == "https://generativelanguage.googleapis.com/v1beta/files/test123?key=test-api-key"
+        )

        # CRITICAL: params should be empty dict, not contain Content-Type or any other params
        # These would be incorrectly interpreted as query parameters
        assert params == {}, f"Expected empty params dict, got: {params}"
-        assert "Content-Type" not in params, "Content-Type should not be in query params"
+        assert (
+            "Content-Type" not in params
+        ), "Content-Type should not be in query params"

    def test_transform_retrieve_file_request_with_file_name_only(self):
        """
@ -59,17 +64,44 @@ class TestGoogleAIStudioFilesTransformation:
            litellm_params=litellm_params,
        )

-        # Verify URL is constructed correctly
-        assert "generativelanguage.googleapis.com" in url
-        assert file_id in url
-        assert "key=test-api-key" in url
+        # Verify URL is constructed exactly as required:
+        # https://generativelanguage.googleapis.com/v1beta/files/{file_id}?key=API_KEY
+        assert (
+            url
+            == "https://generativelanguage.googleapis.com/v1beta/files/test123?key=test-api-key"
+        )

        # CRITICAL: params should be empty dict
        assert params == {}, f"Expected empty params dict, got: {params}"
-        assert "Content-Type" not in params, "Content-Type should not be in query params"
+        assert (
+            "Content-Type" not in params
+        ), "Content-Type should not be in query params"

-    @patch.dict('os.environ', {}, clear=True)
-    @patch('litellm.llms.gemini.common_utils.get_secret_str', return_value=None)
+    def test_transform_retrieve_file_request_with_raw_id_only(self):
+        """
+        Regression guard for the exact retrieval URL format.
+
+        If someone changes the method and stops producing:
+        https://generativelanguage.googleapis.com/v1beta/files/{file_id}?key=API_KEY
+        this test should fail.
+        """
+        file_id = "cctqueckiggb"
+        litellm_params = {"api_key": "test-api-key"}
+
+        url, params = self.handler.transform_retrieve_file_request(
+            file_id=file_id,
+            optional_params={},
+            litellm_params=litellm_params,
+        )
+
+        assert (
+            url
+            == "https://generativelanguage.googleapis.com/v1beta/files/cctqueckiggb?key=test-api-key"
+        )
+        assert params == {}
+
+    @patch.dict("os.environ", {}, clear=True)
+    @patch("litellm.llms.gemini.common_utils.get_secret_str", return_value=None)
    def test_transform_retrieve_file_request_missing_api_key(self, mock_get_secret):
        """Test that transform_retrieve_file_request raises error when API key is missing"""
        file_id = "files/test123"
@ -178,7 +210,7 @@ class TestGoogleAIStudioFilesTransformation:
    def test_transform_retrieve_file_response_missing_createTime(self):
        """
        Test that transform_retrieve_file_response raises proper error when createTime is missing
-        
+
        This tests the error scenario that occurs when API returns an error response
        without the expected file metadata fields.
        """
@ -221,14 +253,15 @@ class TestGoogleAIStudioFilesTransformation:
        assert "x-goog-api-key" in result_headers
        assert result_headers["x-goog-api-key"] == api_key

-    @patch.dict('os.environ', {}, clear=True)
-    @patch('litellm.llms.gemini.common_utils.get_secret_str', return_value=None)
+    @patch.dict("os.environ", {}, clear=True)
+    @patch("litellm.llms.gemini.common_utils.get_secret_str", return_value=None)
    def test_validate_environment_missing_api_key(self, mock_get_secret):
        """Test that validate_environment raises error when API key is missing"""
        headers = {}

        with pytest.raises(
-            ValueError, match="GEMINI_API_KEY is required for Google AI Studio file operations"
+            ValueError,
+            match="GEMINI_API_KEY is required for Google AI Studio file operations",
        ):
            self.handler.validate_environment(
                headers=headers,
@ -243,7 +276,7 @@ class TestGoogleAIStudioFilesTransformation:
        """Test that get_complete_url constructs proper upload URL"""
        api_base = "https://generativelanguage.googleapis.com"
        api_key = "test-api-key"
-        
+
        url = self.handler.get_complete_url(
            api_base=api_base,
            api_key=api_key,
@ -274,7 +307,7 @@ class TestGoogleAIStudioFilesTransformation:
        # Verify URL extraction
        assert "files/test123" in url
        assert "generativelanguage.googleapis.com" in url
-        
+
        # Params should be empty (API key goes in header via validate_environment)
        assert params == {}