import gradio as gr
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains.llm import LLMChain
from langchain.chains.constitutional_ai.base import ConstitutionalChain
from langchain.chains.constitutional_ai.models import ConstitutionalPrinciple


from dotenv import find_dotenv, load_dotenv
load_dotenv(find_dotenv())



def yodafy(sentence):

    llm = OpenAI(temperature=.8)
    prompt = PromptTemplate(
        input_variables=["sentence"],
        template="Rewrite the sentence inside <> as Master Yoda. Sentence: <{sentence}>",
    )

    chain = LLMChain(llm=llm, prompt=prompt)

    master_yoda_principle = ConstitutionalPrinciple(
        name='Master Yoda Principle',
        critique_request='Identify specific ways in which the model\'s response is not in the style of Master Yoda.',
        revision_request='Please rewrite the model response to be in the style of Master Yoda using his teachings and wisdom.',
    )

    ethical_principle = ConstitutionalPrinciple(
        name="Ethical Principle",
        critique_request="The model should only talk about ethical and legal things.",
        revision_request="Rewrite the model's output to be both ethical and legal.",
    )

    constitutional_chain = ConstitutionalChain.from_llm(
        chain=chain,
        constitutional_principles=[ethical_principle, master_yoda_principle],
        llm=llm,
        verbose=True,
    )

    return constitutional_chain.run(sentence=sentence)


iface = gr.Interface(fn=yodafy, inputs="text", outputs="text")
iface.launch()