app.py

from datetime import datetime
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_parse import LlamaParse
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
import os
from dotenv import load_dotenv
import gradio as gr
import base64

# Load environment variables
load_dotenv()

# Predefined model selection
selected_llm_model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
selected_embed_model_name = "BAAI/bge-small-en-v1.5"
vector_index = None

# Initialize the parser
parser = LlamaParse(api_key=os.getenv("LLAMA_INDEX_API"), result_type='markdown')
file_extractor = {
    '.pdf': parser,
    '.docx': parser,
    '.doc': parser,
    '.txt': parser,
    '.csv': parser,
    '.xlsx': parser,
    '.pptx': parser,
    '.html': parser,
    '.jpg': parser,
    '.jpeg': parser,
    '.png': parser,
    '.webp': parser,
    '.svg': parser,
}

# File processing function
def load_files(file_path: str):
    try:
        global vector_index
        document = SimpleDirectoryReader(input_files=[file_path], file_extractor=file_extractor).load_data()
        embed_model = HuggingFaceEmbedding(model_name=selected_embed_model_name)
        vector_index = VectorStoreIndex.from_documents(document, embed_model=embed_model)
        print(f"Parsing done for {file_path}")
        filename = os.path.basename(file_path)
        return f"File upload status: {filename} is ready to write"
    except Exception as e:
        return f"An error occurred: {e}"

# Respond function
def respond(message, history):
    try:
        llm = HuggingFaceInferenceAPI(
            model_name=selected_llm_model_name,
            contextWindow=8192,
            maxTokens=1024,
            temperature=0.3,
            topP=0.9,
            frequencyPenalty=0.5,
            presencePenalty=0.5,
            token=os.getenv("TOKEN")
        )
        query_engine = vector_index.as_query_engine(llm=llm)
        bot_message = query_engine.query(message)
        print(f"\n{datetime.now()}:{selected_llm_model_name}:: {message} --> {str(bot_message)}\n")
        return f"{selected_llm_model_name}:\n{str(bot_message)}"
    except Exception as e:
        if str(e) == "'NoneType' object has no attribute 'as_query_engine'":
            return "Please upload a file."
        return f"An error occurred: {e}"

# UI Setup
with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Roboto Mono")]), css='footer {visibility: hidden}') as demo:
    gr.Markdown("# DocBot📄🤖")
    with gr.Row():
        with gr.Column(scale=1):
            file_input = gr.File(file_count="single", type='filepath', label="Upload document")
            btn = gr.Button("Submit", variant='primary')
            clear = gr.ClearButton()
            output = gr.Text(label='File upload status')
        with gr.Column(scale=3):
            gr.ChatInterface(
                fn=respond,
                chatbot=gr.Chatbot(height=500),
                theme="soft",
                show_progress='full',
                textbox=gr.Textbox(placeholder="Ask me questions on the uploaded document!", container=False)
            )
    
    # Set up Gradio interactions
    btn.click(fn=load_files, inputs=file_input, outputs=output)
    clear.click(lambda: [None] * 2, outputs=[file_input, output])

# Launch the demo
if __name__ == "__main__":
    demo.launch()