from gpt_index import SimpleDirectoryReader, GPTListIndex, GPTSimpleVectorIndex, LLMPredictor, PromptHelper
from langchain.chat_models import ChatOpenAI
import gradio as gr
import sys
import os

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

def construct_index(directory_path):
    max_input_size = 4096
    num_outputs = 512
    max_chunk_overlap = 20
    chunk_size_limit = 600

    prompt_helper = PromptHelper(max_input_size, num_outputs, max_chunk_overlap, chunk_size_limit=chunk_size_limit)

    llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.7, model_name="gpt-3.5-turbo", max_tokens=num_outputs))

    documents = SimpleDirectoryReader(directory_path).load_data()

    index = GPTSimpleVectorIndex(documents, llm_predictor=llm_predictor, prompt_helper=prompt_helper)

    index.save_to_disk('index.json')

    return index

def chatbot(input_text):
    index = GPTSimpleVectorIndex.load_from_disk('index.json')
    prompt_text = "I want you to take the statement at the start of this query and first only answer it using information contained in documents in the 'doc1' directory and say 'this is what I find in Doc1'. Then I want you to do the same but only answer it using information contained in documents in the 'doc2' directory and say this is what I find in Doc2." 
    prompt = input_text + prompt_text 
    response = index.query(prompt, response_mode="compact")
    return response.response

iface = gr.Interface(fn=chatbot,
                     inputs=gr.components.Textbox(lines=7, label="What would you like to ask?"),
                     outputs="text",
                     title="Loss Adjuster HelpBot")

index = construct_index("doc1")
iface.launch()