app.py

import os
import gradio as gr
from typing import List, Tuple
import numpy as np
from pypdf import PdfReader
from docx import Document
from sentence_transformers import SentenceTransformer
import faiss
from huggingface_hub import InferenceClient

# --- Configuration & Styling ---
# SPJIMR Branding Colors
PRIMARY_BLUE = "#004890"
ACCENT_GOLD = "#C5A059"
BG_LIGHT = "#F8F9FA"
TEXT_DARK = "#1E1E1E"

CUSTOM_CSS = f"""
.container {{
    max-width: 900px;
    margin: auto;
    padding: 20px;
    font-family: 'Inter', sans-serif;
}}
.header {{
    background-color: {PRIMARY_BLUE};
    padding: 20px;
    border-radius: 12px 12px 0 0;
    color: white;
    text-align: center;
    border-bottom: 4px solid {ACCENT_GOLD};
    margin-bottom: 20px;
}}
.chat-container {{
    border-radius: 12px;
    box-shadow: 0 4px 20px rgba(0,0,0,0.08);
    background: white;
}}
.gr-button-primary {{
    background-color: {PRIMARY_BLUE} !important;
    border: none !important;
    border-radius: 8px !important;
    color: white !important;
}}
.gr-button-secondary {{
    background-color: white !important;
    border: 1px solid {PRIMARY_BLUE} !important;
    border-radius: 8px !important;
    color: {PRIMARY_BLUE} !important;
}}
#footer {{
    text-align: center;
    color: #666;
    margin-top: 20px;
    font-size: 0.9em;
}}
.user-bubble {{
    background-color: {PRIMARY_BLUE} !important;
    color: white !important;
    border-radius: 15px 15px 0 15px !important;
}}
.bot-bubble {{
    background-color: #F1F1F1 !important;
    color: {TEXT_DARK} !important;
    border-radius: 15px 15px 15px 0 !important;
}}
"""

# --- Backend Logic ---

# Initialize LLM Client (Hugging Face Inference API)
# Get token from environment variable 'HF_TOKEN'
hf_token = os.getenv("HF_TOKEN")
client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3", token=hf_token)

# Initialize Embedding Model
embed_model = SentenceTransformer("all-MiniLM-L6-v2")

class DocumentProcessor:
    def __init__(self):
        self.index = None
        self.chunks = []
        self.file_name = ""

    def extract_text(self, file_path: str) -> str:
        ext = os.path.splitext(file_path)[1].lower()
        text = ""
        if ext == ".pdf":
            reader = PdfReader(file_path)
            for page in reader.pages:
                text += page.extract_text() + "\n"
        elif ext == ".docx":
            doc = Document(file_path)
            for para in doc.paragraphs:
                text += para.text + "\n"
        elif ext == ".txt":
            with open(file_path, "r", encoding="utf-8") as f:
                text = f.read()
        return text

    def chunk_text(self, text: str, chunk_size: int = 500, overlap: int = 50) -> List[str]:
        words = text.split()
        chunks = []
        for i in range(0, len(words), chunk_size - overlap):
            chunk = " ".join(words[i : i + chunk_size])
            chunks.append(chunk)
        return chunks

    def process_document(self, file):
        if file is None:
            return "No file uploaded."
        
        self.file_name = os.path.basename(file.name)
        text = self.extract_text(file.name)
        if not text.strip():
            return "Could not extract text from document."
        
        self.chunks = self.chunk_text(text)
        embeddings = embed_model.encode(self.chunks)
        
        # Create FAISS Index
        dimension = embeddings.shape[1]
        self.index = faiss.IndexFlatL2(dimension)
        self.index.add(np.array(embeddings).astype("float32"))
        
        return f"Successfully processed: {self.file_name} ({len(self.chunks)} chunks)"

    def retrieve(self, query: str, k: int = 3) -> str:
        if self.index is None:
            return ""
        
        query_vec = embed_model.encode([query])
        distances, indices = self.index.search(np.array(query_vec).astype("float32"), k)
        
        context = ""
        for idx in indices[0]:
            if idx < len(self.chunks):
                context += self.chunks[idx] + "\n---\n"
        return context

# Global state for document processing
doc_proc = DocumentProcessor()

def chat_response(message, history):
    # System Prompt
    system_message = "You are a helpful assistant for SPJIMR. "
    
    # Check if we have context from a document
    context = doc_proc.retrieve(message)
    
    if context:
        prompt = f"{system_message}Answer the user question based ONLY on the provided context. If the answer is not in the context, say you don't know based on the document.\n\nContext:\n{context}\n\nQuestion: {message}\nAnswer:"
    else:
        prompt = f"{system_message}You are a general assistant. User: {message}\nAssistant:"

    # Call LLM
    response = ""
    try:
        # Mistral-7B format
        formatted_prompt = f"<s>[INST] {prompt} [/INST]"
        for msg in client.text_generation(formatted_prompt, max_new_tokens=512, stream=True):
            response += msg
            yield response
    except Exception as e:
        yield f"Error connecting to LLM: {str(e)}. Please ensure HF_TOKEN is set."

def clear_session():
    global doc_proc
    doc_proc = DocumentProcessor()
    return None, [], "Session cleared."

# --- UI Construction ---

with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft(primary_hue="blue")) as demo:
    with gr.Div(elem_classes="container"):
        # Header
        gr.HTML(f"""
            <div class='header'>
                <h1 style='margin:0; font-weight: 700;'>Document Assistant</h1>
                <p style='margin:5px 0 0 0; opacity: 0.9;'>SPJIMR Decision Support System</p>
            </div>
        """)
        
        with gr.Row():
            # Left Column: Upload & Controls
            with gr.Column(scale=1):
                file_input = gr.File(
                    label="Upload Document (PDF, TXT, DOCX)",
                    file_types=[".pdf", ".docx", ".txt"],
                    elem_id="upload-box"
                )
                status_msg = gr.Markdown("Status: Ready", elem_id="status")
                
                with gr.Row():
                    process_btn = gr.Button("Analyze Document", variant="primary")
                    clear_btn = gr.Button("Clear All", variant="secondary")
                
                gr.Markdown("""
                ### Instructions:
                1. Upload a document.
                2. Click 'Analyze Document'.
                3. Ask questions in the chat!
                """)

            # Right Column: Chat Interface
            with gr.Column(scale=2):
                chatbot = gr.Chatbot(
                    label="Chat Interface",
                    height=500,
                    bubble_full_width=False,
                    elem_id="chat-window"
                )
                msg_input = gr.Textbox(
                    placeholder="Type your question here...",
                    label=None,
                    container=False,
                    lines=1
                )
                
                with gr.Row():
                    submit_btn = gr.Button("Send", variant="primary")

        # Footer
        gr.HTML("<div id='footer'>Built with Gradio | SPJIMR Branding Applied</div>")

    # --- Event Handlers ---
    
    def handle_upload(file):
        if file is None:
            return "No file selected."
        msg = doc_proc.process_document(file)
        return msg

    process_btn.click(
        handle_upload,
        inputs=[file_input],
        outputs=[status_msg],
        show_progress=True
    )

    submit_btn.click(
        chat_response,
        inputs=[msg_input, chatbot],
        outputs=[chatbot],
        show_progress=True
    ).then(
        lambda: "", None, [msg_input]
    )
    
    msg_input.submit(
        chat_response,
        inputs=[msg_input, chatbot],
        outputs=[chatbot],
        show_progress=True
    ).then(
        lambda: "", None, [msg_input]
    )

    clear_btn.click(
        clear_session,
        outputs=[file_input, chatbot, status_msg]
    )

if __name__ == "__main__":
    demo.launch()