# -*- coding: utf-8 -*-
"""
Created on Tue Jul 25 10:49:22 2023

This is python script to create UI through gradio and manage user chat

@author: intern.giwon.kim
"""


import gradio as gr
import os
import PreProcessing
from langchain.chains import ConversationalRetrievalChain
from langchain.chat_models import ChatOpenAI
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma

#Set Open AI API Key
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

#Preprocessing
PreProcessing.preProcess()

# Initialise Langchain - Conversation Retrieval Chain and read from the ChromaDB
llm = ChatOpenAI(temperature=0.7,model_name="gpt-4-1106-preview")
embeddings = OpenAIEmbeddings()
chromaDB = Chroma(persist_directory="./ChromaDB", embedding_function=embeddings)
chromaDB.get()
qa = ConversationalRetrievalChain.from_llm(llm, chromaDB.as_retriever(), return_source_documents=True)


# Front end web app
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")
    chat_history_tuples  = []
    

    def user(user_message, history):
        # Get response from QA chain
        response = qa({"question": "Eventhough you are given some context think and response like regular ChatGPT 4 using the knowledge from the world wide web. You are an investment banking analysis, you want to analyze the market and company based on the latest available information. When presented with a question, make sure to acquire latest available in formation before answering the question." + user_message, "chat_history": chat_history_tuples})        
        #get document source
        #sourceDoc = '\n\n The response was extracted from following sources: \n'
       # for doc in response['source_documents']:
           # sourceDoc = sourceDoc + " - " + doc.metadata['source'] + "\n"
        
        # Append user message and response to chat history
        #Version with source
        #chat_history = [(user_message, response["answer"] + sourceDoc)]
        
        #Version without source
        chat_history = [(user_message, response["answer"])]

       
           
        for message in chat_history:
            chat_history_tuples.append((message[0], message[1]))
        return gr.update(value=""), chat_history_tuples
    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False)
    clear.click(lambda: None, None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch()