import openai
import time
from openai import OpenAI
import anthropic
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage

claude_api_key = 'sk-...'
openai_api_key = 'sk-...'
mixtral_api_key = 'sk-...'

def GPT_response(messages_input, model_name):
  max_try_times = 5

  if model_name in ['gpt-4-1106-preview', 'gpt-4', 'gpt-4-32k', 'gpt-3.5-turbo-0301', 'gpt-4-0613', 'gpt-4-32k-0613', 'gpt-3.5-turbo-16k-0613']:
    openai_api_key_name = openai_api_key
    messages = [{"role": "system", "content": "You are a helpful assistant."}] + messages_input
    for iteration_num in range(max_try_times):
      try:
        client = OpenAI(api_key = openai_api_key_name)
        response = client.chat.completions.create(
          model=model_name,
          messages=messages,
          temperature = 0.0,
          top_p=1,
          frequency_penalty=0,
          presence_penalty=0
        )
        break
      except:
        print(f'{model_name} Waiting 60 seconds for API query')
        time.sleep(60)
    return response.choices[0].message.content
  elif model_name in ["claude-3-sonnet-20240229", "claude-3-opus-20240229", "claude-3-haiku-20240307"]:
    claude_api_key_name = claude_api_key

    for iteration_num in range(max_try_times):
      try:
        client = anthropic.Anthropic(
          # defaults to os.environ.get("ANTHROPIC_API_KEY")
          api_key=claude_api_key_name,
        )
        message = client.messages.create(
          model=model_name,  # claude-3-sonnet-20240229, claude-3-opus-20240229, claude-3-haiku-20240307
          max_tokens=1000,
          temperature=0.0,
          system="You are a helpful assistant.",
          messages=messages_input
        )
        break
      except:
        print(f'{model_name} Waiting 60 seconds for API query')
        time.sleep(60)
    return message.content[0].text
  elif model_name in ['open-mixtral-8x7b', "mistral-large-latest"]:
    api_key = mixtral_api_key
    for iteration_num in range(max_try_times):
      try:
        client = MistralClient(api_key=api_key)
        messages = messages_input
        # No streaming
        chat_response = client.chat(
          model=model_name,
          messages=messages,
        )
        break
      except:
        print(f'{model_name} Waiting 60 seconds for API query')
        time.sleep(60)
    return chat_response.choices[0].message.content
