litellmlope / litellm /tests /test_logging.py
ka1kuk's picture
Upload 235 files
7db0ae4 verified
raw
history blame
13.7 kB
# #### What this tests ####
# # This tests error logging (with custom user functions) for the raw `completion` + `embedding` endpoints
# # Test Scenarios (test across completion, streaming, embedding)
# ## 1: Pre-API-Call
# ## 2: Post-API-Call
# ## 3: On LiteLLM Call success
# ## 4: On LiteLLM Call failure
# import sys, os, io
# import traceback, logging
# import pytest
# import dotenv
# dotenv.load_dotenv()
# # Create logger
# logger = logging.getLogger(__name__)
# logger.setLevel(logging.DEBUG)
# # Create a stream handler
# stream_handler = logging.StreamHandler(sys.stdout)
# logger.addHandler(stream_handler)
# # Create a function to log information
# def logger_fn(message):
# logger.info(message)
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # Adds the parent directory to the system path
# import litellm
# from litellm import embedding, completion
# from openai.error import AuthenticationError
# litellm.set_verbose = True
# score = 0
# user_message = "Hello, how are you?"
# messages = [{"content": user_message, "role": "user"}]
# # 1. On Call Success
# # normal completion
# # test on openai completion call
# def test_logging_success_completion():
# global score
# try:
# # Redirect stdout
# old_stdout = sys.stdout
# sys.stdout = new_stdout = io.StringIO()
# response = completion(model="gpt-3.5-turbo", messages=messages)
# # Restore stdout
# sys.stdout = old_stdout
# output = new_stdout.getvalue().strip()
# if "Logging Details Pre-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details Post-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details LiteLLM-Success Call" not in output:
# raise Exception("Required log message not found!")
# score += 1
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# pass
# # ## test on non-openai completion call
# # def test_logging_success_completion_non_openai():
# # global score
# # try:
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # response = completion(model="claude-instant-1", messages=messages)
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Success Call" not in output:
# # raise Exception("Required log message not found!")
# # score += 1
# # except Exception as e:
# # pytest.fail(f"Error occurred: {e}")
# # pass
# # streaming completion
# ## test on openai completion call
# def test_logging_success_streaming_openai():
# global score
# try:
# # litellm.set_verbose = False
# def custom_callback(
# kwargs, # kwargs to completion
# completion_response, # response from completion
# start_time, end_time # start/end time
# ):
# if "complete_streaming_response" in kwargs:
# print(f"Complete Streaming Response: {kwargs['complete_streaming_response']}")
# # Assign the custom callback function
# litellm.success_callback = [custom_callback]
# # Redirect stdout
# old_stdout = sys.stdout
# sys.stdout = new_stdout = io.StringIO()
# response = completion(model="gpt-3.5-turbo", messages=messages, stream=True)
# for chunk in response:
# pass
# # Restore stdout
# sys.stdout = old_stdout
# output = new_stdout.getvalue().strip()
# if "Logging Details Pre-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details Post-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details LiteLLM-Success Call" not in output:
# raise Exception("Required log message not found!")
# elif "Complete Streaming Response:" not in output:
# raise Exception("Required log message not found!")
# score += 1
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# pass
# # test_logging_success_streaming_openai()
# ## test on non-openai completion call
# def test_logging_success_streaming_non_openai():
# global score
# try:
# # litellm.set_verbose = False
# def custom_callback(
# kwargs, # kwargs to completion
# completion_response, # response from completion
# start_time, end_time # start/end time
# ):
# # print(f"streaming response: {completion_response}")
# if "complete_streaming_response" in kwargs:
# print(f"Complete Streaming Response: {kwargs['complete_streaming_response']}")
# # Assign the custom callback function
# litellm.success_callback = [custom_callback]
# # Redirect stdout
# old_stdout = sys.stdout
# sys.stdout = new_stdout = io.StringIO()
# response = completion(model="claude-instant-1", messages=messages, stream=True)
# for idx, chunk in enumerate(response):
# pass
# # Restore stdout
# sys.stdout = old_stdout
# output = new_stdout.getvalue().strip()
# if "Logging Details Pre-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details Post-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details LiteLLM-Success Call" not in output:
# raise Exception("Required log message not found!")
# elif "Complete Streaming Response:" not in output:
# raise Exception(f"Required log message not found! {output}")
# score += 1
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# pass
# # test_logging_success_streaming_non_openai()
# # embedding
# def test_logging_success_embedding_openai():
# try:
# # Redirect stdout
# old_stdout = sys.stdout
# sys.stdout = new_stdout = io.StringIO()
# response = embedding(model="text-embedding-ada-002", input=["good morning from litellm"])
# # Restore stdout
# sys.stdout = old_stdout
# output = new_stdout.getvalue().strip()
# if "Logging Details Pre-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details Post-API Call" not in output:
# raise Exception("Required log message not found!")
# elif "Logging Details LiteLLM-Success Call" not in output:
# raise Exception("Required log message not found!")
# except Exception as e:
# pytest.fail(f"Error occurred: {e}")
# # ## 2. On LiteLLM Call failure
# # ## TEST BAD KEY
# # # normal completion
# # ## test on openai completion call
# # try:
# # temporary_oai_key = os.environ["OPENAI_API_KEY"]
# # os.environ["OPENAI_API_KEY"] = "bad-key"
# # temporary_anthropic_key = os.environ["ANTHROPIC_API_KEY"]
# # os.environ["ANTHROPIC_API_KEY"] = "bad-key"
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # try:
# # response = completion(model="gpt-3.5-turbo", messages=messages)
# # except AuthenticationError:
# # print(f"raised auth error")
# # pass
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # print(output)
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Failure Call" not in output:
# # raise Exception("Required log message not found!")
# # os.environ["OPENAI_API_KEY"] = temporary_oai_key
# # os.environ["ANTHROPIC_API_KEY"] = temporary_anthropic_key
# # score += 1
# # except Exception as e:
# # print(f"exception type: {type(e).__name__}")
# # pytest.fail(f"Error occurred: {e}")
# # pass
# # ## test on non-openai completion call
# # try:
# # temporary_oai_key = os.environ["OPENAI_API_KEY"]
# # os.environ["OPENAI_API_KEY"] = "bad-key"
# # temporary_anthropic_key = os.environ["ANTHROPIC_API_KEY"]
# # os.environ["ANTHROPIC_API_KEY"] = "bad-key"
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # try:
# # response = completion(model="claude-instant-1", messages=messages)
# # except AuthenticationError:
# # pass
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Failure Call" not in output:
# # raise Exception("Required log message not found!")
# # os.environ["OPENAI_API_KEY"] = temporary_oai_key
# # os.environ["ANTHROPIC_API_KEY"] = temporary_anthropic_key
# # score += 1
# # except Exception as e:
# # print(f"exception type: {type(e).__name__}")
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # print(output)
# # pytest.fail(f"Error occurred: {e}")
# # # streaming completion
# # ## test on openai completion call
# # try:
# # temporary_oai_key = os.environ["OPENAI_API_KEY"]
# # os.environ["OPENAI_API_KEY"] = "bad-key"
# # temporary_anthropic_key = os.environ["ANTHROPIC_API_KEY"]
# # os.environ["ANTHROPIC_API_KEY"] = "bad-key"
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # try:
# # response = completion(model="gpt-3.5-turbo", messages=messages)
# # except AuthenticationError:
# # pass
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # print(output)
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Failure Call" not in output:
# # raise Exception("Required log message not found!")
# # os.environ["OPENAI_API_KEY"] = temporary_oai_key
# # os.environ["ANTHROPIC_API_KEY"] = temporary_anthropic_key
# # score += 1
# # except Exception as e:
# # print(f"exception type: {type(e).__name__}")
# # pytest.fail(f"Error occurred: {e}")
# # ## test on non-openai completion call
# # try:
# # temporary_oai_key = os.environ["OPENAI_API_KEY"]
# # os.environ["OPENAI_API_KEY"] = "bad-key"
# # temporary_anthropic_key = os.environ["ANTHROPIC_API_KEY"]
# # os.environ["ANTHROPIC_API_KEY"] = "bad-key"
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # try:
# # response = completion(model="claude-instant-1", messages=messages)
# # except AuthenticationError:
# # pass
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # print(output)
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Failure Call" not in output:
# # raise Exception("Required log message not found!")
# # score += 1
# # except Exception as e:
# # print(f"exception type: {type(e).__name__}")
# # pytest.fail(f"Error occurred: {e}")
# # # embedding
# # try:
# # temporary_oai_key = os.environ["OPENAI_API_KEY"]
# # os.environ["OPENAI_API_KEY"] = "bad-key"
# # temporary_anthropic_key = os.environ["ANTHROPIC_API_KEY"]
# # os.environ["ANTHROPIC_API_KEY"] = "bad-key"
# # # Redirect stdout
# # old_stdout = sys.stdout
# # sys.stdout = new_stdout = io.StringIO()
# # try:
# # response = embedding(model="text-embedding-ada-002", input=["good morning from litellm"])
# # except AuthenticationError:
# # pass
# # # Restore stdout
# # sys.stdout = old_stdout
# # output = new_stdout.getvalue().strip()
# # print(output)
# # if "Logging Details Pre-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details Post-API Call" not in output:
# # raise Exception("Required log message not found!")
# # elif "Logging Details LiteLLM-Failure Call" not in output:
# # raise Exception("Required log message not found!")
# # except Exception as e:
# # print(f"exception type: {type(e).__name__}")
# # pytest.fail(f"Error occurred: {e}")