# ##### THESE TESTS CAN ONLY RUN LOCALLY WITH THE OLLAMA SERVER RUNNING ######
# # https://ollama.ai/

# import sys, os
# import traceback
# from dotenv import load_dotenv

# load_dotenv()
# import os

# sys.path.insert(
#     0, os.path.abspath("../..")
# )  # Adds the parent directory to the system path
# import pytest
# import litellm
# from litellm import embedding, completion
# import asyncio


# user_message = "respond in 20 words. who are you?"
# messages = [{"content": user_message, "role": "user"}]


# def test_ollama_streaming():
#     try:
#         litellm.set_verbose = False
#         messages = [{"role": "user", "content": "What is the weather like in Boston?"}]
#         functions = [
#             {
#                 "name": "get_current_weather",
#                 "description": "Get the current weather in a given location",
#                 "parameters": {
#                     "type": "object",
#                     "properties": {
#                         "location": {
#                             "type": "string",
#                             "description": "The city and state, e.g. San Francisco, CA",
#                         },
#                         "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
#                     },
#                     "required": ["location"],
#                 },
#             }
#         ]
#         response = litellm.completion(
#             model="ollama_chat/mistral",
#             messages=messages,
#             functions=functions,
#             stream=True,
#         )
#         for chunk in response:
#             print(f"CHUNK: {chunk}")
#     except Exception as e:
#         print(e)


# # test_ollama_streaming()


# async def test_async_ollama_streaming():
#     try:
#         litellm.set_verbose = True
#         response = await litellm.acompletion(
#             model="ollama_chat/llama2",
#             messages=[{"role": "user", "content": "Hey, how's it going?"}],
#             stream=True,
#         )
#         async for chunk in response:
#             print(f"CHUNK: {chunk}")
#     except Exception as e:
#         print(e)


# # asyncio.run(test_async_ollama_streaming())

# async def test_async_ollama():
#     try:
#         litellm.set_verbose = True
#         response = await litellm.acompletion(
#             model="ollama_chat/llama2",
#             messages=[{"role": "user", "content": "Hey, how's it going?"}],
#         )
#         print("\n response", response)
#     except Exception as e:
#         print(e)


# # asyncio.run(test_async_ollama())


# def test_completion_ollama():
#     try:
#         litellm.set_verbose = True
#         response = completion(
#             model="ollama_chat/mistral",
#             messages=[{"role": "user", "content": "Hey, how's it going?"}],
#             max_tokens=200,
#             request_timeout=10,
#             stream=True,
#         )
#         for chunk in response:
#             print(chunk)
#         print(response)
#     except Exception as e:
#         pytest.fail(f"Error occurred: {e}")


# # test_completion_ollama()


# def test_completion_ollama_function_calling():
#     try:
#         litellm.set_verbose = True
#         messages = [{"role": "user", "content": "What is the weather like in Boston?"}]
#         functions = [
#             {
#                 "name": "get_current_weather",
#                 "description": "Get the current weather in a given location",
#                 "parameters": {
#                     "type": "object",
#                     "properties": {
#                         "location": {
#                             "type": "string",
#                             "description": "The city and state, e.g. San Francisco, CA",
#                         },
#                         "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
#                     },
#                     "required": ["location"],
#                 },
#             }
#         ]
#         response = completion(
#             model="ollama_chat/mistral",
#             messages=messages,
#             functions=functions,
#             max_tokens=200,
#             request_timeout=10,
#         )
#         for chunk in response:
#             print(chunk)
#         print(response)
#     except Exception as e:
#         pytest.fail(f"Error occurred: {e}")


# test_completion_ollama_function_calling()


# async def async_test_completion_ollama_function_calling():
#     try:
#         litellm.set_verbose = True
#         messages = [{"role": "user", "content": "What is the weather like in Boston?"}]
#         functions = [
#             {
#                 "name": "get_current_weather",
#                 "description": "Get the current weather in a given location",
#                 "parameters": {
#                     "type": "object",
#                     "properties": {
#                         "location": {
#                             "type": "string",
#                             "description": "The city and state, e.g. San Francisco, CA",
#                         },
#                         "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
#                     },
#                     "required": ["location"],
#                 },
#             }
#         ]
#         response = await litellm.acompletion(
#             model="ollama/mistral",
#             messages=messages,
#             functions=functions,
#             max_tokens=200,
#             request_timeout=10,
#         )
#         print(response)
#     except Exception as e:
#         pytest.fail(f"Error occurred: {e}")


