import langchain as lc
from langchain import PromptTemplate
from langchain.prompts import load_prompt
import wikipedia
from model import model

# save templates to a file


# An example prompt with multiple input variables
multiple_input_prompt = PromptTemplate(
    input_variables=["adjective", "content"],
    template="Tell me a {adjective} joke about {content}.",
)
multiple_input_prompt.save("awesome_prompt.json")  # Save to JSON file
# multiple_input_prompt.format(adjective="funny", content="chickens")
# -> "Tell me a funny joke about chickens."


prompt = load_prompt("awesome_prompt.json")


def pipeline(text, word):
    model_output = ""
    input_text = prompt.format(adjective="funny", content=text)
    while word not in model_output:
        model_output = model(input_text)
        wikipedia_entry = wikipedia.search(word)[1]
        wiki = wikipedia.summary(wikipedia_entry, auto_suggest=False, redirect=True)
        input_text += model_output + wiki
    return model_output


if __name__ == "__main__":
    print("pipline test")
    pipeline("This is the input text.", "input")