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

os.environ['OPENAI_API_KEY'] = ''

# Load data
documents = SimpleDirectoryReader('data').load_data()

# Configure predictor and prompt helper
llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.5, model_name='gpt-3.5-turbo'))
max_input_size = 4096
num_output = 256
max_chunk_overlap = 20
prompt_helper = PromptHelper(max_input_size, num_output, max_chunk_overlap)

# Configure service context and create index
service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
index = GPTSimpleVectorIndex.from_documents(documents, service_context=service_context)
index.save_to_disk('index.json')

# Define chatbot function
def chatbot(input_text):
    index = GPTSimpleVectorIndex.load_from_disk('index.json', service_context=service_context)
    response = index.query(input_text, response_mode='default')
    return response.response

# Create Gradio interface for chatbot
iface = gr.Interface(fn=chatbot,
                     inputs=gr.components.Textbox(lines=7, label="Enter your text"),
                     outputs="text",
                     title="Custom-trained AI Chatbot")
iface.launch(share=True)