import gradio as gr
import os
import time
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS, Chroma
from langchain.chat_models.openai import ChatOpenAI
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferWindowMemory

# RetrievalQAWithSourcesChain.from_llm()
# ConversationalRetrievalChain()


persist_direcory = "db_index"

embeddings = OpenAIEmbeddings()
# db = FAISS.load_local(persist_directory, embeddings)
chroma = Chroma(embedding_function=embeddings, persist_directory=persist_direcory)
# retriever = chroma.as_retriever(search_type="mmr", search_kwargs={"k": 10})
retriever = chroma.as_retriever(search_kwargs={"k": 10})
query = "what were the net sales of aws in the first quarter of 2023?"
print(retriever.get_relevant_documents(query))

memory = ConversationBufferWindowMemory(
    memory_key="chat_history", return_messages=False
)

qa = ConversationalRetrievalChain.from_llm(
    llm=ChatOpenAI(model_name="gpt-4", temperature=0, streaming=True),
    chain_type="stuff",
    retriever=retriever,
    memory=memory,
    get_chat_history=lambda h: h,
    verbose=False,
    # return_source_documents=True,
)


# res = qa({"question": "what is AD790?", "chat_history": []})
# print(res["answer"])

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("clear")

    def user(user_message, history):
        return "", history + [[user_message, None]]

    def bot(history):
        bot_message = qa.run({"question": history[-1][0], "chat_history": history[:-1]})
        history[-1][1] = ""
        for character in bot_message:
            history[-1][1] += character
            time.sleep(0.05)
            yield history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    clear.click(lambda: None, None, chatbot, queue=False)

demo.queue()
demo.launch(share=False, server_name="0.0.0.0")
"""
Total net revenues of approximately $14.6 billion was up 6.1% on an
operational basis

what was the net revenue during the quarter?
The net revenue during the quarter was $14.6 billion.
how much did it increase percentage on opeational basis?
The percentage increase in net revenue on an operational basis during the quarter was 6.1%.
what was the 2022 adjusted earnings per share guidance ?
"""