import gradio as gr
from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain.vectorstores import Chroma
import openai
import os

openai.api_key = os.environ["OPENAI_API_KEY"]

embedding = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
vectordb = Chroma(persist_directory='db',
                  embedding_function=embedding)
retriever = vectordb.as_retriever(search_kwargs={"k": 4})

def answer(message):
    ctx = retriever.get_relevant_documents(message)

    prompt = """Use the following pieces of context to answer the question at the end. 
    If you don't know the answer, just say that you don't know, don't try to 
    make up an answer.

    {context}

    Question: {question}
    Helpful Answer: """.format(context = " ".join([doc.page_content for doc in ctx]), question = message)
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages= [
            {'role': 'user', 'content': prompt}
        ],
        temperature=0
    )
    sources = "\n\nSources:\n" + "\n".join([doc.metadata['source'] for doc in ctx])

    ans = response['choices'][0]['message']['content'] + sources
    return ans

iface = gr.Interface(fn=answer, inputs="text", outputs="text", examples=["When was the declaration of independence signed?", "What year did the United States enter World War I?", "What was considered the peak of tension within the Cold War?"])
iface.launch()