gradio_app.py

import os
import requests
import gradio as gr
from typing import Optional
import uuid
import json
import time
from loguru import logger

# Use environment variables or default API base URL
API_BASE_URL = os.getenv("API_BASE_URL", "http://localhost:8000")

def upload_pdf_file(file_obj):
    """Upload a PDF file to the API"""
    if not file_obj:
        return "Please select a PDF file to upload", None

    try:
        # Add a small delay to ensure API is ready if running in container
        time.sleep(1)
        with open(file_obj.name, 'rb') as f:
            files = {'file': f}
            response = requests.post(f"{API_BASE_URL}/upload", files=files)

        if response.status_code == 200:
            result = response.json()
            pdf_id = result['pdf_id']
            filename = result['filename']
            return f"Successfully uploaded {filename} with ID: {pdf_id}", pdf_id
        else:
            return f"Error: {response.json().get('detail', 'Upload failed')}", None
    except Exception as e:
        logger.error(f"Error uploading PDF: {e}")
        return f"Error uploading PDF: {str(e)}", None

def chat_with_pdf(message: str, pdf_id: str, session_id: Optional[str] = None):
    """Chat with the PDF via the API - simplified for ChatInterface"""
    if not pdf_id:
        return "Please upload a PDF file first and use its ID"

    if not message.strip():
        return "Please enter a message"

    try:
        # Add a small delay to ensure API is ready if running in container
        time.sleep(0.5)
        # Update session ID if not provided
        if not session_id:
            session_id = str(uuid.uuid4())

        # Prepare the chat request
        payload = {
            "message": message,
            "pdf_id": pdf_id,
            "session_id": session_id
        }

        # Call the API
        response = requests.post(f"{API_BASE_URL}/chat",
                                json=payload,
                                headers={"Content-Type": "application/json"})

        if response.status_code == 200:
            result = response.json()
            bot_response = result['response']

            return bot_response
        else:
            error_detail = response.json().get('detail', 'Chat request failed')
            return f"Error: {error_detail}"
    except Exception as e:
        logger.error(f"Error chatting with PDF: {e}")
        return f"Error: {str(e)}"

def delete_pdf(pdf_id: str):
    """Delete a PDF from the server"""
    if not pdf_id:
        return "Please provide a PDF ID"

    try:
        # Add a small delay to ensure API is ready if running in container
        time.sleep(0.2)
        response = requests.delete(f"{API_BASE_URL}/pdf/{pdf_id}")

        if response.status_code == 200:
            return f"Successfully deleted PDF with ID: {pdf_id}"
        else:
            error_detail = response.json().get('detail', 'Delete failed')
            return f"Error: {error_detail}"
    except Exception as e:
        logger.error(f"Error deleting PDF: {e}")
        return f"Error deleting PDF: {str(e)}"

def query_pdf_chat(message, history, pdf_id):
    """Function for the chat interface to query PDF content"""
    if not pdf_id:
        return "Please upload a PDF file first and use its ID"
    try:
        response = chat_with_pdf(message, pdf_id)
        return response
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio interface
with gr.Blocks(title="PageIndex RAG Chatbot", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📄 PageIndex RAG Chatbot")
    gr.Markdown("Upload a PDF file and ask questions about its content using Vectorless RAG with PageIndex and Memory management using the Memori Library .")

    # State variables
    pdf_id_state = gr.State(value=None)
    session_id_state = gr.State(value=None)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📁 PDF Upload")
            pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
            upload_btn = gr.Button("Upload PDF")
            upload_status = gr.Textbox(label="Status", interactive=False)

            gr.Markdown("### 🔄 Use Existing Document")
            existing_pdf_id = gr.Textbox(label="Existing PDF ID", placeholder="Enter existing PDF ID to chat", max_lines=1)
            use_existing_btn = gr.Button("Use Existing PDF")
            use_existing_status = gr.Textbox(label="Status", interactive=False)

            gr.Markdown("### 🗑️ Delete PDF")
            delete_input = gr.Textbox(label="PDF ID to Delete")
            delete_btn = gr.Button("Delete PDF")
            delete_status = gr.Textbox(label="Delete Status", interactive=False)

        with gr.Column(scale=2):
            chat_interface = gr.ChatInterface(
                fn=query_pdf_chat,
                title="💬 PageIndex RAG Chatbot",
                description="""
                This interface allows you to interact with PDF documents using RAG (Retrieval Augmented Generation).
                Upload a PDF, then ask questions about its content.

                **Tip**: The bot supports markdown formatting including tables, lists, and code blocks.
                """,
                additional_inputs=[pdf_id_state]  # Pass the PDF ID as an additional input
            )

    # Event handling
    def use_existing_pdf(pdf_id):
        if not pdf_id or not pdf_id.strip():
            return "Please enter a valid PDF ID", None
        return f"Set to use existing PDF ID: {pdf_id}", pdf_id.strip()

    upload_btn.click(
        fn=upload_pdf_file,
        inputs=[pdf_input],
        outputs=[upload_status, pdf_id_state]
    )

    use_existing_btn.click(
        fn=use_existing_pdf,
        inputs=[existing_pdf_id],
        outputs=[use_existing_status, pdf_id_state]
    )

    delete_btn.click(
        fn=delete_pdf,
        inputs=[delete_input],
        outputs=[delete_status]
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=os.getenv("GRADIO_SHARE", "False").lower() == "true"
    )