# # asyncio.run(async_test_completion_ollama_function_calling())


# def test_completion_ollama_with_api_base():
#     try:
#         response = completion(
#             model="ollama/llama2", messages=messages, api_base="http://localhost:11434"
#         )
#         print(response)
#     except Exception as e:
#         pytest.fail(f"Error occurred: {e}")


# # test_completion_ollama_with_api_base()


# def test_completion_ollama_custom_prompt_template():
#     user_message = "what is litellm?"
#     litellm.register_prompt_template(
#         model="ollama/llama2",
#         roles={
#             "system": {"pre_message": "System: "},
#             "user": {"pre_message": "User: "},
#             "assistant": {"pre_message": "Assistant: "},
#         },
#     )
#     messages = [{"content": user_message, "role": "user"}]
#     litellm.set_verbose = True
#     try:
#         response = completion(model="ollama/llama2", messages=messages, stream=True)
#         print(response)
#         for chunk in response:
#             print(chunk)
#             # print(chunk['choices'][0]['delta'])

#     except Exception as e:
#         traceback.print_exc()
#         pytest.fail(f"Error occurred: {e}")


# # test_completion_ollama_custom_prompt_template()


# async def test_completion_ollama_async_stream():
#     user_message = "what is the weather"
#     messages = [{"content": user_message, "role": "user"}]
#     try:
#         response = await litellm.acompletion(
#             model="ollama/llama2",
#             messages=messages,
#             api_base="http://localhost:11434",
#             stream=True,
#         )
#         async for chunk in response:
#             print(chunk["choices"][0]["delta"])

#         print("TEST ASYNC NON Stream")
#         response = await litellm.acompletion(
#             model="ollama/llama2",
#             messages=messages,
#             api_base="http://localhost:11434",
#         )
#         print(response)
#     except Exception as e:
#         pytest.fail(f"Error occurred: {e}")


# # import asyncio
# # asyncio.run(test_completion_ollama_async_stream())


# def prepare_messages_for_chat(text: str) -> list:
#     messages = [
#         {"role": "user", "content": text},
#     ]
#     return messages


# async def ask_question():
#     params = {
#         "messages": prepare_messages_for_chat(
#             "What is litellm? tell me 10 things about it who is sihaan.write an essay"
#         ),
#         "api_base": "http://localhost:11434",
#         "model": "ollama/llama2",
#         "stream": True,
#     }
#     response = await litellm.acompletion(**params)
#     return response


# async def main():
#     response = await ask_question()
#     async for chunk in response:
#         print(chunk)

#     print("test async completion without streaming")
#     response = await litellm.acompletion(
#         model="ollama/llama2",
#         messages=prepare_messages_for_chat("What is litellm? respond in 2 words"),
#     )
#     print("response", response)


# def test_completion_expect_error():
#     # this tests if we can exception map correctly for ollama
#     print("making ollama request")
#     # litellm.set_verbose=True
#     user_message = "what is litellm?"
#     messages = [{"content": user_message, "role": "user"}]
#     try:
#         response = completion(model="ollama/invalid", messages=messages, stream=True)
#         print(response)
#         for chunk in response:
#             print(chunk)
#             # print(chunk['choices'][0]['delta'])

#     except Exception as e:
#         pass
#         pytest.fail(f"Error occurred: {e}")


# # test_completion_expect_error()


# def test_ollama_llava():
#     litellm.set_verbose = True
#     # same params as gpt-4 vision
#     response = completion(
#         model="ollama/llava",
#         messages=[
#             {
#                 "role": "user",
#                 "content": [
#                     {"type": "text", "text": "What is in this picture"},
#                     {
#                         "type": "image_url",
#                         "image_url": {
#                             "url": " "
#                         },
#                     },
#                 ],
#             }
#         ],
#     )
#     print("Response from ollama/llava")
#     print(response)


# # test_ollama_llava()


# # PROCESSED CHUNK PRE CHUNK CREATOR
