import os
import json
from gpt_index import GPTSimpleVectorIndex
from gpt_index.readers.file.docs_parser import PDFParser
import gradio as gr
from gpt_index.readers.schema.base import Document

def save_to_file(index, file):
        out_dict = {
            "index_struct": index.index_struct.to_dict(),
            "docstore": index.docstore.to_dict(),
        }
        with open(file, "w") as f:
            json.dump(out_dict, f)

def load_data(file):
        data = ""
        data_list = []
        parser = PDFParser()
        data = parser.parse_file(file)
        data_list.append(data)
        return [Document(d) for d in data_list]

def index(file, key):
    if key:
        os.environ["OPENAI_API_KEY"] = key
        documents = load_data(file.name)
        index = GPTSimpleVectorIndex(documents)
        save_to_file(index, 'index.json')
        os.environ["OPENAI_API_KEY"] = ""
    return "index.json"

key=gr.Textbox(
            placeholder="Paste your OpenAI API key (sk-...)",
            show_label=False,
            lines=1,
            type="password",
        )

demo = gr.Interface(
    index,[
    gr.File(label="PDF files only. Submit will convert to a ready to download GPT-Index index.json file with vector embeddings", file_count="single", file_types=["file"]), key],
    "file"
)

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