import os
from pathlib import Path
from dotenv import load_dotenv
from openai import OpenAI
from typing import List

# TODO: merge the .env file
env_path = Path(__file__).resolve().parents[0] / '.env'
load_dotenv(env_path)
client = OpenAI()

def openai_call(
    messages: List[dict], llm_model_name: str = os.environ.get("MODEL_NAME"), json_format=False
) -> str:
    """
    Call the OpenAI API to get the response
    Args:
        messages (List[dict]): The messages to send to the OpenAI API
        llm_model_name (str): The name of the LLM model


    Returns:
        response_json_str (str): The response from the OpenAI API
        cost (float): The cost of the response
    """
    result_json_str = ""
    while True:
        if json_format:
            response = client.chat.completions.create(
                model=llm_model_name,
                response_format={"type": "json_object"},
                messages=messages,
                temperature=0.0,
            )
        else:
            response = client.chat.completions.create(
                model=llm_model_name,
                messages=messages,
                temperature=0.0,
                # logprobs = True,
                # top_logprobs = 2
            )
        content = response.choices[0].message.content
        result_json_str += content
        # if finish_reason is length, then it is not complete
        if response.choices[0].finish_reason != "length":
            break
        else:
            messages.append(
                {
                    "role": "assistant",
                    "content": content,
                }
            )
            messages.append(
                {
                    "role": "user",
                    "content": "Continue the response",
                }
            )

    return result_json_str

if __name__ == '__main__':
    openai_call([{"role": "user", "content": "say it is a test"}])