#### 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("Required log message not found!")
        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}")