"""
Data Scientist.: Dr.Eddy Giusepe Chirinos Isidro

Este script só funciona com langchain versão 0.0.142
Você pode tentar executar ele com uma versão mais recente.
"""
from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI
from langchain.chains import LLMChain
# WE ARE NOW IMPORTING THE SEQUENTIAL CHAIN
from langchain.chains import SequentialChain

# Importar para respostas de streaming:
from langchain.callbacks.base import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler


from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv()) # read local .env file
#openai.api_key  = os.environ['OPENAI_API_KEY']


llm = OpenAI(streaming=True,
             callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]),
             verbose=True, temperature=0
            )

# Para usar o LLMChain, primeiro criamos um modelo de prompt.
# DEFINIMOS TRÊS VARIÁVEIS DE ENTRADA: technology, tone, duração e 
expand_prompt = PromptTemplate(
    input_variables=["technology", "tone", "how long"],
    template="Escreva uma descrição de {how long} no tom de {tone} explicando a seguinte tecnologia: {technology}",
)

# We can now create a very simple chain that will take user input, format the prompt with it, and then send it to the LLM.
# WE DEFINE THE OUTPUT KEY AS text
expand_chain = LLMChain(llm=llm, prompt=expand_prompt, output_key="text")

# Defina o prompt para a segunda cadeia:
review_prompt = PromptTemplate(
    input_variables=["text"],
    template="""Escreva uma resenha de quão bem escrito o texto a seguir está: {text}?
    Revisão do texto quanto à gramática, ortografia e tom:""",
)
# Defina a segunda chain:
review_chain = LLMChain(llm=llm,
                        prompt=review_prompt,
                        output_key="review"
                       )


multi_chain_multi_input_multi_output = SequentialChain(chains=[expand_chain, review_chain], verbose=False, input_variables=["technology", "tone", "how long"],  output_variables=["text", "review"])

# Execute a cadeia especificando apenas as variáveis de entrada COMO DICIONÁRIO.NO RUN METHOD THIS TIME
resp = multi_chain_multi_input_multi_output({"technology": "Docker", "tone": "Informal", "how long": "20 palavras como máximo"})