feat gcs log user api key metadata

This commit is contained in:
Ishaan Jaff 2024-08-12 16:06:10 -07:00 • committed by Krrish Dholakia
parent c9ea8cdf41
commit 5569e8790a

View file

@ -14,6 +14,7 @@ from litellm.litellm_core_utils.logging_utils import (
)
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import CommonProxyErrors, SpendLogsMetadata, SpendLogsPayload
from litellm.proxy._types import CommonProxyErrors, SpendLogsMetadata, SpendLogsPayload
class RequestKwargs(TypedDict):
@ -29,6 +30,8 @@ class GCSBucketPayload(TypedDict):
end_time: str
response_cost: Optional[float]
spend_log_metadata: str
response_cost: Optional[float]
spend_log_metadata: str
class GCSBucketLogger(CustomLogger):
@ -81,12 +84,7 @@ class GCSBucketLogger(CustomLogger):
)
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']}"
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}",
@ -133,6 +131,10 @@ class GCSBucketLogger(CustomLogger):
get_logging_payload,
)
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),
@ -151,6 +153,14 @@ class GCSBucketLogger(CustomLogger):
end_user_id=kwargs.get("end_user_id", None),
)
_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,
@ -158,6 +168,8 @@ class GCSBucketLogger(CustomLogger):
end_time=end_time,
spend_log_metadata=_spend_log_payload["metadata"],
response_cost=kwargs.get("response_cost", None),
spend_log_metadata=_spend_log_payload["metadata"],
response_cost=kwargs.get("response_cost", None),
)
return gcs_payload