from langchain.llms import OpenAI, GooglePalm
from Api_Key import openapi_key,google_plam
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
import PromptHelper

google_llm = GooglePalm(google_api_key=google_plam,temperature=0.1)

def Book_Name1(main_cat):

    """
    Prompt P1
    Return only Education Related
    """
    prompt_template_name = PromptTemplate(
        input_variables=['main_cat'],
        template=PromptHelper.P1
    )
    chain = LLMChain(llm=google_llm, prompt=prompt_template_name)
    response = chain.run(main_cat=main_cat)
    return response


def Book_Name1_1(main_cat):
    """
    Prompt P1_1
    Return only Education Related for given year
    """
    prompt_template_name = PromptTemplate(
        input_variables=['main_cat'],
        template=PromptHelper.P1_1
    )
    chain = LLMChain(llm=google_llm, prompt=prompt_template_name)
    response = chain.run(main_cat=main_cat)
    return response

def Book_Name2(main_cat,topic):
    """
    Prompt P2
    Return only Education Related for ant specific topic
    """
    prompt_template_name = PromptTemplate(
        input_variables=['main_cat', 'topic'],
        template=PromptHelper.P2
    )
    chain = LLMChain(llm=google_llm, prompt=prompt_template_name)
    response = chain.run(main_cat=main_cat,topic=topic)
    return response

def Book_Name3(main_cat,topic,p_year):
    """
    Prompt P3
    Return only Education Related for ant specific topic for given year
    """
    prompt_template_name = PromptTemplate(
        input_variables=['main_cat', 'topic', 'p_year'],
        template=PromptHelper.P3
    )
    chain = LLMChain(llm=google_llm, prompt=prompt_template_name)
    response = chain.run(main_cat=main_cat,topic=topic,p_year=p_year)
    return response

def Book_Name4(main_cat,genres):
    """
    Prompt P4
    Return only Non Education Related for any specific list of genres
    """
    prompt_template_name = PromptTemplate(
        input_variables=['main_cat', 'genres'],
        template=PromptHelper.P4
    )
    chain = LLMChain(llm=google_llm, prompt=prompt_template_name)
    response = chain.run(main_cat=main_cat,genres=genres)
    return response

def Book_Name5(main_cat,genres,p_year):
    """
    Prompt P5
    Return only Non Education Related for any specific list of genres for given year
    """
    prompt_template_name = PromptTemplate(
        input_variables=['main_cat', 'list_sub_cat', 'p_year'],
        template=PromptHelper.P5
    )
    chain = LLMChain(llm=google_llm, prompt=prompt_template_name)
    response = chain.run(main_cat=main_cat,genres=genres,p_year=p_year)
    return response

if __name__ == "__main__":
    print(Book_Name5('Non Education', 'Horror', '2002'))