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()

llm_models = [
    "mistralai/Mixtral-8x7B-Instruct-v0.1",
    "meta-llama/Meta-Llama-3-8B-Instruct",
    "mistralai/Mistral-7B-Instruct-v0.2",
    "tiiuae/falcon-7b-instruct",
]

embed_models = [
    "BAAI/bge-small-en-v1.5",  # 33.4M
    "NeuML/pubmedbert-base-embeddings",
    "BAAI/llm-embedder", # 109M
    "BAAI/bge-large-en" # 335M
]

# Global variable for selected model
selected_llm_model_name = llm_models[0]  # Default to the first model in the list
selected_embed_model_name = embed_models[0] # Default to the first model in the list
vector_index = None

# Initialize the parser
parser = LlamaParse(api_key=os.getenv("LLAMA_INDEX_API"), result_type='markdown')

# Define file extractor with various common extensions
file_extractor = {
    '.pdf': parser,  # PDF documents
    '.docx': parser,  # Microsoft Word documents
    '.doc': parser,  # Older Microsoft Word documents
    '.txt': parser,  # Plain text files
    '.csv': parser,  # Comma-separated values files
    '.xlsx': parser,  # Microsoft Excel files
    '.pptx': parser,  # Microsoft PowerPoint files
    '.html': parser,  # HTML files
    '.jpg': parser,  # JPEG images
    '.jpeg': parser,  # JPEG images
    '.png': parser,  # PNG images
    '.webp': parser,  # WebP images
    '.svg': parser,  # SVG files
}

# Markdown content definitions
description = """
## Welcome to DocBot 📄🤖
DocBot is an intelligent document analysis tool that can help you extract insights from various document formats including:
- PDF documents
- Word documents (.docx, .doc)
- Text files
- CSV files
- Excel files
- PowerPoint presentations
- HTML files
- Images with text (JPG, PNG, WebP, SVG)
Simply upload your document, select your preferred embedding model and LLM, then start asking questions!
"""

guide = """
### How to Use DocBot:
1. **Upload Document**: Choose any supported file format
2. **Select Embedding Model**: Choose from available embedding models (BAAI/bge-small-en-v1.5 is recommended for most cases)
3. **Submit**: Click submit to process your document
4. **Select LLM**: Choose your preferred language model
5. **Ask Questions**: Start chatting with your document!
### Tips:
- Smaller embedding models (like bge-small-en-v1.5) are faster but may be less accurate
- Larger models provide better understanding but take more time
- Be specific in your questions for better results
"""

footer = """
<div style="text-align: center; margin-top: 20px; padding: 20px; border-top: 1px solid #ddd;">
    <p>Built with ❤️ using LlamaIndex and Gradio</p>
    <div style="display: flex; justify-content: center; gap: 20px; margin-top: 10px;">
        <a href="https://github.com" target="_blank">
            <img src="data:image/png;base64,{0}" alt="GitHub" style="width: 24px; height: 24px;">
        </a>
        <a href="https://linkedin.com" target="_blank">
            <img src="data:image/png;base64,{1}" alt="LinkedIn" style="width: 24px; height: 24px;">
        </a>
        <a href="https://your-website.com" target="_blank">
            <img src="data:image/png;base64,{2}" alt="Website" style="width: 24px; height: 24px;">
        </a>
    </div>
</div>
"""

# File processing function
def load_files(file_path: str, embed_model_name: str):
    try:
        if not file_path:
            return "Please select a file first."
        
        if not embed_model_name:
            return "Please select an embedding model."
            
        global vector_index
        document = SimpleDirectoryReader(input_files=[file_path], file_extractor=file_extractor).load_data()
        embed_model = HuggingFaceEmbedding(model_name=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"✅ Ready to answer questions about: {filename}"
    except Exception as e:
        return f"❌ An error occurred: {str(e)}"

# Function to handle the selected model from dropdown
def set_llm_model(selected_model):
    global selected_llm_model_name
    if selected_model:
        selected_llm_model_name = selected_model
    return f"LLM set to: {selected_model}"

# Respond function that uses the globally set selected model
def respond(message, history):
    try:
        if not vector_index:
            return "Please upload and process a document first."
            
        if not message.strip():
            return "Please enter a question."
            
        # Initialize the LLM with the selected model
        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")
        )

        # Set up the query engine with the selected LLM
        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\n{str(bot_message)}"
    except Exception as e:
        return f"❌ An error occurred: {str(e)}"

