Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions litellm/llms/gemini/files/transformation.py
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,7 @@ def transform_retrieve_file_request(
We expect file_id to be the URI (e.g. https://generativelanguage.googleapis.com/v1beta/files/...)
as returned by the upload response.
"""
api_key = litellm_params.get("api_key")
api_key = litellm_params.get("api_key") or self.get_api_key()
if not api_key:
raise ValueError("api_key is required")

Expand All @@ -222,7 +222,8 @@ def transform_retrieve_file_request(
api_base = api_base.rstrip("/")
url = "{}/v1beta/{}?key={}".format(api_base, file_id, api_key)

return url, {"Content-Type": "application/json"}
# Return empty params dict - API key is already in URL, no query params needed
return url, {}

def transform_retrieve_file_response(
self,
Expand Down Expand Up @@ -299,7 +300,7 @@ def transform_delete_file_request(
# Extract the file path from full URI
file_name = file_id.split("/v1beta/")[-1]
else:
file_name = file_id
file_name = file_id if file_id.startswith("files/") else f"files/{file_id}"

# Construct the delete URL
url = f"{api_base}/v1beta/{file_name}"
Expand Down
1 change: 1 addition & 0 deletions tests/test_litellm/llms/gemini/files/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
"""Tests for Gemini files functionality"""
Original file line number Diff line number Diff line change
@@ -0,0 +1,298 @@
"""
Test Google AI Studio (Gemini) files transformation functionality
"""

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

import httpx

from litellm.llms.gemini.files.transformation import GoogleAIStudioFilesHandler
from litellm.types.llms.openai import OpenAIFileObject


class TestGoogleAIStudioFilesTransformation:
"""Test Google AI Studio files transformation"""

def setup_method(self):
"""Setup test method"""
self.handler = GoogleAIStudioFilesHandler()

def test_transform_retrieve_file_request_with_full_uri(self):
"""
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.
"""
file_id = "https://generativelanguage.googleapis.com/v1beta/files/test123"
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,
)

# Verify URL is constructed correctly with API key
assert "key=test-api-key" in url
assert file_id in url

# 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"

def test_transform_retrieve_file_request_with_file_name_only(self):
"""
Test that transform_retrieve_file_request handles file_id without full URI
"""
file_id = "files/test123"
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,
)

# Verify URL is constructed correctly
assert "generativelanguage.googleapis.com" in url

Check failure

Code scanning / CodeQL

Incomplete URL substring sanitization High test

The string
generativelanguage.googleapis.com
may be at an arbitrary position in the sanitized URL.

Copilot Autofix

AI 8 months ago

In general, the fix is to avoid treating the URL as an opaque string when validating its host or structure. Instead, the URL should be parsed using a standard URL parser (for example, urllib.parse.urlparse in Python), and assertions or checks should be made against structured components like scheme, netloc, and hostname. This prevents accidentally accepting URLs that merely contain an allowed host as a substring in an unsafe location.

For this specific test, we want to keep the existing intent—verifying that the constructed URL targets the Gemini API endpoint—but replace the substring assertion with a structured check on the parsed hostname. We will import urlparse from urllib.parse at the top of the file, and in test_transform_retrieve_file_request_with_file_name_only, we will replace assert "generativelanguage.googleapis.com" in url with code that parses url and asserts that parsed.hostname == "generativelanguage.googleapis.com". This keeps behavior equivalent (or stricter) while eliminating the substring-host pattern that CodeQL flags. No other test logic or functional behavior needs to change.

Concretely:

  • Add from urllib.parse import urlparse alongside the existing imports in tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py.
  • In test_transform_retrieve_file_request_with_file_name_only, replace the line at/around 63 with parsing url = ... and asserting on hostname. For example:
    • parsed_url = urlparse(url)
    • assert parsed_url.hostname == "generativelanguage.googleapis.com"
      This keeps the rest of the assertions (for file_id and API key presence, and empty params) unchanged.
Suggested changeset 1
tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py

Autofix patch

Autofix patch
Run the following command in your local git repository to apply this patch
cat << 'EOF' | git apply
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
--- a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py
+++ b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py
@@ -7,6 +7,7 @@
 from unittest.mock import Mock, patch
 
 import httpx
+from urllib.parse import urlparse
 
 from litellm.llms.gemini.files.transformation import GoogleAIStudioFilesHandler
 from litellm.types.llms.openai import OpenAIFileObject
@@ -60,7 +61,8 @@
         )
 
         # Verify URL is constructed correctly
-        assert "generativelanguage.googleapis.com" in url
+        parsed_url = urlparse(url)
+        assert parsed_url.hostname == "generativelanguage.googleapis.com"
         assert file_id in url
         assert "key=test-api-key" in url
 
EOF
@@ -7,6 +7,7 @@
from unittest.mock import Mock, patch

import httpx
from urllib.parse import urlparse

from litellm.llms.gemini.files.transformation import GoogleAIStudioFilesHandler
from litellm.types.llms.openai import OpenAIFileObject
@@ -60,7 +61,8 @@
)

# Verify URL is constructed correctly
assert "generativelanguage.googleapis.com" in url
parsed_url = urlparse(url)
assert parsed_url.hostname == "generativelanguage.googleapis.com"
assert file_id in url
assert "key=test-api-key" in url

Copilot is powered by AI and may make mistakes. Always verify output.
assert file_id in url
assert "key=test-api-key" in url

# 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"

@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"
litellm_params = {}

with pytest.raises(ValueError, match="api_key is required"):
self.handler.transform_retrieve_file_request(
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
)

def test_transform_retrieve_file_response_success(self):
"""Test successful transformation of Gemini file retrieval response"""
# Mock response data from Gemini API
mock_response_data = {
"name": "files/test123",
"displayName": "test_file.pdf",
"mimeType": "application/pdf",
"sizeBytes": "1024",
"createTime": "2024-01-15T10:30:00.123456Z",
"updateTime": "2024-01-15T10:30:00.123456Z",
"expirationTime": "2024-01-17T10:30:00.123456Z",
"sha256Hash": "abcd1234",
"uri": "https://generativelanguage.googleapis.com/v1beta/files/test123",
"state": "ACTIVE",
}

# Create mock httpx response
mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = mock_response_data

# Create mock logging object
mock_logging_obj = Mock()

# Transform response
result = self.handler.transform_retrieve_file_response(
raw_response=mock_response,
logging_obj=mock_logging_obj,
litellm_params={},
)

# Verify transformation
assert isinstance(result, OpenAIFileObject)
assert result.id == mock_response_data["uri"]
assert result.filename == mock_response_data["displayName"]
assert result.bytes == int(mock_response_data["sizeBytes"])
assert result.object == "file"
assert result.purpose == "user_data"
assert result.status == "processed" # ACTIVE state maps to processed
assert result.status_details is None

def test_transform_retrieve_file_response_failed_state(self):
"""Test transformation of Gemini file retrieval response with FAILED state"""
mock_response_data = {
"name": "files/test123",
"displayName": "test_file.pdf",
"mimeType": "application/pdf",
"sizeBytes": "1024",
"createTime": "2024-01-15T10:30:00.123456Z",
"uri": "https://generativelanguage.googleapis.com/v1beta/files/test123",
"state": "FAILED",
"error": {"message": "Upload failed", "code": "INTERNAL"},
}

mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = mock_response_data
mock_logging_obj = Mock()

result = self.handler.transform_retrieve_file_response(
raw_response=mock_response,
logging_obj=mock_logging_obj,
litellm_params={},
)

# Verify error state handling
assert result.status == "error"
assert result.status_details is not None
assert "message" in result.status_details

def test_transform_retrieve_file_response_processing_state(self):
"""Test transformation of Gemini file retrieval response with PROCESSING state"""
mock_response_data = {
"name": "files/test123",
"displayName": "test_file.pdf",
"mimeType": "application/pdf",
"sizeBytes": "1024",
"createTime": "2024-01-15T10:30:00.123456Z",
"uri": "https://generativelanguage.googleapis.com/v1beta/files/test123",
"state": "PROCESSING",
}

mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = mock_response_data
mock_logging_obj = Mock()

result = self.handler.transform_retrieve_file_response(
raw_response=mock_response,
logging_obj=mock_logging_obj,
litellm_params={},
)

# PROCESSING state should map to "uploaded" status
assert result.status == "uploaded"

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.
"""
# Mock error response from Gemini API (missing createTime)
mock_response_data = {
"error": {
"code": 400,
"message": "Invalid request",
"status": "INVALID_ARGUMENT",
}
}

mock_response = Mock(spec=httpx.Response)
mock_response.json.return_value = mock_response_data
mock_logging_obj = Mock()

# Should raise ValueError with helpful message
with pytest.raises(ValueError, match="Error parsing file retrieve response"):
self.handler.transform_retrieve_file_response(
raw_response=mock_response,
logging_obj=mock_logging_obj,
litellm_params={},
)

def test_validate_environment(self):
"""Test that validate_environment properly adds API key to headers"""
headers = {}
api_key = "test-gemini-api-key"

result_headers = self.handler.validate_environment(
headers=headers,
model="gemini-pro",
messages=[],
optional_params={},
litellm_params={},
api_key=api_key,
)

# Verify API key is added to headers
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)
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"
):
self.handler.validate_environment(
headers=headers,
model="gemini-pro",
messages=[],
optional_params={},
litellm_params={},
api_key=None,
)

