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-DA',
                  embedding_function=embedding)
retriever = vectordb.as_retriever(search_kwargs={"k": 3})

def remove_duplicates(input_list):
    return list(dict.fromkeys(input_list))

def format_sources(doc):
    string = doc.metadata['source'][13:-4]
    if string == "home":
        return "https://dupperanalytics.com/"
    else:
        return "https://dupperanalytics.com/" + string

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_list = remove_duplicates([format_sources(doc) for doc in ctx])

    sources = "\n\nSources:\n" + "\n".join(sources_list)

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

iface = gr.Interface(fn=answer, inputs="text", outputs="text", examples=["Why might I choose DupperAnalytics over other providers?", "What makes the CEO unique?", "What cloud servers do Dupper Analytics use?"])
iface.launch()