import openai
import random
import time
import gradio as gr
import os
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import DeepLake
from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationalRetrievalChain

def set_api_key(key):
    os.environ["OPENAI_API_KEY"] = key
    return f"Key Successfully Set to: {key}"
    
def get_api_key():
    api_key = os.getenv("OPENAI_API_KEY")
    return api_key

def respond(message, chat_history):
    
    # Get embeddings
    embeddings = OpenAIEmbeddings()
    
    #Connect to existing vectorstore
    db = DeepLake(dataset_path="./documentation_db", embedding_function=embeddings, read_only=True)
    #Set retriever settings
    retriever = db.as_retriever(search_kwargs={"distance_metric":'cos',
                                               "fetch_k":20,
                                               "maximal_marginal_relevance":True,
                                               "k":20})

    # Create ChatOpenAI and ConversationalRetrievalChain
    model = ChatOpenAI(model='gpt-3.5-turbo')
    qa = ConversationalRetrievalChain.from_llm(model, retriever)

    chat_history=[]
    bot_message = qa({"question": message, "chat_history": chat_history})
    chat_history.append((message, bot_message['answer']))
    time.sleep(1)
    return "", chat_history

with gr.Blocks() as demo:
    with gr.Tab("OpenAI API Key Submission"):
        api_input = gr.Textbox(label = "API Key", placeholder = "Please provide your OpenAI API key here")
        api_submission_conf = gr.Textbox(label = "Submission Confirmation")
        api_submit_button = gr.Button("Submit")
    
    with gr.Tab("Coding Assistant"):
        api_check_button = gr.Button("Get API Key")
        api_print = gr.Textbox(label = "OpenAI API Key - Please ensure the API Key is set correctly")
        chatbot = gr.Chatbot(label="ChatGPT Powered Coding Assistant")
        msg = gr.Textbox(label="User Prompt", placeholder="Your Query Here")
        clear = gr.Button("Clear")
        
    api_submit_button.click(set_api_key, inputs=api_input, outputs=api_submission_conf)
    api_check_button.click(get_api_key, outputs=api_print)
    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)
    
demo.launch()