def test_get_complete_url(self):
"""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,
model="gemini-pro",
optional_params={},
litellm_params={},
)

# Verify URL structure
assert api_base in url
assert "upload/v1beta/files" in url
assert f"key={api_key}" in url

def test_transform_delete_file_request_with_full_uri(self):
"""Test delete file request transformation with full URI"""
file_id = "https://generativelanguage.googleapis.com/v1beta/files/test123"
litellm_params = {
"api_key": "test-api-key",
"api_base": "https://generativelanguage.googleapis.com",
}

url, params = self.handler.transform_delete_file_request(
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
)

# Verify URL extraction
assert "files/test123" in url
assert "generativelanguage.googleapis.com" in url

Check failure

Code scanning / CodeQL

Incomplete URL substring sanitization High test

The string
generativelanguage.googleapis.com
may be at an arbitrary position in the sanitized URL.

Copilot Autofix

AI 8 months ago

In general, instead of checking that a trusted host appears as a substring of a URL, you should parse the URL and verify the hostname (and optionally scheme and path) explicitly. In Python, urllib.parse.urlparse provides a robust way to extract the hostname from a URL string and compare it against an expected value.

For this concrete test file, we only need to adjust the assertions that currently perform substring checks on the full URL. We can keep the behavior under test exactly the same (i.e., we don’t change GoogleAIStudioFilesHandler), but assert on parsed components instead of substring membership:

  • In test_get_complete_url, replace assert api_base in url with a parse-based check that the URL’s scheme and hostname match the expected base, and that the path contains the expected upload endpoint.
  • In test_transform_delete_file_request_with_full_uri, replace assert "generativelanguage.googleapis.com" in url with an assertion that urlparse(url).hostname == "generativelanguage.googleapis.com".
  • In test_transform_delete_file_request_with_file_name_only, replace assert "generativelanguage.googleapis.com" in url with a similar hostname assertion, and replace assert file_id in url with a more specific path-based assertion.

To implement this, we need to:

  • Import urlparse (or urllib.parse) at the top of tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py.
  • Update the relevant assertions to use urlparse(url) and compare hostname and path.

No changes to production code are required, and there is no change in the intended behavior being tested; we are only tightening how the tests verify the domain and path.

Suggested changeset 1
tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py

Autofix patch

Autofix patch
Run the following command in your local git repository to apply this patch
cat << 'EOF' | git apply
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
--- a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py
+++ b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py
@@ -5,6 +5,7 @@
 import os
 import pytest
 from unittest.mock import Mock, patch
+from urllib.parse import urlparse
 
 import httpx
 
@@ -253,8 +254,11 @@
         )
 
         # Verify URL structure
-        assert api_base in url
-        assert "upload/v1beta/files" in url
+        parsed = urlparse(url)
+        parsed_api_base = urlparse(api_base)
+        assert parsed.scheme == parsed_api_base.scheme
+        assert parsed.hostname == parsed_api_base.hostname
+        assert "upload/v1beta/files" in parsed.path
         assert f"key={api_key}" in url
 
     def test_transform_delete_file_request_with_full_uri(self):
@@ -272,8 +276,9 @@
         )
 
         # Verify URL extraction
-        assert "files/test123" in url
-        assert "generativelanguage.googleapis.com" in url
+        parsed = urlparse(url)
+        assert parsed.hostname == "generativelanguage.googleapis.com"
+        assert parsed.path.endswith("/files/test123")
         
         # Params should be empty (API key goes in header via validate_environment)
         assert params == {}
@@ -293,6 +298,7 @@
         )
 
         # Verify URL construction
-        assert file_id in url
-        assert "generativelanguage.googleapis.com" in url
+        parsed = urlparse(url)
+        assert parsed.hostname == "generativelanguage.googleapis.com"
+        assert parsed.path.endswith("/files/test123")
         assert params == {}
EOF
@@ -5,6 +5,7 @@
import os
import pytest
from unittest.mock import Mock, patch
from urllib.parse import urlparse

import httpx

@@ -253,8 +254,11 @@
)

# Verify URL structure
assert api_base in url
assert "upload/v1beta/files" in url
parsed = urlparse(url)
parsed_api_base = urlparse(api_base)
assert parsed.scheme == parsed_api_base.scheme
assert parsed.hostname == parsed_api_base.hostname
assert "upload/v1beta/files" in parsed.path
assert f"key={api_key}" in url

def test_transform_delete_file_request_with_full_uri(self):
@@ -272,8 +276,9 @@
)

# Verify URL extraction
assert "files/test123" in url
assert "generativelanguage.googleapis.com" in url
parsed = urlparse(url)
assert parsed.hostname == "generativelanguage.googleapis.com"
assert parsed.path.endswith("/files/test123")

# Params should be empty (API key goes in header via validate_environment)
assert params == {}
@@ -293,6 +298,7 @@
)

# Verify URL construction
assert file_id in url
assert "generativelanguage.googleapis.com" in url
parsed = urlparse(url)
assert parsed.hostname == "generativelanguage.googleapis.com"
assert parsed.path.endswith("/files/test123")
assert params == {}
Copilot is powered by AI and may make mistakes. Always verify output.

# Params should be empty (API key goes in header via validate_environment)
assert params == {}

def test_transform_delete_file_request_with_file_name_only(self):
"""Test delete file request transformation with file name only"""
file_id = "files/test123"
litellm_params = {
"api_key": "test-api-key",
"api_base": "https://generativelanguage.googleapis.com",
}

url, params = self.handler.transform_delete_file_request(
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
)

# Verify URL construction
assert file_id in url
assert "generativelanguage.googleapis.com" in url

Check failure

Code scanning / CodeQL

Incomplete URL substring sanitization High test

The string
generativelanguage.googleapis.com
may be at an arbitrary position in the sanitized URL.

Copilot Autofix

AI 8 months ago

In general, the problem with substring-based checks is that they do not ensure the host component of a URL is what you expect; the string can appear anywhere. The robust way to validate or assert URL structure is to parse the URL using a URL parser (such as urllib.parse.urlparse) and then inspect the specific components—scheme, netloc (host and port), path, and query parameters.

For this test, instead of asserting that "generativelanguage.googleapis.com" in url, we should parse url and assert that its netloc (or hostname) is exactly generativelanguage.googleapis.com. Similarly, for other structure-related checks (like assert file_id in url), we can tighten the test to assert on parsed.path.endswith(file_id) or similar, which still verifies the same behavior but with structured checks. Concretely:

  • Add an import of urllib.parse.urlparse near the top of tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py.
  • In test_transform_delete_file_request_with_full_uri, parse url and replace the existing host substring assertion with an assertion on parsed_url.netloc.
  • In test_transform_delete_file_request_with_file_name_only, parse url and:
    • Replace assert file_id in url with an assertion based on parsed_url.path (e.g., endswith(file_id)).
    • Replace assert "generativelanguage.googleapis.com" in url with an assertion on parsed_url.netloc.

This preserves the intended checks (correct host and inclusion of file id in the URL) but removes reliance on unsafe substring matches.

Suggested changeset 1
tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py

Autofix patch

Autofix patch
Run the following command in your local git repository to apply this patch
cat << 'EOF' | git apply
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
--- a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py
+++ b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py
@@ -7,6 +7,7 @@
 from unittest.mock import Mock, patch
 
 import httpx
+from urllib.parse import urlparse
 
 from litellm.llms.gemini.files.transformation import GoogleAIStudioFilesHandler
 from litellm.types.llms.openai import OpenAIFileObject
@@ -272,8 +273,9 @@
         )
 
         # Verify URL extraction
-        assert "files/test123" in url
-        assert "generativelanguage.googleapis.com" in url
+        parsed_url = urlparse(url)
+        assert parsed_url.netloc == "generativelanguage.googleapis.com"
+        assert parsed_url.path.endswith("/v1beta/files/test123")
         
         # Params should be empty (API key goes in header via validate_environment)
         assert params == {}
@@ -293,6 +295,7 @@
         )
 
         # Verify URL construction
-        assert file_id in url
-        assert "generativelanguage.googleapis.com" in url
+        parsed_url = urlparse(url)
+        assert parsed_url.netloc == "generativelanguage.googleapis.com"
+        assert parsed_url.path.endswith(f"/{file_id}")
         assert params == {}
EOF
@@ -7,6 +7,7 @@
from unittest.mock import Mock, patch

import httpx
from urllib.parse import urlparse

from litellm.llms.gemini.files.transformation import GoogleAIStudioFilesHandler
from litellm.types.llms.openai import OpenAIFileObject
@@ -272,8 +273,9 @@
)

# Verify URL extraction
assert "files/test123" in url
assert "generativelanguage.googleapis.com" in url
parsed_url = urlparse(url)
assert parsed_url.netloc == "generativelanguage.googleapis.com"
assert parsed_url.path.endswith("/v1beta/files/test123")

# Params should be empty (API key goes in header via validate_environment)
assert params == {}
@@ -293,6 +295,7 @@
)

# Verify URL construction
assert file_id in url
assert "generativelanguage.googleapis.com" in url
parsed_url = urlparse(url)
assert parsed_url.netloc == "generativelanguage.googleapis.com"
assert parsed_url.path.endswith(f"/{file_id}")
assert params == {}
Copilot is powered by AI and may make mistakes. Always verify output.
assert params == {}
Loading