import streamlit as st
from llama_index import SimpleDirectoryReader, GPTVectorStoreIndex
from langchain import OpenAI
import gradio as gr
import os
import shutil

os.environ['OPENAI_API_KEY'] = "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"

def ingest(docs_dir):
    documents = SimpleDirectoryReader(docs_dir).load_data()
    index = GPTVectorStoreIndex.from_documents(documents)
    return index

def query(index, query_text):
    query_engine = index.as_query_engine()
    response = query_engine.query(query_text)
    return response
    
def get_answer(index, message):
    response = query(index, message)
    return [('Chatbot', ''.join(response.response))]

if os.environ['OPENAI_API_KEY']:
    # Initialize chatbot history
    chatbot = []

    index = ingest('temp_docs')

    # Display message input component
    message = st.text_input('Enter message')

    # If message is entered, ingest documents and get chatbot response
    if message:
        chatbot.append(('You', message))
        chatbot += get_answer(index, message)

    # Display chat history
    st.text_area('Chatbot:', value='\n'.join(
        [f'{x[0]}: {x[1]}' for x in chatbot]), height=250)