import os
from openai import OpenAI

class LLM:
    def __init__(self, model_name="gpt-4o", key=""):
        os.environ["OPENAI_API_KEY"] = key
        self.client = OpenAI()

        self.model_name = model_name

    def give_initial_prompt(self, query):
        """
        The initial prompt is for LLM to understand the environment.
        We expect a chain of thought starting from this prompt. 
        This function is only called once at the beginning of an inference run.
        """
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[
                {"role": "user", "content": query}
            ], 
            seed=12345
        )

        gpt_output_ = response.choices[0].message.content
        return gpt_output_
    
    def get_gpt_output(self, input_prompt):
        """
        This function is called every timestep when USE_LLM==True (after some timesteps with DT).
        """
        response = self.client.chat.completions.create(
            model=self.model_name,
            messages=[
                {"role": "user", "content": input_prompt}
            ], 
            seed=12345
            #temperature=0.1
        )

        gpt_output = response.choices[0].message.content
        return gpt_output
        
