mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-27 03:34:10 +00:00
Add Google AI Studio /v1/files upload API support (#9645)
* test: fix import for test
* fix: fix bad error string
* docs: cleanup files docs
* fix(files/main.py): cleanup error string
* style: initial commit with a provider/config pattern for files api
google ai studio files api onboarding
* fix: test
* feat(gemini/files/transformation.py): support gemini files api response transformation
* fix(gemini/files/transformation.py): return file id as gemini uri
allows id to be passed in to chat completion request, just like openai
* feat(llm_http_handler.py): support async route for files api on llm_http_handler
* fix: fix linting errors
* fix: fix model info check
* fix: fix ruff errors
* fix: fix linting errors
* Revert "fix: fix linting errors"
This reverts commit 926a5a527f.
* fix: fix linting errors
* test: fix test
* test: fix tests
This commit is contained in:
Krish Dholakia
GitHub
parent
d1abb9b68b
commit
0519c0c507
40 changed files with 1006 additions and 245 deletions
|
|
@ -7,11 +7,13 @@ import litellm
|
|||
import litellm.litellm_core_utils
|
||||
import litellm.types
|
||||
import litellm.types.utils
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.base_llm.audio_transcription.transformation import (
|
||||
BaseAudioTranscriptionConfig,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig
|
||||
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
|
||||
from litellm.llms.base_llm.files.transformation import BaseFilesConfig
|
||||
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
|
|
@ -26,7 +28,12 @@ from litellm.responses.streaming_iterator import (
|
|||
ResponsesAPIStreamingIterator,
|
||||
SyncResponsesAPIStreamingIterator,
|
||||
)
|
||||
from litellm.types.llms.openai import ResponseInputParam, ResponsesAPIResponse
|
||||
from litellm.types.llms.openai import (
|
||||
CreateFileRequest,
|
||||
OpenAIFileObject,
|
||||
ResponseInputParam,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.types.rerank import OptionalRerankParams, RerankResponse
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import EmbeddingResponse, FileTypes, TranscriptionResponse
|
||||
|
|
@ -240,6 +247,7 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
api_base = provider_config.get_complete_url(
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
stream=stream,
|
||||
|
|
@ -611,6 +619,7 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
api_base = provider_config.get_complete_url(
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
|
|
@ -884,6 +893,7 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
complete_url = provider_config.get_complete_url(
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
|
|
@ -1185,6 +1195,188 @@ class BaseLLMHTTPHandler:
|
|||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
def create_file(
|
||||
self,
|
||||
create_file_data: CreateFileRequest,
|
||||
litellm_params: dict,
|
||||
provider_config: BaseFilesConfig,
|
||||
headers: dict,
|
||||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
_is_async: bool = False,
|
||||
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]:
|
||||
"""
|
||||
Creates a file using Gemini's two-step upload process
|
||||
"""
|
||||
# get config from model, custom llm provider
|
||||
headers = provider_config.validate_environment(
|
||||
api_key=api_key,
|
||||
headers=headers,
|
||||
model="",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
)
|
||||
|
||||
api_base = provider_config.get_complete_url(
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
model="",
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
# Get the transformed request data for both steps
|
||||
transformed_request = provider_config.transform_create_file_request(
|
||||
model="",
|
||||
create_file_data=create_file_data,
|
||||
litellm_params=litellm_params,
|
||||
optional_params={},
|
||||
)
|
||||
|
||||
if _is_async:
|
||||
return self.async_create_file(
|
||||
transformed_request=transformed_request,
|
||||
litellm_params=litellm_params,
|
||||
provider_config=provider_config,
|
||||
headers=headers,
|
||||
api_base=api_base,
|
||||
logging_obj=logging_obj,
|
||||
client=client,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
if client is None or not isinstance(client, HTTPHandler):
|
||||
sync_httpx_client = _get_httpx_client()
|
||||
else:
|
||||
sync_httpx_client = client
|
||||
|
||||
try:
|
||||
# Step 1: Initial request to get upload URL
|
||||
initial_response = sync_httpx_client.post(
|
||||
url=api_base,
|
||||
headers={
|
||||
**headers,
|
||||
**transformed_request["initial_request"]["headers"],
|
||||
},
|
||||
data=json.dumps(transformed_request["initial_request"]["data"]),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
# Extract upload URL from response headers
|
||||
upload_url = initial_response.headers.get("X-Goog-Upload-URL")
|
||||
|
||||
if not upload_url:
|
||||
raise ValueError("Failed to get upload URL from initial request")
|
||||
|
||||
# Step 2: Upload the actual file
|
||||
upload_response = sync_httpx_client.post(
|
||||
url=upload_url,
|
||||
headers=transformed_request["upload_request"]["headers"],
|
||||
data=transformed_request["upload_request"]["data"],
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
return provider_config.transform_create_file_response(
|
||||
model=None,
|
||||
raw_response=upload_response,
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise self._handle_error(
|
||||
e=e,
|
||||
provider_config=provider_config,
|
||||
)
|
||||
|
||||
async def async_create_file(
|
||||
self,
|
||||
transformed_request: dict,
|
||||
litellm_params: dict,
|
||||
provider_config: BaseFilesConfig,
|
||||
headers: dict,
|
||||
api_base: str,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
):
|
||||
"""
|
||||
Creates a file using Gemini's two-step upload process
|
||||
"""
|
||||
if client is None or not isinstance(client, AsyncHTTPHandler):
|
||||
async_httpx_client = get_async_httpx_client(
|
||||
llm_provider=provider_config.custom_llm_provider
|
||||
)
|
||||
else:
|
||||
async_httpx_client = client
|
||||
|
||||
try:
|
||||
# Step 1: Initial request to get upload URL
|
||||
initial_response = await async_httpx_client.post(
|
||||
url=api_base,
|
||||
headers={
|
||||
**headers,
|
||||
**transformed_request["initial_request"]["headers"],
|
||||
},
|
||||
data=json.dumps(transformed_request["initial_request"]["data"]),
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
# Extract upload URL from response headers
|
||||
upload_url = initial_response.headers.get("X-Goog-Upload-URL")
|
||||
|
||||
if not upload_url:
|
||||
raise ValueError("Failed to get upload URL from initial request")
|
||||
|
||||
# Step 2: Upload the actual file
|
||||
upload_response = await async_httpx_client.post(
|
||||
url=upload_url,
|
||||
headers=transformed_request["upload_request"]["headers"],
|
||||
data=transformed_request["upload_request"]["data"],
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
return provider_config.transform_create_file_response(
|
||||
model=None,
|
||||
raw_response=upload_response,
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error creating file: {e}")
|
||||
raise self._handle_error(
|
||||
e=e,
|
||||
provider_config=provider_config,
|
||||
)
|
||||
|
||||
def list_files(self):
|
||||
"""
|
||||
Lists all files
|
||||
"""
|
||||
pass
|
||||
|
||||
def delete_file(self):
|
||||
"""
|
||||
Deletes a file
|
||||
"""
|
||||
pass
|
||||
|
||||
def retrieve_file(self):
|
||||
"""
|
||||
Returns the metadata of the file
|
||||
"""
|
||||
pass
|
||||
|
||||
def retrieve_file_content(self):
|
||||
"""
|
||||
Returns the content of the file
|
||||
"""
|
||||
pass
|
||||
|
||||
def _prepare_fake_stream_request(
|
||||
self,
|
||||
stream: bool,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue