ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat gcs log user api key metadata

Browse Source
This commit is contained in:
Ishaan Jaff 2024-08-12 16:06:10 -07:00 committed by Krrish Dholakia
parent a8644d8a7d
commit 3714b83ee4
1 changed files with 19 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
litellm/integrations/gcs_bucket.py
View File

@ -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