support for later embeddings

This commit is contained in:
Vince Lwt 2023-08-22 15:42:56 +02:00
parent 342b83544f
commit 0cb4816ab2
3 changed files with 61 additions and 22 deletions

View file

@ -12,8 +12,10 @@ dotenv.load_dotenv() # Loading env variables using dotenv
# convert to {completion: xx, tokens: xx} # convert to {completion: xx, tokens: xx}
def parse_usage(usage): def parse_usage(usage):
return { return {
"completion": usage["completion_tokens"], "completion":
"prompt": usage["prompt_tokens"], usage["completion_tokens"] if "completion_tokens" in usage else 0,
"prompt":
usage["prompt_tokens"] if "prompt_tokens" in usage else 0,
} }
@ -23,6 +25,10 @@ def parse_messages(input):
return None return None
def clean_message(message): def clean_message(message):
#if is strin, return as is
if isinstance(message, str):
return message
if "message" in message: if "message" in message:
return clean_message(message["message"]) return clean_message(message["message"])
@ -51,10 +57,11 @@ class LLMonitorLogger:
def log_event( def log_event(
self, self,
type, type,
event,
run_id, run_id,
model, model,
print_verbose, print_verbose,
messages=None, input=None,
user_id=None, user_id=None,
response_obj=None, response_obj=None,
start_time=datetime.datetime.now(), start_time=datetime.datetime.now(),
@ -67,8 +74,11 @@ class LLMonitorLogger:
f"LLMonitor Logging - Logging request for model {model}") f"LLMonitor Logging - Logging request for model {model}")
if response_obj: if response_obj:
usage = parse_usage(response_obj['usage']) print(response_obj)
output = response_obj['choices'] usage = parse_usage(
response_obj['usage']) if 'usage' in response_obj else None
output = response_obj[
'choices'] if 'choices' in response_obj else None
else: else:
usage = None usage = None
output = None output = None
@ -80,19 +90,19 @@ class LLMonitorLogger:
error_obj = None error_obj = None
data = [{ data = [{
"type": "llm", "type": type,
"name": model, "name": model,
"runId": run_id, "runId": run_id,
"app": self.app_id, "app": self.app_id,
'event': 'start', 'event': 'start',
"timestamp": start_time.isoformat(), "timestamp": start_time.isoformat(),
"userId": user_id, "userId": user_id,
"input": parse_messages(messages), "input": parse_messages(input),
}, { }, {
"type": "llm", "type": type,
"runId": run_id, "runId": run_id,
"app": self.app_id, "app": self.app_id,
"event": type, "event": event,
"error": error_obj, "error": error_obj,
"timestamp": end_time.isoformat(), "timestamp": end_time.isoformat(),
"userId": user_id, "userId": user_id,
@ -100,6 +110,8 @@ class LLMonitorLogger:
"tokensUsage": usage, "tokensUsage": usage,
}] }]
print(data)
# print_verbose(f"LLMonitor Logging - final data object: {data}") # print_verbose(f"LLMonitor Logging - final data object: {data}")
response = requests.post( response = requests.post(

View file

@ -6,7 +6,7 @@ import os
sys.path.insert(0, os.path.abspath('../..')) sys.path.insert(0, os.path.abspath('../..'))
from litellm import completion from litellm import completion, embedding
import litellm import litellm
litellm.success_callback = ["llmonitor"] litellm.success_callback = ["llmonitor"]
@ -15,10 +15,23 @@ litellm.failure_callback = ["llmonitor"]
litellm.set_verbose = True litellm.set_verbose = True
# openai call # openai call
first_success_test = completion(model="gpt-3.5-turbo", # first_success_test = completion(model="gpt-3.5-turbo",
messages=[{ # messages=[{
"role": "user", # "role": "user",
"content": "Hi 👋 - i'm openai" # "content": "Hi 👋 - i'm openai"
}]) # }])
print(first_success_test) # print(first_success_test)
def test_embedding_openai():
try:
response = embedding(model="text-embedding-ada-002", input=['test'])
# Add any assertions here to check the response
print(f"response: {str(response)[:50]}")
except Exception as e:
print(e)
# pytest.fail(f"Error occurred: {e}")
test_embedding_openai()

View file

@ -842,14 +842,21 @@ def handle_failure(exception, traceback_exception, start_time, end_time, args,
print_verbose=print_verbose, print_verbose=print_verbose,
) )
elif callback == "llmonitor": elif callback == "llmonitor":
print_verbose("reaches llmonitor for logging!") print_verbose("reaches llmonitor for logging error!")
model = args[0] if len(args) > 0 else kwargs["model"] model = args[0] if len(args) > 0 else kwargs["model"]
messages = args[1] if len(args) > 1 else kwargs["messages"]
input = args[1] if len(args) > 1 else kwargs.get(
"messages", kwargs.get("input", None))
type = 'embed' if 'input' in kwargs else 'llm'
llmonitorLogger.log_event( llmonitorLogger.log_event(
type="error", type=type,
event="error",
user_id=litellm._thread_context.user, user_id=litellm._thread_context.user,
model=model, model=model,
input=input,
error=traceback_exception, error=traceback_exception,
run_id=kwargs["litellm_call_id"], run_id=kwargs["litellm_call_id"],
start_time=start_time, start_time=start_time,
@ -969,11 +976,18 @@ def handle_success(args, kwargs, result, start_time, end_time):
elif callback == "llmonitor": elif callback == "llmonitor":
print_verbose("reaches llmonitor for logging!") print_verbose("reaches llmonitor for logging!")
model = args[0] if len(args) > 0 else kwargs["model"] model = args[0] if len(args) > 0 else kwargs["model"]
messages = args[1] if len(args) > 1 else kwargs["messages"]
input = args[1] if len(args) > 1 else kwargs.get(
"messages", kwargs.get("input", None))
#if contains input, it's 'embedding', otherwise 'llm'
type = 'embed' if 'input' in kwargs else 'llm'
llmonitorLogger.log_event( llmonitorLogger.log_event(
type="end", type=type,
event="end",
model=model, model=model,
messages=messages, input=input,
user_id=litellm._thread_context.user, user_id=litellm._thread_context.user,
response_obj=result, response_obj=result,
start_time=start_time, start_time=start_time,