Spaces:
Sleeping
Sleeping
File size: 5,430 Bytes
469eae6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 |
"""
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 List, Optional
import httpx
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 (
CreateFileRequest,
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 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")
if not api_key:
raise ValueError("api_key is required")
url = "{}/{}?key={}".format(api_base, endpoint, 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)}")
|