diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index bdfb0ee1e52..a29ed66e63d 100644 --- 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 diff --git a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py index a5f72fc08c3..6cc97cd95e6 100644 --- 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 == {}