import logging
from openai import OpenAI
from openai import AsyncOpenAI

class LLMModel():
    # "https://api.deepseek.com/v1", "https://api.openai.com/v1/"
    def __init__(self, api_key, model="gpt-3.5-turbo", base_url="https://api.openai.com/v1/", use_web_search=False):
        if api_key is None:
            self.api_key = "sk-rifpc-2Gg7xjJ4qrwzWY7hUhZKT3BlbkFJBkz9CHkx9LkVsSciz9Tg"
        else:
            self.api_key = api_key
        self.model = model
        self.client = OpenAI(api_key=self.api_key, base_url=base_url)
        self.client_async = AsyncOpenAI(api_key=self.api_key, base_url=base_url)
        self.use_web_search = use_web_search
    
    def LLM_response(self, prompt, gen_kwargs={}, model=None):
        if model is None:
            model = self.model

        if type(prompt) == str:
            input_messages = [
                {"role": "user", "content": prompt}
                ]
        elif type(prompt) == list:
            input_messages = prompt
        else:
            logging.error("prompt must be a string or a list of messages, current type: ", type(prompt))
            raise ValueError("prompt must be a string or a list of messages")
        
        if self.use_web_search:
            model = "gpt-4o-search-preview"
            gen_kwargs["web_search_options"] = {}

        completion = self.client.chat.completions.create(
            model=model,
            messages=input_messages,
            **gen_kwargs
            )

        return completion.choices[0].message.content
    
    async def LLM_response_async(self, prompt, gen_kwargs={}, model=None):
        if model is None:
            model = self.model

        if type(prompt) == str:
            prompt = " ".join(prompt)
            input_messages = [
                {"role": "user", "content": prompt}
                ]
        elif type(prompt) == list:
            input_messages = prompt
        else:
            logging.error("prompt must be a string or a list of messages, current type: ", type(prompt))
            raise ValueError("prompt must be a string or a list of messages")
        
        if self.use_web_search:
            model = "gpt-4o-search-preview"
            gen_kwargs["web_search_options"] = {}

        completion = await self.client_async.chat.completions.create(
            model=model,
            messages=input_messages,
            **gen_kwargs
            )
        return completion.choices[0].message.content