litellm/tests/test_logging.py

# #### 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}")