diff --git a/litellm/integrations/gcs_bucket.py b/litellm/integrations/gcs_bucket.py index 3fb778e2428..46f55f8f01c 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,17 @@ class GCSBucketLogger(CustomLogger): kwargs, response_obj, start_time_str, end_time_str ) - object_name = response_obj["id"] + json_logged_payload = json.dumps(logging_payload) + + # Get the current date + current_date = datetime.now().strftime("%Y-%m-%d") + + # Modify the object_name to include the date-based folder + object_name = f"{current_date}/{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 +129,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 +143,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 diff --git a/litellm/tests/test_gcs_bucket.py b/litellm/tests/test_gcs_bucket.py index c5a6fb76acc..b30978bad5d 100644 --- a/litellm/tests/test_gcs_bucket.py +++ b/litellm/tests/test_gcs_bucket.py @@ -9,6 +9,7 @@ import json import logging import tempfile import uuid +from datetime import datetime import pytest @@ -63,7 +64,7 @@ def load_vertex_ai_credentials(): @pytest.mark.asyncio async def test_basic_gcs_logger(): - load_vertex_ai_credentials() + # load_vertex_ai_credentials() gcs_logger = GCSBucketLogger() print("GCSBucketLogger", gcs_logger) @@ -75,19 +76,66 @@ async def test_basic_gcs_logger(): max_tokens=10, user="ishaan-2", mock_response="Hi!", + metadata={ + "tags": ["model-anthropic-claude-v2.1", "app-ishaan-prod"], + "user_api_key": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b", + "user_api_key_alias": None, + "user_api_end_user_max_budget": None, + "litellm_api_version": "0.0.0", + "global_max_parallel_requests": None, + "user_api_key_user_id": "116544810872468347480", + "user_api_key_org_id": None, + "user_api_key_team_id": None, + "user_api_key_team_alias": None, + "user_api_key_metadata": {}, + "requester_ip_address": "127.0.0.1", + "spend_logs_metadata": {"hello": "world"}, + "headers": { + "content-type": "application/json", + "user-agent": "PostmanRuntime/7.32.3", + "accept": "*/*", + "postman-token": "92300061-eeaa-423b-a420-0b44896ecdc4", + "host": "localhost:4000", + "accept-encoding": "gzip, deflate, br", + "connection": "keep-alive", + "content-length": "163", + }, + "endpoint": "http://localhost:4000/chat/completions", + "model_group": "gpt-3.5-turbo", + "deployment": "azure/chatgpt-v-2", + "model_info": { + "id": "4bad40a1eb6bebd1682800f16f44b9f06c52a6703444c99c7f9f32e9de3693b4", + "db_model": False, + }, + "api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/", + "caching_groups": None, + "raw_request": "\n\nPOST Request Sent from LiteLLM:\ncurl -X POST \\\nhttps://openai-gpt-4-test-v-1.openai.azure.com//openai/ \\\n-H 'Authorization: *****' \\\n-d '{'model': 'chatgpt-v-2', 'messages': [{'role': 'system', 'content': 'you are a helpful assistant.\\n'}, {'role': 'user', 'content': 'bom dia'}], 'stream': False, 'max_tokens': 10, 'user': '116544810872468347480', 'extra_body': {}}'\n", + }, ) print("response", response) await asyncio.sleep(5) - # Check if object landed on GCS - object_from_gcs = await gcs_logger.download_gcs_object(object_name=response.id) - # convert object_from_gcs from bytes to DICT - object_from_gcs = json.loads(object_from_gcs) - print("object_from_gcs", object_from_gcs) + # Get the current date + # Get the current date + current_date = datetime.now().strftime("%Y-%m-%d") - gcs_payload = GCSBucketPayload(**object_from_gcs) + # Modify the object_name to include the date-based folder + object_name = f"{current_date}%2F{response.id}" + + print("object_name", object_name) + + # Check if object landed on GCS + object_from_gcs = await gcs_logger.download_gcs_object(object_name=object_name) + print("object from gcs=", object_from_gcs) + # convert object_from_gcs from bytes to DICT + parsed_data = json.loads(object_from_gcs) + print("object_from_gcs as dict", parsed_data) + + print("type of object_from_gcs", type(parsed_data)) + + gcs_payload = GCSBucketPayload(**parsed_data) print("gcs_payload", gcs_payload) @@ -97,6 +145,19 @@ async def test_basic_gcs_logger(): ] assert gcs_payload["response_obj"]["choices"][0]["message"]["content"] == "Hi!" + assert gcs_payload["response_cost"] > 0.0 + + gcs_payload["spend_log_metadata"] = json.loads(gcs_payload["spend_log_metadata"]) + + assert ( + gcs_payload["spend_log_metadata"]["user_api_key"] + == "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" + ) + assert ( + gcs_payload["spend_log_metadata"]["user_api_key_user_id"] + == "116544810872468347480" + ) + # Delete Object from GCS print("deleting object from GCS") - await gcs_logger.delete_gcs_object(object_name=response.id) + await gcs_logger.delete_gcs_object(object_name=object_name)