o ưi"@sddlZddlZddlmZmZmZmZmZddl Z ddl m Z ddl m Z mZddlmZddlmZmZmZmZddlmZdd lmZeZGd d d e ZdS) N)Any CoroutineOptionalTupleUnion) LlmProviders) GCSBucketBaseGCSLoggingConfig)get_async_httpx_client)CreateFileRequestFileContentRequestHttpxBinaryResponseContentOpenAIFileObject)VERTEX_CREDENTIALS_TYPES) VertexAIJsonlFilesTransformationcs~eZdZdZfddZdedeedeedeedeed e e e j fd ee d efd d Zdededeedeedeedeed e e e j fd ee d e eeeeefffddZded eeeffddZdedeedeedeed e e e j fd ee d efddZdededeedeedeedeed e e e j fd ee d e eeeeefffddZZS)VertexAIFilesHandlerz~ Handles Calling VertexAI in OpenAI Files API format v1/files/* This implementation uploads files on GCS Buckets cstttjd|_dS)N)Z llm_provider)super__init__r rZ VERTEX_AIZasync_httpx_client)self __class__[/home/app/Keep/.python/lib/python3.10/site-packages/litellm/llms/vertex_ai/files/handler.pyr!s  zVertexAIFilesHandler.__init__create_file_dataapi_basevertex_credentialsvertex_projectvertex_locationtimeout max_retriesreturncst|jidIdH}|j|d|ddIdH} |d} tj|dd\} } |j| | | | dIdH} tj|| d S) N)kwargsvertex_instanceZpath_service_account)r#Zservice_account_json bucket_namefile)Zopenai_file_content)headersr$ object_namelogging_payload)rgcs_upload_response)Zget_gcs_logging_configZconstruct_request_headersvertex_ai_files_transformationZ7transform_openai_file_content_to_vertex_ai_file_contentgetZ_log_json_data_on_gcsZ3transform_gcs_bucket_response_to_openai_file_object)rrrrrrrr Zgcs_logging_configr&r$r(r'r)rrrasync_create_file's2    z&VertexAIFilesHandler.async_create_file _is_asyncc Cs:|r|j|||||||dSt|j|||||||dS)zn Creates a file on VertexAI GCS Bucket Only supported for Async litellm.acreate_file )rrrrrrr )r,asynciorun) rr-rrrrrrr rrr create_fileKs* z VertexAIFilesHandler.create_filefile_idcCsbtj|}|dr|dd}n|}d|vr!|dd\}}n|}d}tjj|dd}||fS)az Extract bucket name and object path from URL-encoded file_id. Expected format: gs%3A%2F%2Fbucket-name%2Fpath%2Fto%2Ffile Which decodes to: gs://bucket-name/path/to/file Returns: tuple: (bucket_name, url_encoded_object_path) - bucket_name: "bucket-name" - url_encoded_object_path: "path%2Fto%2Ffile" zgs://N/r)safe)urllibparseunquote startswithsplitquote)rr1 decoded_path full_pathr$Z object_pathencoded_object_pathrrr'_extract_bucket_and_object_from_file_idss zZdownload_kwargs file_contentr<Z mock_responserrr afile_contents2      z"VertexAIFilesHandler.afile_contentc Cs6|r |j||||||dSt|j||||||dS)a Download file content from GCS bucket for VertexAI files. Supports both sync and async operations. Args: _is_async: Whether to run asynchronously file_content_request: Contains file_id (URL-encoded GCS path) api_base: API base (unused for GCS operations) vertex_credentials: VertexAI credentials vertex_project: VertexAI project ID vertex_location: VertexAI location timeout: Request timeout max_retries: Max retry attempts Returns: HttpxBinaryResponseContent or Coroutine: Binary content wrapped in compatible response format )r@rrrrr )rNr.r/) rr-r@rrrrrr rrrrMs& z!VertexAIFilesHandler.file_content)__name__ __module__ __qualname____doc__rr rstrrrfloatrJTimeoutintrr,boolrrr0rr?r r rNrM __classcell__rrrrrs   $   (  5   r)r. urllib.parser6typingrrrrrrJZlitellmrZ/litellm.integrations.gcs_bucket.gcs_bucket_baserr Z&litellm.llms.custom_httpx.http_handlerr Zlitellm.types.llms.openair r r rZlitellm.types.llms.vertex_airZtransformationrr*rrrrrs