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

def yodafy(sentence, selection):
    llm = OpenAI(model_name="gpt-3.5-turbo", temperature=.8)
    prompt = PromptTemplate(
        input_variables=["sentence", "selection"],
        template="You are Master Yoda. A young apprentice has came to tell ysou this sentence: {sentence}. {selection} the sentence as Master Yoda would. You may refer to the Star Wars and use punctuation. Remove any \n.",
    )

    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, his wisdom and the Force.',
    )

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

    return constitutional_chain.run(sentence=sentence, selection=selection)


iface = gr.Interface(fn=yodafy, 
                     inputs=["text", 
                             gr.inputs.Radio(["Rewrite", "Reply"])], 
                     outputs="text",
                     examples=[
                         ["May the Force be with you!", "Rewrite"],
                         ]
                    )
iface.launch()