mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-24 10:14:26 +00:00
Merge 98817b4cea into b82af5b826
This commit is contained in:
Lei (Raymond) Chi
GitHub
commit
afd705329c
2 changed files with 198 additions and 68 deletions
|
|
@ -1,5 +1,6 @@
|
|||
#### What this does ####
|
||||
# On success, logs events to Promptlayer
|
||||
import json
|
||||
import os
|
||||
import traceback
|
||||
|
||||
|
|
@ -13,79 +14,127 @@ class PromptLayerLogger:
|
|||
def __init__(self):
|
||||
# Instance variables
|
||||
self.key = os.getenv("PROMPTLAYER_API_KEY")
|
||||
self._streaming_content = []
|
||||
|
||||
def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose):
|
||||
# Method definition
|
||||
try:
|
||||
new_kwargs = {}
|
||||
new_kwargs["model"] = kwargs["model"]
|
||||
new_kwargs["messages"] = kwargs["messages"]
|
||||
|
||||
# add kwargs["optional_params"] to new_kwargs
|
||||
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}"
|
||||
)
|
||||
|
||||
# python-openai >= 1.0.0 returns Pydantic objects instead of jsons
|
||||
# Convert pydantic object to dict if necessary
|
||||
if isinstance(response_obj, BaseModel):
|
||||
response_obj = response_obj.model_dump()
|
||||
|
||||
request_response = litellm.module_level_client.post(
|
||||
"https://api.promptlayer.com/rest/track-request",
|
||||
json={
|
||||
"function_name": "openai.ChatCompletion.create",
|
||||
"kwargs": new_kwargs,
|
||||
"tags": tags,
|
||||
"request_response": dict(response_obj),
|
||||
"request_start_time": int(start_time.timestamp()),
|
||||
"request_end_time": int(end_time.timestamp()),
|
||||
"api_key": self.key,
|
||||
# Optional params for PromptLayer
|
||||
# "prompt_id": "<PROMPT ID>",
|
||||
# "prompt_input_variables": "<Dictionary of variables for prompt>",
|
||||
# "prompt_version":1,
|
||||
},
|
||||
)
|
||||
# Handle metadata and tags
|
||||
tags = []
|
||||
metadata = {}
|
||||
total_cost = 0
|
||||
|
||||
response_json = request_response.json()
|
||||
if not request_response.json().get("success", False):
|
||||
raise Exception("Promptlayer did not successfully log the response!")
|
||||
if kwargs.get("litellm_params"):
|
||||
metadata_dict = kwargs["litellm_params"].get("metadata", {})
|
||||
if isinstance(metadata_dict, dict):
|
||||
if "pl_tags" in metadata_dict:
|
||||
tags = metadata_dict["pl_tags"]
|
||||
metadata = {
|
||||
k: v for k, v in metadata_dict.items() if k != "pl_tags"
|
||||
}
|
||||
# Get cost from hidden_params if it exists
|
||||
if "hidden_params" in metadata:
|
||||
total_cost = metadata["hidden_params"].get("response_cost", 0)
|
||||
metadata["hidden_params"] = json.dumps(
|
||||
metadata["hidden_params"]
|
||||
)
|
||||
|
||||
# Handle streaming vs non-streaming responses
|
||||
if kwargs.get("stream", False):
|
||||
for choice in response_obj.get("choices", []):
|
||||
delta = choice.get("delta", {})
|
||||
content = delta.get("content", "")
|
||||
if content:
|
||||
self._streaming_content.append(content)
|
||||
|
||||
is_final_chunk = (
|
||||
response_obj.get("choices")
|
||||
and response_obj.get("choices")[0].get("finish_reason") == "stop"
|
||||
)
|
||||
|
||||
if not is_final_chunk:
|
||||
return None
|
||||
|
||||
full_content = "".join(self._streaming_content)
|
||||
self._streaming_content = [] # Reset for next stream
|
||||
output_messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": full_content}],
|
||||
}
|
||||
]
|
||||
else:
|
||||
output_content = (
|
||||
response_obj.get("choices", [{}])[0]
|
||||
.get("message", {})
|
||||
.get("content", "")
|
||||
)
|
||||
output_messages = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": output_content}],
|
||||
}
|
||||
]
|
||||
|
||||
# Format input messages
|
||||
input_messages = [
|
||||
{
|
||||
"role": msg["role"],
|
||||
"content": [{"type": "text", "text": msg["content"]}],
|
||||
}
|
||||
for msg in kwargs["messages"]
|
||||
]
|
||||
|
||||
# Construct request payload
|
||||
payload = {
|
||||
"provider": "openai",
|
||||
"model": kwargs["model"],
|
||||
"input": {"type": "chat", "messages": input_messages},
|
||||
"output": {"type": "chat", "messages": output_messages},
|
||||
"request_start_time": start_time.timestamp(),
|
||||
"request_end_time": end_time.timestamp(),
|
||||
"parameters": kwargs.get("optional_params", {}),
|
||||
"prompt_name": kwargs.get("prompt_name", ""),
|
||||
"prompt_version_number": kwargs.get("prompt_version_number", 1),
|
||||
"prompt_input_variables": kwargs.get("prompt_input_variables", {}),
|
||||
"input_tokens": response_obj.get("usage", {}).get("prompt_tokens", 0),
|
||||
"output_tokens": response_obj.get("usage", {}).get(
|
||||
"completion_tokens", 0
|
||||
),
|
||||
"function_name": "openai.chat.completions.create",
|
||||
"tags": tags,
|
||||
"metadata": metadata,
|
||||
"price": total_cost,
|
||||
"score": 0,
|
||||
}
|
||||
|
||||
print_verbose(
|
||||
f"Prompt Layer Logging: success - final response object: {request_response.text}"
|
||||
f"Prompt Layer Logging - Sending payload: {json.dumps(payload, indent=2)}"
|
||||
)
|
||||
|
||||
if "request_id" in response_json:
|
||||
if metadata:
|
||||
response = litellm.module_level_client.post(
|
||||
"https://api.promptlayer.com/rest/track-metadata",
|
||||
json={
|
||||
"request_id": response_json["request_id"],
|
||||
"api_key": self.key,
|
||||
"metadata": metadata,
|
||||
},
|
||||
)
|
||||
print_verbose(
|
||||
f"Prompt Layer Logging: success - metadata post response object: {response.text}"
|
||||
)
|
||||
request_response = litellm.module_level_client.post(
|
||||
"https://api.promptlayer.com/log-request",
|
||||
json=payload,
|
||||
headers={"X-API-KEY": self.key, "Content-Type": "application/json"},
|
||||
)
|
||||
|
||||
except Exception:
|
||||
request_response.raise_for_status()
|
||||
response_json = request_response.json()
|
||||
|
||||
if "id" in response_json:
|
||||
print_verbose(
|
||||
f"Prompt Layer Logging: success - request ID: {response_json['id']}"
|
||||
)
|
||||
return response_json
|
||||
|
||||
print_verbose(
|
||||
f"PromptLayer API response missing 'id' field: {response_json}"
|
||||
)
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print_verbose(f"error: Prompt Layer Error - {traceback.format_exc()}")
|
||||
pass
|
||||
return None
|
||||
|
|
@ -4,12 +4,20 @@ import io
|
|||
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
from litellm import completion
|
||||
import litellm
|
||||
|
||||
import pytest
|
||||
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from openai.types.chat import ChatCompletionMessage
|
||||
from openai.types.chat.chat_completion import ChatCompletion, Choice
|
||||
from respx import MockRouter
|
||||
|
||||
import litellm
|
||||
from litellm import completion
|
||||
from litellm.integrations.prompt_layer import PromptLayerLogger
|
||||
|
||||
# def test_promptlayer_logging():
|
||||
# try:
|
||||
|
|
@ -114,3 +122,76 @@ def test_promptlayer_logging_with_metadata_tags():
|
|||
# print(e)
|
||||
|
||||
# test_chat_openai()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.respx
|
||||
async def test_promptlayer_logging_with_mocked_request(respx_mock: MockRouter):
|
||||
promptlayer_logger = PromptLayerLogger()
|
||||
|
||||
mock_response = AsyncMock()
|
||||
obj = ChatCompletion(
|
||||
id="foo",
|
||||
model="gpt-4",
|
||||
object="chat.completion",
|
||||
choices=[
|
||||
Choice(
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
message=ChatCompletionMessage(
|
||||
content="Hello world!",
|
||||
role="assistant",
|
||||
),
|
||||
)
|
||||
],
|
||||
created=int(datetime.now().timestamp()),
|
||||
)
|
||||
litellm.set_verbose = True
|
||||
|
||||
mock_request = respx_mock.post(url__regex=r".*/chat/completions.*").mock(
|
||||
return_value=httpx.Response(200, json=obj.model_dump(mode="json"))
|
||||
)
|
||||
|
||||
mock_promptlayer_response = respx_mock.post(
|
||||
"https://api.promptlayer.com/log-request"
|
||||
).mock(return_value=httpx.Response(200, json={"id": "mock_promptlayer_id"}))
|
||||
|
||||
response = completion(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "Hello, can you provide a response?"}],
|
||||
temperature=0.2,
|
||||
max_tokens=20,
|
||||
metadata={"model": "ai21", "pl_tags": ["env:dev"]},
|
||||
)
|
||||
|
||||
status_code = promptlayer_logger.log_event(
|
||||
kwargs={
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello, can you provide a response?"}
|
||||
],
|
||||
},
|
||||
response_obj=response,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
print_verbose=print,
|
||||
)
|
||||
|
||||
respx_mock.assert_all_called()
|
||||
|
||||
for call in mock_request.calls:
|
||||
print(call)
|
||||
print(call.request.content)
|
||||
|
||||
json_body = json.loads(call.request.content)
|
||||
|
||||
print(json_body)
|
||||
|
||||
for call in mock_promptlayer_response.calls:
|
||||
print(call)
|
||||
print(call.request.content)
|
||||
|
||||
json_body = json.loads(call.request.content)
|
||||
print(json_body)
|
||||
|
||||
assert status_code == {"id": "mock_promptlayer_id"}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue