From 98817b4cea231de2f8707925d1106200fa2f7513 Mon Sep 17 00:00:00 2001 From: raymond-chii Date: Thu, 30 Jan 2025 12:30:47 -0500 Subject: [PATCH] Update promptlayer integration with test based on review feedback --- litellm/integrations/prompt_layer.py | 175 +++++++++++------- .../test_promptlayer_integration.py | 91 ++++++++- 2 files changed, 198 insertions(+), 68 deletions(-) diff --git a/litellm/integrations/prompt_layer.py b/litellm/integrations/prompt_layer.py index 190b995fa4..56e3fef667 100644 --- a/litellm/integrations/prompt_layer.py +++ b/litellm/integrations/prompt_layer.py @@ -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_input_variables": "", - # "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 \ No newline at end of file diff --git a/tests/local_testing/test_promptlayer_integration.py b/tests/local_testing/test_promptlayer_integration.py index d2e2268e61..6615b4b78e 100644 --- a/tests/local_testing/test_promptlayer_integration.py +++ b/tests/local_testing/test_promptlayer_integration.py @@ -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"}