diff --git a/litellm/__init__.py b/litellm/__init__.py index 4f0b0a16be7..d2b51696640 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -122,6 +122,9 @@ langsmith_batch_size: Optional[int] = None prometheus_initialize_budget_metrics: Optional[bool] = False argilla_batch_size: Optional[int] = None datadog_use_v1: Optional[bool] = False # if you want to use v1 datadog logged payload +gcs_pub_sub_use_v1: Optional[bool] = ( + False # if you want to use v1 gcs pubsub logged payload +) argilla_transformation_object: Optional[Dict[str, Any]] = None _async_input_callback: List[Union[str, Callable, CustomLogger]] = ( [] diff --git a/litellm/integrations/gcs_pubsub/pub_sub.py b/litellm/integrations/gcs_pubsub/pub_sub.py index e94c853f3ff..1b078df7bc8 100644 --- a/litellm/integrations/gcs_pubsub/pub_sub.py +++ b/litellm/integrations/gcs_pubsub/pub_sub.py @@ -10,13 +10,16 @@ import asyncio import json import os import traceback -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from litellm.types.utils import StandardLoggingPayload if TYPE_CHECKING: from litellm.proxy._types import SpendLogsPayload else: SpendLogsPayload = Any +import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.llms.custom_httpx.http_handler import ( @@ -61,7 +64,7 @@ class GcsPubSubLogger(CustomBatchLogger): self.flush_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) asyncio.create_task(self.periodic_flush()) - self.log_queue: List[SpendLogsPayload] = [] + self.log_queue: List[Union[SpendLogsPayload, StandardLoggingPayload]] = [] async def construct_request_headers(self) -> Dict[str, str]: """Construct authorization headers using Vertex AI auth""" @@ -115,13 +118,20 @@ class GcsPubSubLogger(CustomBatchLogger): verbose_logger.debug( "PubSub: Logging - Enters logging function for model %s", kwargs ) - spend_logs_payload = get_logging_payload( - kwargs=kwargs, - response_obj=response_obj, - start_time=start_time, - end_time=end_time, - ) - self.log_queue.append(spend_logs_payload) + standard_logging_payload = kwargs.get("standard_logging_object", None) + + # Backwards compatibility with old logging payload + if litellm.gcs_pub_sub_use_v1 is True: + spend_logs_payload = get_logging_payload( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + ) + self.log_queue.append(spend_logs_payload) + else: + # New logging payload, StandardLoggingPayload + self.log_queue.append(standard_logging_payload) if len(self.log_queue) >= self.batch_size: await self.async_send_batch() @@ -155,7 +165,7 @@ class GcsPubSubLogger(CustomBatchLogger): self.log_queue.clear() async def publish_message( - self, message: SpendLogsPayload + self, message: Union[SpendLogsPayload, StandardLoggingPayload] ) -> Optional[Dict[str, Any]]: """ Publish message to Google Cloud Pub/Sub using REST API