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