chore: initial snapshot for gitea/github upload
This commit is contained in:
Your Name
@@ -0,0 +1,389 @@
|
||||
"""
|
||||
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"
|
||||
)
|
||||
Reference in New Issue
Block a user