import os
from dotenv import load_dotenv
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate


# Set your OpenAI API key (using .env file)
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")


llm = OpenAI(temperature=0)

def generate_json(sentence):

    prompt_template_name = PromptTemplate(
    input_variables =['sentence'],
    template = '''Act as a information extractor from the sentences.You should extract the information like Agent Type, Agent Profile, PAN details, Email ID, Contact Number, First Name, Last Name, Gender, Date of Birth, Address, Pincode, City, State, Occupation from the sentences and return all the fields and corresponding value in json format. 
    If any field is missing in the sentence then ask user to provide that information and add it to json schema. If any of the field is not provided in the sentence then don't add the things by your own.
    Provided sentence is: {sentence}'''
    )
    
    p = prompt_template_name.format(sentence=sentence)
    print(p)

    response = llm.predict(p)

    return response

if __name__ == "__main__":
    print(generate_json("My name is Mr x and my address is H24 house, my pan details is XYz."))