from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader, LLMPredictor, ServiceContext, StorageContext, load_index_from_storage
from langchain import OpenAI
import gradio
import os

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

def construct_index(directory_path):
    num_outputs = 512

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

    service_context = ServiceContext.from_defaults(llm_predictor=_llm_predictor)

    docs = SimpleDirectoryReader(directory_path).load_data()

    index = GPTVectorStoreIndex.from_documents(docs, service_context=service_context)
    
    index.storage_context.persist(persist_dir="indexes")

    return index

def chatbot(input_text):
    
    storage_context = StorageContext.from_defaults(persist_dir="indexes")
     
    query_engne = load_index_from_storage(storage_context).as_query_engine()
    
    response = query_engne.query(input_text)
    
    return response.response

iface = gradio.Interface(fn=chatbot,
                     inputs=gradio.inputs.Textbox(lines=4, label="Enter your question here"),
                     outputs=gradio.outputs.Textbox(label="Generated Text"),
                     title="My Custom trained AI Chatbot")

index = construct_index("trainingData")

iface.launch()