def encode_image_safe(image_path):
    """Safely encode image, return empty string if file doesn't exist"""
    try:
        if os.path.exists(image_path):
            with open(image_path, "rb") as image_file:
                return base64.b64encode(image_file.read()).decode('utf-8')
    except Exception:
        pass
    return ""

# Clear function for file processing components
def clear_file_components():
    return None, embed_models[0], ""

# Encode the images (with fallback for missing images)
github_logo_encoded = encode_image_safe("Images/github-logo.png")
linkedin_logo_encoded = encode_image_safe("Images/linkedin-logo.png")
website_logo_encoded = encode_image_safe("Images/ai-logo.png")

# UI Setup
with gr.Blocks(
    theme=gr.themes.Soft(),
    css='footer {visibility: hidden}',
    title="DocBot - Document Analysis Assistant"
) as demo:
    
    gr.Markdown("# DocBot 📄🤖")
    gr.Markdown("*Intelligent Document Analysis Assistant*")
    
    with gr.Tabs():
        with gr.TabItem("📖 Introduction"):
            gr.Markdown(description)

        with gr.TabItem("🤖 DocBot"):
            with gr.Accordion("📋 Quick Start Guide", open=False):
                gr.Markdown(guide)
                
            with gr.Row():
                with gr.Column(scale=1):
                    with gr.Group():
                        gr.Markdown("### Document Processing")
                        file_input = gr.File(
                            file_count="single", 
                            type='filepath', 
                            label="Step 1: Upload Document",
                            file_types=['.pdf', '.docx', '.doc', '.txt', '.csv', '.xlsx', '.pptx', '.html', '.jpg', '.jpeg', '.png', '.webp', '.svg']
                        )
                        
                        embed_model_dropdown = gr.Dropdown(
                            choices=embed_models, 
                            label="Step 2: Select Embedding Model", 
                            interactive=True,
                            value=embed_models[0]
                        )
                        
                        with gr.Row():
                            btn = gr.Button("🚀 Process Document", variant='primary', size="lg")
                            clear_btn = gr.Button("🗑️ Clear", size="lg")
                        
                        output = gr.Textbox(
                            label='Processing Status', 
                            interactive=False,
                            placeholder="Upload a document and click 'Process Document' to begin..."
                        )
                    
                    with gr.Group():
                        gr.Markdown("### Model Selection")
                        llm_model_dropdown = gr.Dropdown(
                            choices=llm_models, 
                            label="Step 3: Select Language Model", 
                            interactive=True,
                            value=llm_models[0]
                        )
                        llm_status = gr.Textbox(
                            label="Selected Model",
                            interactive=False,
                            value=f"LLM set to: {llm_models[0]}"
                        )
                
                with gr.Column(scale=2):
                    gr.Markdown("### Chat with Your Document")
                    chatbot = gr.Chatbot(
                        height=600,
                        placeholder="Process a document first, then start asking questions!",
                        show_label=False
                    )
                    
                    msg = gr.Textbox(
                        placeholder="Step 4: Ask questions about your document...",
                        container=False,
                        scale=7
                    )
                    
                    with gr.Row():
                        submit_btn = gr.Button("Send", variant="primary")
                        clear_chat_btn = gr.Button("Clear Chat")

    # Add footer if images exist
    if any([github_logo_encoded, linkedin_logo_encoded, website_logo_encoded]):
        gr.HTML(footer.format(github_logo_encoded, linkedin_logo_encoded, website_logo_encoded))
    
    # Set up event handlers
    def chat_respond(message, history):
        if not message.strip():
            return history, ""
        
        response = respond(message, history)
        history.append([message, response])
        return history, ""
    
    def clear_chat():
        return [], ""
    
    # Event bindings
    llm_model_dropdown.change(
        fn=set_llm_model, 
        inputs=[llm_model_dropdown], 
        outputs=[llm_status]
    )
    
    btn.click(
        fn=load_files, 
        inputs=[file_input, embed_model_dropdown], 
        outputs=[output]
    )
    
    submit_btn.click(
        fn=chat_respond,
        inputs=[msg, chatbot],
        outputs=[chatbot, msg]
    )
    
    msg.submit(
        fn=chat_respond,
        inputs=[msg, chatbot], 
        outputs=[chatbot, msg]
    )
    
    clear_btn.click(
        fn=clear_file_components,
        outputs=[file_input, embed_model_dropdown, output]
    )
    
    clear_chat_btn.click(
        fn=clear_chat,
        outputs=[chatbot, msg]
    )

# Launch the demo
if __name__ == "__main__":
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )