adding helicone integration

This commit is contained in:
Krrish Dholakia 2023-08-02 22:25:08 -07:00
parent 13e800493c
commit aa598d506f
8 changed files with 74 additions and 8 deletions

View file

@ -0,0 +1,49 @@
#### What this does ####
# On success, logs events to Helicone
import dotenv, os
import requests
dotenv.load_dotenv() # Loading env variables using dotenv
import traceback
class HeliconeLogger:
# Class variables or attributes
helicone_model_list = ["gpt", "claude"]
def __init__(self):
# Instance variables
self.provider_url = "https://api.openai.com/v1"
self.key = os.getenv('HELICONE_API_KEY')
def log_success(self, model, messages, response_obj, start_time, end_time):
# Method definition
try:
model = model if any(accepted_model in model for accepted_model in self.helicone_model_list) else "gpt-3.5-turbo"
provider_request = {"model": model, "messages": messages}
providerResponse = {
"json": response_obj,
"headers": {"openai-version": "2020-10-01"},
"status": 200
}
# Code to be executed
url = "https://api.hconeai.com/oai/v1/log"
headers = {
'Authorization': f'Bearer {self.key}',
'Content-Type': 'application/json'
}
start_time_seconds = int(start_time.timestamp())
start_time_milliseconds = int((start_time.timestamp() - start_time_seconds) * 1000)
end_time_seconds = int(end_time.timestamp())
end_time_milliseconds = int((end_time.timestamp() - end_time_seconds) * 1000)
data = {
"providerRequest": {"url": self.provider_url, "json": provider_request, "meta": {"Helicone-Auth": f"Bearer {self.key}"}},
"providerResponse": providerResponse,
"timing": {"startTime": {"seconds": start_time_seconds, "milliseconds": start_time_milliseconds}, "endTime": {"seconds": end_time_seconds, "milliseconds": end_time_milliseconds}} # {"seconds": .., "milliseconds": ..}
}
response = requests.post(url, headers=headers, json=data)
# if response.status_code == 200:
# print("Success!")
# else:
# print("Request was not successful. Status Code:", response.status_code)
except:
# traceback.print_exc()
pass

View file

@ -9,10 +9,10 @@ sys.path.insert(0, os.path.abspath('../..')) # Adds the parent directory to the
import litellm import litellm
from litellm import embedding, completion from litellm import embedding, completion
litellm.success_callback = ["posthog"] litellm.success_callback = ["posthog", "helicone"]
litellm.failure_callback = ["slack", "sentry", "posthog"] litellm.failure_callback = ["slack", "sentry", "posthog"]
# litellm.set_verbose = True litellm.set_verbose = True
def logger_fn(model_call_object: dict): def logger_fn(model_call_object: dict):
# print(f"model call details: {model_call_object}") # print(f"model call details: {model_call_object}")
@ -23,11 +23,14 @@ messages = [{ "content": user_message,"role": "user"}]
def test_completion_openai(): def test_completion_openai():
try: try:
print("running query")
response = completion(model="gpt-3.5-turbo", messages=messages, logger_fn=logger_fn) response = completion(model="gpt-3.5-turbo", messages=messages, logger_fn=logger_fn)
print(f"response: {response}")
# Add any assertions here to check the response # Add any assertions here to check the response
except Exception as e: except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
test_completion_openai()
def test_completion_non_openai(): def test_completion_non_openai():
try: try:
response = completion(model="claude-instant-1", messages=messages, logger_fn=logger_fn) response = completion(model="claude-instant-1", messages=messages, logger_fn=logger_fn)

View file

