From 19240b6cfd4a9b44ba21dad59548ddc33f9abfeb Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 12 Aug 2024 16:06:10 -0700 Subject: [PATCH] feat gcs log user api key metadata --- litellm/integrations/gcs_bucket.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/litellm/integrations/gcs_bucket.py b/litellm/integrations/gcs_bucket.py index 3fb778e2428..a16d952861e 100644 --- a/litellm/integrations/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket.py @@ -13,7 +13,7 @@ from litellm.litellm_core_utils.logging_utils import ( convert_litellm_response_object_to_dict, ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.proxy._types import CommonProxyErrors, SpendLogsPayload +from litellm.proxy._types import CommonProxyErrors, SpendLogsMetadata, SpendLogsPayload class RequestKwargs(TypedDict): @@ -27,6 +27,8 @@ class GCSBucketPayload(TypedDict): response_obj: Optional[Dict] start_time: str end_time: str + response_cost: Optional[float] + spend_log_metadata: str class GCSBucketLogger(CustomLogger): @@ -78,11 +80,12 @@ class GCSBucketLogger(CustomLogger): kwargs, response_obj, start_time_str, end_time_str ) + json_logged_payload = json.dumps(logging_payload) object_name = response_obj["id"] response = await self.async_httpx_client.post( headers=headers, url=f"https://storage.googleapis.com/upload/storage/v1/b/{self.BUCKET_NAME}/o?uploadType=media&name={object_name}", - json=logging_payload, + data=json_logged_payload, ) if response.status_code != 200: @@ -121,6 +124,10 @@ class GCSBucketLogger(CustomLogger): async def get_gcs_payload( self, kwargs, response_obj, start_time, end_time ) -> GCSBucketPayload: + from litellm.proxy.spend_tracking.spend_tracking_utils import ( + get_logging_payload, + ) + request_kwargs = RequestKwargs( model=kwargs.get("model", None), messages=kwargs.get("messages", None), @@ -131,11 +138,21 @@ class GCSBucketLogger(CustomLogger): response_obj=response_obj ) + _spend_log_payload: SpendLogsPayload = get_logging_payload( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + end_user_id=kwargs.get("end_user_id", None), + ) + gcs_payload: GCSBucketPayload = GCSBucketPayload( request_kwargs=request_kwargs, response_obj=response_dict, start_time=start_time, end_time=end_time, + spend_log_metadata=_spend_log_payload["metadata"], + response_cost=kwargs.get("response_cost", None), ) return gcs_payload