Spaces:
Running
Running
| import json | |
| import os | |
| from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union | |
| from litellm._logging import verbose_logger | |
| from litellm.integrations.custom_batch_logger import CustomBatchLogger | |
| from litellm.llms.custom_httpx.http_handler import ( | |
| get_async_httpx_client, | |
| httpxSpecialProvider, | |
| ) | |
| from litellm.types.integrations.gcs_bucket import * | |
| from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload | |
| if TYPE_CHECKING: | |
| from litellm.llms.vertex_ai.vertex_llm_base import VertexBase | |
| else: | |
| VertexBase = Any | |
| IAM_AUTH_KEY = "IAM_AUTH" | |
| class GCSBucketBase(CustomBatchLogger): | |
| def __init__(self, bucket_name: Optional[str] = None, **kwargs) -> None: | |
| self.async_httpx_client = get_async_httpx_client( | |
| llm_provider=httpxSpecialProvider.LoggingCallback | |
| ) | |
| _path_service_account = os.getenv("GCS_PATH_SERVICE_ACCOUNT") | |
| _bucket_name = bucket_name or os.getenv("GCS_BUCKET_NAME") | |
| self.path_service_account_json: Optional[str] = _path_service_account | |
| self.BUCKET_NAME: Optional[str] = _bucket_name | |
| self.vertex_instances: Dict[str, VertexBase] = {} | |
| super().__init__(**kwargs) | |
| async def construct_request_headers( | |
| self, | |
| service_account_json: Optional[str], | |
| vertex_instance: Optional[VertexBase] = None, | |
| ) -> Dict[str, str]: | |
| from litellm import vertex_chat_completion | |
| if vertex_instance is None: | |
| vertex_instance = vertex_chat_completion | |
| _auth_header, vertex_project = await vertex_instance._ensure_access_token_async( | |
| credentials=service_account_json, | |
| project_id=None, | |
| custom_llm_provider="vertex_ai", | |
| ) | |
| auth_header, _ = vertex_instance._get_token_and_url( | |
| model="gcs-bucket", | |
| auth_header=_auth_header, | |
| vertex_credentials=service_account_json, | |
| vertex_project=vertex_project, | |
| vertex_location=None, | |
| gemini_api_key=None, | |
| stream=None, | |
| custom_llm_provider="vertex_ai", | |
| api_base=None, | |
| ) | |
| verbose_logger.debug("constructed auth_header %s", auth_header) | |
| headers = { | |
| "Authorization": f"Bearer {auth_header}", # auth_header | |
| "Content-Type": "application/json", | |
| } | |
| return headers | |
| def sync_construct_request_headers(self) -> Dict[str, str]: | |
| from litellm import vertex_chat_completion | |
| _auth_header, vertex_project = vertex_chat_completion._ensure_access_token( | |
| credentials=self.path_service_account_json, | |
| project_id=None, | |
| custom_llm_provider="vertex_ai", | |
| ) | |
| auth_header, _ = vertex_chat_completion._get_token_and_url( | |
| model="gcs-bucket", | |
| auth_header=_auth_header, | |
| vertex_credentials=self.path_service_account_json, | |
| vertex_project=vertex_project, | |
| vertex_location=None, | |
| gemini_api_key=None, | |
| stream=None, | |
| custom_llm_provider="vertex_ai", | |
| api_base=None, | |
| ) | |
| verbose_logger.debug("constructed auth_header %s", auth_header) | |
| headers = { | |
| "Authorization": f"Bearer {auth_header}", # auth_header | |
| "Content-Type": "application/json", | |
| } | |
| return headers | |
| def _handle_folders_in_bucket_name( | |
| self, | |
| bucket_name: str, | |
| object_name: str, | |
| ) -> Tuple[str, str]: | |
| """ | |
| Handles when the user passes a bucket name with a folder postfix | |
| Example: | |
| - Bucket name: "my-bucket/my-folder/dev" | |
| - Object name: "my-object" | |
| - Returns: bucket_name="my-bucket", object_name="my-folder/dev/my-object" | |
| """ | |
| if "/" in bucket_name: | |
| bucket_name, prefix = bucket_name.split("/", 1) | |
| object_name = f"{prefix}/{object_name}" | |
| return bucket_name, object_name | |
| return bucket_name, object_name | |
| async def get_gcs_logging_config( | |
| self, kwargs: Optional[Dict[str, Any]] = {} | |
| ) -> GCSLoggingConfig: | |
| """ | |
| This function is used to get the GCS logging config for the GCS Bucket Logger. | |
| It checks if the dynamic parameters are provided in the kwargs and uses them to get the GCS logging config. | |
| If no dynamic parameters are provided, it uses the default values. | |
| """ | |
| if kwargs is None: | |
| kwargs = {} | |
| standard_callback_dynamic_params: Optional[ | |
| StandardCallbackDynamicParams | |
| ] = kwargs.get("standard_callback_dynamic_params", None) | |
| bucket_name: str | |
| path_service_account: Optional[str] | |
| if standard_callback_dynamic_params is not None: | |
| verbose_logger.debug("Using dynamic GCS logging") | |
| verbose_logger.debug( | |
| "standard_callback_dynamic_params: %s", standard_callback_dynamic_params | |
| ) | |
| _bucket_name: Optional[str] = ( | |
| standard_callback_dynamic_params.get("gcs_bucket_name", None) | |
| or self.BUCKET_NAME | |
| ) | |
| _path_service_account: Optional[str] = ( | |
| standard_callback_dynamic_params.get("gcs_path_service_account", None) | |
| or self.path_service_account_json | |
| ) | |
| if _bucket_name is None: | |
| raise ValueError( | |
| "GCS_BUCKET_NAME is not set in the environment, but GCS Bucket is being used as a logging callback. Please set 'GCS_BUCKET_NAME' in the environment." | |
| ) | |
| bucket_name = _bucket_name | |
| path_service_account = _path_service_account | |
| vertex_instance = await self.get_or_create_vertex_instance( | |
| credentials=path_service_account | |
| ) | |
| else: | |
| # If no dynamic parameters, use the default instance | |
| if self.BUCKET_NAME is None: | |
| raise ValueError( | |
| "GCS_BUCKET_NAME is not set in the environment, but GCS Bucket is being used as a logging callback. Please set 'GCS_BUCKET_NAME' in the environment." | |
| ) | |
| bucket_name = self.BUCKET_NAME | |
| path_service_account = self.path_service_account_json | |
| vertex_instance = await self.get_or_create_vertex_instance( | |
| credentials=path_service_account | |
| ) | |
| return GCSLoggingConfig( | |
| bucket_name=bucket_name, | |
| vertex_instance=vertex_instance, | |
| path_service_account=path_service_account, | |
| ) | |
| async def get_or_create_vertex_instance( | |
| self, credentials: Optional[str] | |
| ) -> VertexBase: | |
| """ | |
| This function is used to get the Vertex instance for the GCS Bucket Logger. | |
| It checks if the Vertex instance is already created and cached, if not it creates a new instance and caches it. | |
| """ | |
| from litellm.llms.vertex_ai.vertex_llm_base import VertexBase | |
| _in_memory_key = self._get_in_memory_key_for_vertex_instance(credentials) | |
| if _in_memory_key not in self.vertex_instances: | |
| vertex_instance = VertexBase() | |
| await vertex_instance._ensure_access_token_async( | |
| credentials=credentials, | |
| project_id=None, | |
| custom_llm_provider="vertex_ai", | |
| ) | |
| self.vertex_instances[_in_memory_key] = vertex_instance | |
| return self.vertex_instances[_in_memory_key] | |
| def _get_in_memory_key_for_vertex_instance(self, credentials: Optional[str]) -> str: | |
| """ | |
| Returns key to use for caching the Vertex instance in-memory. | |
| When using Vertex with Key based logging, we need to cache the Vertex instance in-memory. | |
| - If a credentials string is provided, it is used as the key. | |
| - If no credentials string is provided, "IAM_AUTH" is used as the key. | |
| """ | |
| return credentials or IAM_AUTH_KEY | |
| async def download_gcs_object(self, object_name: str, **kwargs): | |
| """ | |
| Download an object from GCS. | |
| https://cloud.google.com/storage/docs/downloading-objects#download-object-json | |
| """ | |
| try: | |
| gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config( | |
| kwargs=kwargs | |
| ) | |
| headers = await self.construct_request_headers( | |
| vertex_instance=gcs_logging_config["vertex_instance"], | |
| service_account_json=gcs_logging_config["path_service_account"], | |
| ) | |
| bucket_name = gcs_logging_config["bucket_name"] | |
| bucket_name, object_name = self._handle_folders_in_bucket_name( | |
| bucket_name=bucket_name, | |
| object_name=object_name, | |
| ) | |
| url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{object_name}?alt=media" | |
| # Send the GET request to download the object | |
| response = await self.async_httpx_client.get(url=url, headers=headers) | |
| if response.status_code != 200: | |
| verbose_logger.error( | |
| "GCS object download error: %s", str(response.text) | |
| ) | |
| return None | |
| verbose_logger.debug( | |
| "GCS object download response status code: %s", response.status_code | |
| ) | |
| # Return the content of the downloaded object | |
| return response.content | |
| except Exception as e: | |
| verbose_logger.error("GCS object download error: %s", str(e)) | |
| return None | |
| async def delete_gcs_object(self, object_name: str, **kwargs): | |
| """ | |
| Delete an object from GCS. | |
| """ | |
| try: | |
| gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config( | |
| kwargs=kwargs | |
| ) | |
| headers = await self.construct_request_headers( | |
| vertex_instance=gcs_logging_config["vertex_instance"], | |
| service_account_json=gcs_logging_config["path_service_account"], | |
| ) | |
| bucket_name = gcs_logging_config["bucket_name"] | |
| bucket_name, object_name = self._handle_folders_in_bucket_name( | |
| bucket_name=bucket_name, | |
| object_name=object_name, | |
| ) | |
| url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{object_name}" | |
| # Send the DELETE request to delete the object | |
| response = await self.async_httpx_client.delete(url=url, headers=headers) | |
| if (response.status_code != 200) or (response.status_code != 204): | |
| verbose_logger.error( | |
| "GCS object delete error: %s, status code: %s", | |
| str(response.text), | |
| response.status_code, | |
| ) | |
| return None | |
| verbose_logger.debug( | |
| "GCS object delete response status code: %s, response: %s", | |
| response.status_code, | |
| response.text, | |
| ) | |
| # Return the content of the downloaded object | |
| return response.text | |
| except Exception as e: | |
| verbose_logger.error("GCS object download error: %s", str(e)) | |
| return None | |
| async def _log_json_data_on_gcs( | |
| self, | |
| headers: Dict[str, str], | |
| bucket_name: str, | |
| object_name: str, | |
| logging_payload: Union[StandardLoggingPayload, str], | |
| ): | |
| """ | |
| Helper function to make POST request to GCS Bucket in the specified bucket. | |
| """ | |
| if isinstance(logging_payload, str): | |
| json_logged_payload = logging_payload | |
| else: | |
| json_logged_payload = json.dumps(logging_payload, default=str) | |
| bucket_name, object_name = self._handle_folders_in_bucket_name( | |
| bucket_name=bucket_name, | |
| object_name=object_name, | |
| ) | |
| response = await self.async_httpx_client.post( | |
| headers=headers, | |
| url=f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={object_name}", | |
| data=json_logged_payload, | |
| ) | |
| if response.status_code != 200: | |
| verbose_logger.error("GCS Bucket logging error: %s", str(response.text)) | |
| verbose_logger.debug("GCS Bucket response %s", response) | |
| verbose_logger.debug("GCS Bucket status code %s", response.status_code) | |
| verbose_logger.debug("GCS Bucket response.text %s", response.text) | |
| return response.json() | |