@ -2,6 +2,7 @@ import dotenv, json, traceback, threading
import subprocess, os import subprocess, os
import litellm, openai import litellm, openai
import random, uuid, requests import random, uuid, requests
import datetime
from openai.error import AuthenticationError, InvalidRequestError, RateLimitError, ServiceUnavailableError, OpenAIError from openai.error import AuthenticationError, InvalidRequestError, RateLimitError, ServiceUnavailableError, OpenAIError
####### ENVIRONMENT VARIABLES ################### ####### ENVIRONMENT VARIABLES ###################
dotenv.load_dotenv() # Loading env variables using dotenv dotenv.load_dotenv() # Loading env variables using dotenv
@ -11,6 +12,7 @@ add_breadcrumb = None
posthog = None posthog = None
slack_app = None slack_app = None
alerts_channel = None alerts_channel = None
heliconeLogger = None
callback_list = [] callback_list = []
user_logger_fn = None user_logger_fn = None
additional_details = {} additional_details = {}
@ -68,7 +70,7 @@ def client(original_function):
global callback_list, add_breadcrumb global callback_list, add_breadcrumb
if (len(litellm.success_callback) > 0 or len(litellm.failure_callback) > 0) and len(callback_list) == 0: if (len(litellm.success_callback) > 0 or len(litellm.failure_callback) > 0) and len(callback_list) == 0:
callback_list = list(set(litellm.success_callback + litellm.failure_callback)) callback_list = list(set(litellm.success_callback + litellm.failure_callback))
set_callbacks(callback_list=callback_list) set_callbacks(callback_list=callback_list,)
if add_breadcrumb: if add_breadcrumb:
add_breadcrumb( add_breadcrumb(
category="litellm.llm_call", category="litellm.llm_call",
@ -83,9 +85,11 @@ def client(original_function):
try: try:
function_setup(args, kwargs) function_setup(args, kwargs)
## MODEL CALL ## MODEL CALL
start_time = datetime.datetime.now()
result = original_function(*args, **kwargs) result = original_function(*args, **kwargs)
end_time = datetime.datetime.now()
## LOG SUCCESS ## LOG SUCCESS
my_thread = threading.Thread(target=handle_success, args=(args, kwargs)) # don't interrupt execution of main thread my_thread = threading.Thread(target=handle_success, args=(args, kwargs, result, start_time, end_time)) # don't interrupt execution of main thread
my_thread.start() my_thread.start()
return result return result
except Exception as e: except Exception as e:
@ -97,7 +101,7 @@ def client(original_function):
####### HELPER FUNCTIONS ################ ####### HELPER FUNCTIONS ################
def set_callbacks(callback_list): def set_callbacks(callback_list):
global sentry_sdk_instance, capture_exception, add_breadcrumb, posthog, slack_app, alerts_channel global sentry_sdk_instance, capture_exception, add_breadcrumb, posthog, slack_app, alerts_channel, heliconeLogger
try: try:
for callback in callback_list: for callback in callback_list:
if callback == "sentry": if callback == "sentry":
@ -134,6 +138,10 @@ def set_callbacks(callback_list):
) )
alerts_channel = os.environ["SLACK_API_CHANNEL"] alerts_channel = os.environ["SLACK_API_CHANNEL"]
print_verbose(f"Initialized Slack App: {slack_app}") print_verbose(f"Initialized Slack App: {slack_app}")
elif callback == "helicone":
from .integrations.helicone import HeliconeLogger
heliconeLogger = HeliconeLogger()
except: except:
pass pass
@ -200,7 +208,8 @@ def handle_failure(exception, traceback_exception, args, kwargs):
except: except:
pass pass
def handle_success(*args, **kwargs): def handle_success(args, kwargs, result, start_time, end_time):
global heliconeLogger
try: try:
success_handler = additional_details.pop("success_handler", None) success_handler = additional_details.pop("success_handler", None)
failure_handler = additional_details.pop("failure_handler", None) failure_handler = additional_details.pop("failure_handler", None)
@ -223,6 +232,11 @@ def handle_success(*args, **kwargs):
for detail in additional_details: for detail in additional_details:
slack_msg += f"{detail}: {additional_details[detail]}\n" slack_msg += f"{detail}: {additional_details[detail]}\n"
slack_app.client.chat_postMessage(channel=alerts_channel, text=slack_msg) slack_app.client.chat_postMessage(channel=alerts_channel, text=slack_msg)
elif callback == "helicone":
print_verbose("reaches helicone for logging!")
model = args[0] if len(args) > 0 else kwargs["model"]
messages = args[1] if len(args) > 1 else kwargs["messages"]
heliconeLogger.log_success(model=model, messages=messages, response_obj=result, start_time=start_time, end_time=end_time)
except: except:
pass pass

View file

@ -2,7 +2,7 @@ from setuptools import setup, find_packages
setup( setup(
name='litellm', name='litellm',
version='0.1.220', version='0.1.221',
description='Library to easily interface with LLM API providers', description='Library to easily interface with LLM API providers',
author='BerriAI', author='BerriAI',
packages=[ packages=[