from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.llms import OpenAI

import os
os.environ['OPENAI_API_KEY'] = "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"
llm = OpenAI(temperature=0.5)
def book_summary_generator(theme):
  book_name_prompt_template = PromptTemplate(
      input_variables=["theme"],
      template="Please provide a list of ten well-known books that center around the theme of {theme}",
  )

  # Create an LLM chain with the prompt template and LLM
  book_name_chain = LLMChain(llm=llm, prompt=book_name_prompt_template, output_key="book_names_list")

  book_summary_prompt_template = PromptTemplate(
      input_variables=["book_names_list"],
      template="""
      Please take one book from the books list {book_names_list}. Start with the book title in capitals. 
      Please provide a comprehensive summary of the book, in 10 bullet points
      """
  )

  # Create an LLM chain with the new prompt template and LLM
  book_summary_chain = LLMChain(llm=llm, prompt=book_summary_prompt_template, output_key="book_summary")

  from langchain.chains import SequentialChain

  # Create a sequential chain that first gets the book names based on the theme and then gets the summary of a specific book
  book_chain = SequentialChain(
      chains=[book_name_chain, book_summary_chain],
      input_variables=["theme"],
      output_variables=["book_names_list", "book_summary"]
      )

  # Get the book summary for a specific book based on the theme
  book_summary = book_chain.invoke(theme)
  return(book_summary)

if __name__ == "__main__":
  theme = "personality development"
  book_summary = book_summary_generator(theme)
  print(book_summary)