390 lines
13 KiB
Python
390 lines
13 KiB
Python
|
"""
|
||
|
|
Supports writing files to Google AI Studio Files API.
|
||
|
|
|
||
|
|
For vertex ai, check out the vertex_ai/files/handler.py file.
|
||
|
|
"""
|
||
|
|
import time
|
||
|
|
from typing import Any, List, Literal, Optional
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
from openai.types.file_deleted import FileDeleted
|
||
|
|
|
||
|
|
from litellm._logging import verbose_logger
|
||
|
|
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
|
||
|
|
from litellm.llms.base_llm.files.transformation import (
|
||
|
|
BaseFilesConfig,
|
||
|
|
LiteLLMLoggingObj,
|
||
|
|
)
|
||
|
|
from litellm.types.llms.gemini import GeminiCreateFilesResponseObject
|
||
|
|
from litellm.types.llms.openai import (
|
||
|
|
AllMessageValues,
|
||
|
|
CreateFileRequest,
|
||
|
|
HttpxBinaryResponseContent,
|
||
|
|
OpenAICreateFileRequestOptionalParams,
|
||
|
|
OpenAIFileObject,
|
||
|
|
)
|
||
|
|
from litellm.types.utils import LlmProviders
|
||
|
|
|
||
|
|
from ..common_utils import GeminiModelInfo
|
||
|
|
|
||
|
|
|
||
|
|
class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
|
||
|
|
def __init__(self):
|
||
|
|
pass
|
||
|
|
|
||
|
|
@property
|
||
|
|
def custom_llm_provider(self) -> LlmProviders:
|
||
|
|
return LlmProviders.GEMINI
|
||
|
|
|
||
|
|
def validate_environment(
|
||
|
|
self,
|
||
|
|
headers: dict[Any, Any],
|
||
|
|
model: str,
|
||
|
|
messages: List[AllMessageValues],
|
||
|
|
optional_params: dict[Any, Any],
|
||
|
|
litellm_params: dict[Any, Any],
|
||
|
|
api_key: Optional[str] = None,
|
||
|
|
api_base: Optional[str] = None,
|
||
|
|
) -> dict[Any, Any]:
|
||
|
|
"""
|
||
|
|
Validate environment and add Gemini API key to headers.
|
||
|
|
Google AI Studio uses x-goog-api-key header for authentication.
|
||
|
|
"""
|
||
|
|
resolved_api_key = self.get_api_key(api_key)
|
||
|
|
if not resolved_api_key:
|
||
|
|
raise ValueError(
|
||
|
|
"GEMINI_API_KEY is required for Google AI Studio file operations"
|
||
|
|
)
|
||
|
|
|
||
|
|
headers["x-goog-api-key"] = resolved_api_key
|
||
|
|
return headers
|
||
|
|
|
||
|
|
def get_complete_url(
|
||
|
|
self,
|
||
|
|
api_base: Optional[str],
|
||
|
|
api_key: Optional[str],
|
||
|
|
model: str,
|
||
|
|
optional_params: dict,
|
||
|
|
litellm_params: dict,
|
||
|
|
stream: Optional[bool] = None,
|
||
|
|
) -> str:
|
||
|
|
"""
|
||
|
|
OPTIONAL
|
||
|
|
|
||
|
|
Get the complete url for the request
|
||
|
|
|
||
|
|
Some providers need `model` in `api_base`
|
||
|
|
"""
|
||
|
|
endpoint = "upload/v1beta/files"
|
||
|
|
api_base = self.get_api_base(api_base)
|
||
|
|
if not api_base:
|
||
|
|
raise ValueError("api_base is required")
|
||
|
|
|
||
|
|
# Get API key from multiple sources
|
||
|
|
final_api_key = api_key or litellm_params.get("api_key") or self.get_api_key()
|
||
|
|
if not final_api_key:
|
||
|
|
raise ValueError("api_key is required")
|
||
|
|
|
||
|
|
url = "{}/{}?key={}".format(api_base, endpoint, final_api_key)
|
||
|
|
return url
|
||
|
|
|
||
|
|
def get_supported_openai_params(
|
||
|
|
self, model: str
|
||
|
|
) -> List[OpenAICreateFileRequestOptionalParams]:
|
||
|
|
return []
|
||
|
|
|
||
|
|
def map_openai_params(
|
||
|
|
self,
|
||
|
|
non_default_params: dict,
|
||
|
|
optional_params: dict,
|
||
|
|
model: str,
|
||
|
|
drop_params: bool,
|
||
|
|
) -> dict:
|
||
|
|
return optional_params
|
||
|
|
|
||
|
|
def transform_create_file_request(
|
||
|
|
self,
|
||
|
|
model: str,
|
||
|
|
create_file_data: CreateFileRequest,
|
||
|
|
optional_params: dict,
|
||
|
|
litellm_params: dict,
|
||
|
|
) -> dict:
|
||
|
|
"""
|
||
|
|
Transform the OpenAI-style file creation request into Gemini's format
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
dict: Contains both request data and headers for the two-step upload
|
||
|
|
"""
|
||
|
|
# Extract the file information
|
||
|
|
file_data = create_file_data.get("file")
|
||
|
|
if file_data is None:
|
||
|
|
raise ValueError("File data is required")
|
||
|
|
|
||
|
|
# Use the common utility function to extract file data
|
||
|
|
extracted_data = extract_file_data(file_data)
|
||
|
|
|
||
|
|
# Get file size
|
||
|
|
file_size = len(extracted_data["content"])
|
||
|
|
|
||
|
|
# Step 1: Initial resumable upload request
|
||
|
|
headers = {
|
||
|
|
"X-Goog-Upload-Protocol": "resumable",
|
||
|
|
"X-Goog-Upload-Command": "start",
|
||
|
|
"X-Goog-Upload-Header-Content-Length": str(file_size),
|
||
|
|
"X-Goog-Upload-Header-Content-Type": extracted_data["content_type"],
|
||
|
|
"Content-Type": "application/json",
|
||
|
|
}
|
||
|
|
headers.update(extracted_data["headers"]) # Add any custom headers
|
||
|
|
|
||
|
|
# Initial metadata request body
|
||
|
|
initial_data = {
|
||
|
|
"file": {
|
||
|
|
"display_name": extracted_data["filename"] or str(int(time.time()))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
# Step 2: Actual file upload data
|
||
|
|
upload_headers = {
|
||
|
|
"Content-Length": str(file_size),
|
||
|
|
"X-Goog-Upload-Offset": "0",
|
||
|
|
"X-Goog-Upload-Command": "upload, finalize",
|
||
|
|
}
|
||
|
|
|
||
|
|
return {
|
||
|
|
"initial_request": {"headers": headers, "data": initial_data},
|
||
|
|
"upload_request": {
|
||
|
|
"headers": upload_headers,
|
||
|
|
"data": extracted_data["content"],
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
def transform_create_file_response(
|
||
|
|
self,
|
||
|
|
model: Optional[str],
|
||
|
|
raw_response: httpx.Response,
|
||
|
|
logging_obj: LiteLLMLoggingObj,
|
||
|
|
litellm_params: dict,
|
||
|
|
) -> OpenAIFileObject:
|
||
|
|
"""
|
||
|
|
Transform Gemini's file upload response into OpenAI-style FileObject
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
response_json = raw_response.json()
|
||
|
|
|
||
|
|
response_object = GeminiCreateFilesResponseObject(
|
||
|
|
**response_json.get("file", {}) # type: ignore
|
||
|
|
)
|
||
|
|
|
||
|
|
# Extract file information from Gemini response
|
||
|
|
|
||
|
|
return OpenAIFileObject(
|
||
|
|
id=response_object["uri"], # Gemini uses URI as identifier
|
||
|
|
bytes=int(
|
||
|
|
response_object["sizeBytes"]
|
||
|
|
), # Gemini doesn't return file size
|
||
|
|
created_at=int(
|
||
|
|
time.mktime(
|
||
|
|
time.strptime(
|
||
|
|
response_object["createTime"].replace("Z", "+00:00"),
|
||
|
|
"%Y-%m-%dT%H:%M:%S.%f%z",
|
||
|
|
)
|
||
|
|
)
|
||
|
|
),
|
||
|
|
filename=response_object["displayName"],
|
||
|
|
object="file",
|
||
|
|
purpose="user_data", # Default to assistants as that's the main use case
|
||
|
|
status="uploaded",
|
||
|
|
status_details=None,
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
verbose_logger.exception(f"Error parsing file upload response: {str(e)}")
|
||
|
|
raise ValueError(f"Error parsing file upload response: {str(e)}")
|
||
|
|
|
||
|
|
def transform_retrieve_file_request(
|
||
|
|
self,
|
||
|
|
file_id: str,
|
||
|
|
optional_params: dict,
|
||
|
|
litellm_params: dict,
|
||
|
|
) -> tuple[str, dict]:
|
||
|
|
"""
|
||
|
|
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.
|
||
|
|
"""
|
||
|
|
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)
|
||
|
|
|
||
|
|
# Return empty params dict - API key is already in URL, no query params needed
|
||
|
|
return url, {}
|
||
|
|
|
||
|
|
def transform_retrieve_file_response(
|
||
|
|
self,
|
||
|
|
raw_response: httpx.Response,
|
||
|
|
logging_obj: LiteLLMLoggingObj,
|
||
|
|
litellm_params: dict,
|
||
|
|
) -> OpenAIFileObject:
|
||
|
|
"""
|
||
|
|
Transform Gemini's file retrieval response into OpenAI-style FileObject
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
response_json = raw_response.json()
|
||
|
|
|
||
|
|
# Map Gemini state to OpenAI status
|
||
|
|
gemini_state = response_json.get("state", "STATE_UNSPECIFIED")
|
||
|
|
# Explicitly type status as the Literal union
|
||
|
|
if gemini_state == "ACTIVE":
|
||
|
|
status: Literal["uploaded", "processed", "error"] = "processed"
|
||
|
|
elif gemini_state == "FAILED":
|
||
|
|
status = "error"
|
||
|
|
else:
|
||
|
|
status = "uploaded"
|
||
|
|
|
||
|
|
return OpenAIFileObject(
|
||
|
|
id=response_json.get("uri", ""),
|
||
|
|
bytes=int(response_json.get("sizeBytes", 0)),
|
||
|
|
created_at=int(
|
||
|
|
time.mktime(
|
||
|
|
time.strptime(
|
||
|
|
response_json["createTime"].replace("Z", "+00:00"),
|
||
|
|
"%Y-%m-%dT%H:%M:%S.%f%z",
|
||
|
|
)
|
||
|
|
)
|
||
|
|
),
|
||
|
|
filename=response_json.get("displayName", ""),
|
||
|
|
object="file",
|
||
|
|
purpose="user_data",
|
||
|
|
status=status,
|
||
|
|
status_details=str(response_json.get("error", ""))
|
||
|
|
if gemini_state == "FAILED"
|
||
|
|
else None,
|
||
|
|
)
|
||
|
|
except Exception as e:
|
||
|
|
verbose_logger.exception(f"Error parsing file retrieve response: {str(e)}")
|
||
|
|
raise ValueError(f"Error parsing file retrieve response: {str(e)}")
|
||
|
|
|
||
|
|
def transform_delete_file_request(
|
||
|
|
self,
|
||
|
|
file_id: str,
|
||
|
|
optional_params: dict,
|
||
|
|
litellm_params: dict,
|
||
|
|
) -> tuple[str, dict]:
|
||
|
|
"""
|
||
|
|
Transform delete file request for Google AI Studio.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
file_id: The file URI (e.g., "files/abc123" or full URI)
|
||
|
|
optional_params: Optional parameters
|
||
|
|
litellm_params: LiteLLM parameters containing api_key
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
tuple[str, dict]: (url, params) for the DELETE request
|
||
|
|
"""
|
||
|
|
api_base = self.get_api_base(litellm_params.get("api_base"))
|
||
|
|
if not api_base:
|
||
|
|
raise ValueError("api_base is required")
|
||
|
|
|
||
|
|
# Get API key from multiple sources (same pattern as get_complete_url)
|
||
|
|
api_key = litellm_params.get("api_key") or self.get_api_key()
|
||
|
|
if not api_key:
|
||
|
|
raise ValueError("api_key is required")
|
||
|
|
|
||
|
|
# Extract file name from URI if full URI is provided
|
||
|
|
# file_id could be "files/abc123" or "https://generativelanguage.googleapis.com/v1beta/files/abc123"
|
||
|
|
if file_id.startswith("http"):
|
||
|
|
# Extract the file path from full URI
|
||
|
|
file_name = file_id.split("/v1beta/")[-1]
|
||
|
|
else:
|
||
|
|
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}"
|
||
|
|
|
||
|
|
# Add API key as header (Google AI Studio uses x-goog-api-key header)
|
||
|
|
params: dict = {}
|
||
|
|
|
||
|
|
return url, params
|
||
|
|
|
||
|
|
def transform_delete_file_response(
|
||
|
|
self,
|
||
|
|
raw_response: httpx.Response,
|
||
|
|
logging_obj: LiteLLMLoggingObj,
|
||
|
|
litellm_params: dict,
|
||
|
|
) -> FileDeleted:
|
||
|
|
"""
|
||
|
|
Transform Gemini's file delete response into OpenAI-style FileDeleted.
|
||
|
|
|
||
|
|
Google AI Studio returns an empty JSON object {} on successful deletion.
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
# Google AI Studio returns {} on successful deletion
|
||
|
|
if raw_response.status_code == 200:
|
||
|
|
# Extract file ID from the request URL if possible
|
||
|
|
file_id = "deleted"
|
||
|
|
if hasattr(raw_response, "request") and raw_response.request:
|
||
|
|
url = str(raw_response.request.url)
|
||
|
|
if "/files/" in url:
|
||
|
|
file_id = url.split("/files/")[-1].split("?")[0]
|
||
|
|
# Add the files/ prefix if not present
|
||
|
|
if not file_id.startswith("files/"):
|
||
|
|
file_id = f"files/{file_id}"
|
||
|
|
|
||
|
|
return FileDeleted(id=file_id, deleted=True, object="file")
|
||
|
|
else:
|
||
|
|
raise ValueError(f"Failed to delete file: {raw_response.text}")
|
||
|
|
except Exception as e:
|
||
|
|
verbose_logger.exception(f"Error parsing file delete response: {str(e)}")
|
||
|
|
raise ValueError(f"Error parsing file delete response: {str(e)}")
|
||
|
|
|
||
|
|
def transform_list_files_request(
|
||
|
|
self,
|
||
|
|
purpose: Optional[str],
|
||
|
|
optional_params: dict,
|
||
|
|
litellm_params: dict,
|
||
|
|
) -> tuple[str, dict]:
|
||
|
|
raise NotImplementedError(
|
||
|
|
"GoogleAIStudioFilesHandler does not support file listing"
|
||
|
|
)
|
||
|
|
|
||
|
|
def transform_list_files_response(
|
||
|
|
self,
|
||
|
|
raw_response: httpx.Response,
|
||
|
|
logging_obj: LiteLLMLoggingObj,
|
||
|
|
litellm_params: dict,
|
||
|
|
) -> List[OpenAIFileObject]:
|
||
|
|
raise NotImplementedError(
|
||
|
|
"GoogleAIStudioFilesHandler does not support file listing"
|
||
|
|
)
|
||
|
|
|
||
|
|
def transform_file_content_request(
|
||
|
|
self,
|
||
|
|
file_content_request,
|
||
|
|
optional_params: dict,
|
||
|
|
litellm_params: dict,
|
||
|
|
) -> tuple[str, dict]:
|
||
|
|
raise NotImplementedError(
|
||
|
|
"GoogleAIStudioFilesHandler does not support file content retrieval"
|
||
|
|
)
|
||
|
|
|
||
|
|
def transform_file_content_response(
|
||
|
|
self,
|
||
|
|
raw_response: httpx.Response,
|
||
|
|
logging_obj: LiteLLMLoggingObj,
|
||
|
|
litellm_params: dict,
|
||
|
|
) -> HttpxBinaryResponseContent:
|
||
|
|
raise NotImplementedError(
|
||
|
|
"GoogleAIStudioFilesHandler does not support file content retrieval"
|
||
|
|
)
|