Updating PromptLayer callback to support tags + metadata

This commit is contained in:
Samy Chouiti 2024-01-18 09:51:06 +01:00
parent 85b5395692
commit 370e7777e3

View file

@ -2,12 +2,10 @@
# On success, logs events to Promptlayer
import dotenv, os
import requests
import requests
dotenv.load_dotenv() # Loading env variables using dotenv
import traceback
class PromptLayerLogger:
# Class variables or attributes
def __init__(self):
@ -25,6 +23,16 @@ class PromptLayerLogger:
for optional_param in kwargs["optional_params"]:
new_kwargs[optional_param] = kwargs["optional_params"][optional_param]
# Extract PromptLayer tags from metadata, if such exists
tags = []
metadata = {}
if "metadata" in kwargs["litellm_params"]:
if "pl_tags" in kwargs["litellm_params"]["metadata"]:
tags = kwargs["litellm_params"]["metadata"]["pl_tags"]
# Remove "pl_tags" from metadata
metadata = {k:v for k, v in kwargs["litellm_params"]["metadata"].items() if k != "pl_tags"}
print_verbose(
f"Prompt Layer Logging - Enters logging function for model kwargs: {new_kwargs}\n, response: {response_obj}"
)
@ -34,7 +42,7 @@ class PromptLayerLogger:
json={
"function_name": "openai.ChatCompletion.create",
"kwargs": new_kwargs,
"tags": ["hello", "world"],
"tags": tags,
"request_response": dict(response_obj),
"request_start_time": int(start_time.timestamp()),
"request_end_time": int(end_time.timestamp()),
@ -53,14 +61,13 @@ class PromptLayerLogger:
raise Exception("Promptlayer did not successfully log the response!")
if "request_id" in response_json:
print(kwargs["litellm_params"]["metadata"])
if kwargs["litellm_params"]["metadata"] is not None:
if metadata:
response = requests.post(
"https://api.promptlayer.com/rest/track-metadata",
json={
"request_id": response_json["request_id"],
"api_key": self.key,
"metadata": kwargs["litellm_params"]["metadata"],
"metadata": metadata,
},
)
print_verbose(