Updating PromptLayer callback to support tags + metadata

This commit is contained in:
Samy Chouiti 2024-01-18 09:51:06 +01:00
parent 4f503dc736
commit 97bbda01cd

View file

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