import os, gradio
from langchain.document_loaders import UnstructuredPDFLoader
from langchain.indexes import VectorstoreIndexCreator
import glob

os.environ['OPENAI_API_KEY'] = os.getenv("OPENAI_API_KEY")

loaders = [UnstructuredPDFLoader(glob.glob("*.pdf"))]

# Create the index, if it does not exist, and save it
if not os.path.isfile('chroma-embeddings.parquet'):
  from langchain.vectorstores import Chroma
  index = VectorstoreIndexCreator(vectorstore_cls=Chroma, vectorstore_kwargs={ "persist_directory": "VectorStoreIndex/"}).from_loaders(loaders)
  index.vectorstore.persist()

# Load the saved index
index_saved = VectorstoreIndexCreator().from_persistent_index(".")

description = '''This is an AI conversational agent that has studied the Asom Barta newspapers over the last 1 year, from June 2022 to May 2023. You can ask it any question pertaining to this period, and it will answer it.
\n\nThe AI can only frame its answers based on its worldview attained from the Asom Barta newspapers. If you ask it about anything not pertaining to the content of the 
newspapers, it will simply reply with "I don't know". Enjoy!'''

def chat_response(query):
  return index_saved.query(query)

interface = gradio.Interface(fn=chat_response, inputs="text", outputs="text", title='Asom Barta Q&A Bot', description=description)

interface.launch()