Spaces:
Running on Zero
Running on Zero
| import os | |
| import time | |
| # Force CPU mode and prevent HF Spaces ZeroGPU CUDA conflict | |
| os.environ["CUDA_VISIBLE_DEVICES"] = "" | |
| os.environ["TOKENIZERS_PARALLELISM"] = "false" | |
| # Disable Gradio 6 Server-Side Rendering (causes blank page on HF Spaces) | |
| os.environ["GRADIO_SSR_MODE"] = "false" | |
| # Handle spaces import for HF Spaces compatibility | |
| try: | |
| import spaces | |
| except ImportError: | |
| class _MockSpaces: | |
| def GPU(func=None, **kwargs): | |
| if func is not None: | |
| return func | |
| def decorator(f): | |
| return f | |
| return decorator | |
| spaces = _MockSpaces() | |
| def _gpu_init(): | |
| """Dummy GPU function required by HF Spaces ZeroGPU infrastructure.""" | |
| return "GPU Initialized" | |
| try: | |
| _gpu_init() | |
| except Exception as e: | |
| print(f"[ZeroGPU] Init call note: {e}") | |
| import gradio as gr | |
| try: | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| except ImportError: | |
| pass | |
| from rag_engine import RagEngine | |
| try: | |
| rag_engine = RagEngine() | |
| except Exception as e: | |
| print(f"[ERROR] Failed to initialize RAG engine: {e}") | |
| rag_engine = None | |
| UPLOAD_ROLES = [ | |
| "Admin (Full Access)", | |
| "HR Department Head", | |
| "Finance Department Head", | |
| "Staff / General User", | |
| ] | |
| # ════════════════════════════════════════════════════════ | |
| # UPLOAD PERMISSION & INGESTION GOVERNANCE | |
| # ════════════════════════════════════════════════════════ | |
| DOCUMENT_TYPES = [ | |
| "HR Document", | |
| "Finance Document", | |
| "Admin Document", | |
| "Public / General Document", | |
| ] | |
| # Map dropdown label → internal key | |
| _DOC_TYPE_MAP = { | |
| "HR Document": "hr", | |
| "Finance Document": "finance", | |
| "Admin Document": "admin", | |
| "Public / General Document": "public", | |
| } | |
| # Which doc types each role is allowed to upload | |
| _ROLE_ALLOWED_TYPES = { | |
| "admin": {"hr", "finance", "admin", "public"}, | |
| "hr": {"hr", "public"}, | |
| "finance": {"finance", "public"}, | |
| "staff": {"public"}, | |
| } | |
| _ROLE_LABELS = { | |
| "admin": "Admin (Full Access)", | |
| "hr": "HR Department Head", | |
| "finance": "Finance Department Head", | |
| "staff": "Staff / General User", | |
| } | |
| def _get_role_key(uploader_role: str) -> str: | |
| if "HR" in uploader_role: | |
| return "hr" | |
| if "Finance" in uploader_role: | |
| return "finance" | |
| if "Staff" in uploader_role or "General" in uploader_role: | |
| return "staff" | |
| return "admin" | |
| def _validate_upload_permission(role_key: str, doc_type: str) -> tuple[bool, str]: | |
| """ | |
| Validate whether the uploader has the necessary security authorization | |
| to ingest the given document type. | |
| Rules: | |
| - Admin : Full access (HR, Finance, Admin, Public). | |
| - HR Head : HR + Public only. Blocked from Finance / Admin. | |
| - Finance Head: Finance + Public only. Blocked from HR / Admin. | |
| - Staff : Public only. Blocked from all sensitive docs. | |
| """ | |
| allowed = _ROLE_ALLOWED_TYPES.get(role_key, set()) | |
| if doc_type in allowed: | |
| return True, f"Authorized ({_ROLE_LABELS[role_key]})" | |
| return ( | |
| False, | |
| f"⛔ Security Policy Violation: " | |
| f"**{_ROLE_LABELS[role_key]}** is NOT authorized to upload " | |
| f"**{doc_type.upper()}** documents. " | |
| f"This role may only upload: {', '.join(t.upper() for t in sorted(allowed))}." | |
| ) | |
| def ingest_file(files, uploader_role, declared_doc_type_label): | |
| """ | |
| Handle uploaded file(s): enforce strict departmental upload permissions | |
| using the EXPLICITLY DECLARED document type (not guessed from filename). | |
| """ | |
| if rag_engine is None: | |
| return "❌ RAG engine not initialized. Check GROQ_API_KEY secret." | |
| if not files: | |
| return "No file uploaded." | |
| doc_type = _DOC_TYPE_MAP.get(declared_doc_type_label, "public") | |
| role_key = _get_role_key(uploader_role) | |
| is_allowed, reason = _validate_upload_permission(role_key, doc_type) | |
| if not isinstance(files, list): | |
| files = [files] | |
| ingested_names = [] | |
| rejected_files = [] | |
| errors = [] | |
| for file in files: | |
| filepath = file.name if hasattr(file, "name") else file | |
| filename = os.path.basename(filepath) | |
| if not is_allowed: | |
| rejected_files.append(f"{filename} [{doc_type.upper()} Document] → {reason}") | |
| continue | |
| try: | |
| rag_engine.ingest_document(filepath, doc_role=doc_type) | |
| ingested_names.append(f"{filename} [{doc_type.upper()}] (Verified by {uploader_role})") | |
| except Exception as e: | |
| errors.append(f"{filename}: {e}") | |
| stats = rag_engine.get_stats() | |
| status_parts = [] | |
| if ingested_names: | |
| status_parts.append( | |
| f"✅ **Successfully Ingested ({len(ingested_names)} file(s)):**\n" + | |
| "\n".join(f" * {name}" for name in ingested_names) | |
| ) | |
| if rejected_files: | |
| status_parts.append( | |
| f"⛔ **Upload Access Control Blocked ({len(rejected_files)} file(s)):**\n" + | |
| "\n".join(f" * {rej}" for rej in rejected_files) | |
| ) | |
| if errors: | |
| status_parts.append( | |
| f"⚠️ **Failed to Ingest:**\n" + | |
| "\n".join(f" * {err}" for err in errors) | |
| ) | |
| status_parts.append( | |
| f"📊 **Index Status:** Total chunks: **{stats['total_chunks']}** | Unique sources: **{stats['unique_sources']}**" | |
| ) | |
| return "\n\n".join(status_parts) | |
| # ════════════════════════════════════════════════════════ | |
| # AURORA MIDNIGHT CYBER-GLASS THEME (CSS) | |
| # ════════════════════════════════════════════════════════ | |
| CUSTOM_CSS = """ | |
| @import url('https://fonts.googleapis.com/css2?family=Plus+Jakarta+Sans:wght@400;500;600;700;800&family=JetBrains+Mono:wght@400;500;600&display=swap'); | |
| :root { | |
| --primary-cyan: #06b6d4; --primary-indigo: #6366f1; --primary-purple: #a855f7; | |
| --aurora-gradient: linear-gradient(135deg, #06b6d4 0%, #6366f1 50%, #a855f7 100%); | |
| --bg-dark: #050811; --surface-glass: rgba(13, 19, 33, 0.75); --border-glow: rgba(99, 102, 241, 0.25); | |
| } | |
| body, gradio-app { | |
| background-color: #050811 !important; | |
| background-image: radial-gradient(at 15% 15%, rgba(6, 182, 212, 0.08) 0px, transparent 50%), | |
| radial-gradient(at 85% 85%, rgba(168, 85, 247, 0.08) 0px, transparent 50%) !important; | |
| font-family: 'Plus Jakarta Sans', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important; | |
| color: #f1f5f9 !important; | |
| } | |
| .gradio-container { max-width: 1120px !important; margin: 0 auto !important; padding: 24px 16px !important; } | |
| .hero-banner { | |
| text-align: center; padding: 34px 24px 26px 24px; margin-bottom: 24px; | |
| background: radial-gradient(circle at 50% -20%, rgba(99, 102, 241, 0.25), rgba(6, 182, 212, 0.08) 45%, transparent 80%); | |
| border-radius: 22px; border: 1px solid rgba(99, 102, 241, 0.3); backdrop-filter: blur(20px); | |
| box-shadow: 0 16px 40px -15px rgba(0, 0, 0, 0.7), 0 0 20px -5px rgba(99, 102, 241, 0.2); | |
| } | |
| .hero-badge { | |
| display: inline-flex; align-items: center; gap: 8px; padding: 5px 16px; | |
| background: rgba(15, 23, 42, 0.8); border: 1px solid rgba(6, 182, 212, 0.4); | |
| border-radius: 9999px; font-size: 0.82rem; font-weight: 600; color: #38bdf8; | |
| margin-bottom: 12px; letter-spacing: 0.03em; box-shadow: 0 0 12px rgba(6, 182, 212, 0.2); | |
| } | |
| .hero-title { | |
| font-size: 2.35rem !important; font-weight: 800 !important; letter-spacing: -0.035em !important; | |
| background: linear-gradient(135deg, #ffffff 15%, #bae6fd 55%, #c084fc 100%); | |
| -webkit-background-clip: text; -webkit-text-fill-color: transparent; | |
| margin-bottom: 8px !important; line-height: 1.2 !important; | |
| } | |
| .hero-desc { font-size: 0.98rem !important; color: #94a3b8 !important; max-width: 660px; margin: 0 auto !important; line-height: 1.6; } | |
| .chatbot-box { | |
| border-radius: 18px !important; border: 1px solid rgba(255, 255, 255, 0.08) !important; | |
| background: rgba(11, 16, 28, 0.75) !important; backdrop-filter: blur(18px) !important; | |
| box-shadow: 0 14px 40px -10px rgba(0, 0, 0, 0.6) !important; | |
| } | |
| .quick-chip { | |
| border-radius: 12px !important; font-size: 0.82rem !important; padding: 8px 14px !important; | |
| background: rgba(19, 27, 46, 0.8) !important; border: 1px solid rgba(255, 255, 255, 0.1) !important; | |
| color: #cbd5e1 !important; font-weight: 500 !important; | |
| transition: all 0.25s cubic-bezier(0.4, 0, 0.2, 1) !important; | |
| } | |
| .quick-chip:hover { | |
| background: rgba(99, 102, 241, 0.25) !important; border-color: rgba(6, 182, 212, 0.6) !important; | |
| color: #ffffff !important; box-shadow: 0 0 14px rgba(6, 182, 212, 0.3) !important; transform: translateY(-2px); | |
| } | |
| .send-btn { | |
| background: linear-gradient(135deg, #06b6d4 0%, #6366f1 50%, #9333ea 100%) !important; | |
| border: none !important; border-radius: 12px !important; color: #ffffff !important; | |
| font-weight: 700 !important; padding: 0 24px !important; | |
| box-shadow: 0 4px 18px rgba(99, 102, 241, 0.45) !important; transition: all 0.25s ease !important; | |
| } | |
| .send-btn:hover { box-shadow: 0 6px 24px rgba(6, 182, 212, 0.65) !important; transform: translateY(-2px) !important; } | |
| .footer-text { text-align: center; font-size: 0.82rem; color: #64748b; margin-top: 28px; letter-spacing: 0.02em; } | |
| """ | |
| def chat_function(message, history): | |
| """Generate a RAG-powered response with sources and citations.""" | |
| if rag_engine is None: | |
| return "RAG engine not initialized. Check GROQ_API_KEY secret." | |
| try: | |
| chunks = rag_engine.query(message, top_k=8) | |
| if not chunks: | |
| return ( | |
| "I don't have enough information to answer that question. " | |
| "Please upload documents in the Document Ingestion tab to get started." | |
| ) | |
| answer = rag_engine.generate_answer(message, chunks) | |
| sources_info = "\n\n---\n**Sources used:**\n" | |
| for i, chunk in enumerate(chunks, 1): | |
| source = chunk.get("source", "unknown") | |
| page_info = f" (Page {chunk.get('page')})" if chunk.get("page") else "" | |
| score = chunk.get("score", None) | |
| score_str = f" `[RRF: {score:.4f}]`" if score is not None else "" | |
| sources_info += f"- **Source {i}:** `{source}`{page_info}{score_str}\n" | |
| return answer + sources_info | |
| except Exception as e: | |
| print(f"[ERROR] Error in chat function: {e}") | |
| return f"Error: {e}" | |
| def get_index_info(): | |
| """Get current index statistics.""" | |
| if rag_engine is None: | |
| return "RAG engine not initialized. Check GROQ_API_KEY secret." | |
| stats = rag_engine.get_stats() | |
| if stats["total_chunks"] == 0: | |
| return "Index is empty. Upload documents to get started." | |
| sources_list = "\n".join(f" * {s}" for s in stats["sources"]) | |
| return ( | |
| f"**Index Statistics:**\n" | |
| f"- Total chunks: **{stats['total_chunks']}**\n" | |
| f"- Unique sources: **{stats['unique_sources']}**\n" | |
| f"- Active Documents:\n{sources_list}" | |
| ) | |
| def get_governance_dashboard(): | |
| """Governance & trust layer summary: access roles, staleness, upload policies, and live audit log.""" | |
| if rag_engine is None: | |
| return "RAG engine not initialized. Check GROQ_API_KEY secret." | |
| summary = rag_engine.get_audit_summary() | |
| role_lines = "\n".join( | |
| f" * **{r.upper()}**: {c} chunks" for r, c in summary["chunks_by_role"].items() | |
| ) or " * No documents indexed yet" | |
| log_lines = "\n".join( | |
| f" * `{time.strftime('%H:%M:%S', time.localtime(e['timestamp']))}` " | |
| f"\"{e['question'][:50]}\" → {', '.join(e['sources_retrieved'][:3]) or 'no match'}" | |
| for e in summary["recent_log"] | |
| ) or " * No queries logged yet" | |
| return ( | |
| f"### 🛡️ Enterprise Governance & Trust Layer\n\n" | |
| f"**System Statistics:**\n" | |
| f"- Total queries logged: **{summary['total_queries']}**\n" | |
| f"- Stale chunks flagged (>180 days): **{summary['stale_chunk_count']}**\n\n" | |
| f"**Documents Indexed by Department / Role:**\n{role_lines}\n\n" | |
| f"**🔒 Departmental Upload Access Policies:**\n" | |
| f"- **Admin (Full Access)**: Authorized across all documents (Admin, HR, Finance, Public)\n" | |
| f"- **HR Department Head**: Restricted to HR & Public documents\n" | |
| f"- **Finance Department Head**: Restricted to Finance & Public documents\n" | |
| f"- **Staff / General User**: Restricted strictly to Public/General documents\n\n" | |
| f"**Recent Audit Log (last 15 queries):**\n{log_lines}" | |
| ) | |
| def clear_index(): | |
| """Clear the entire RAG index.""" | |
| if rag_engine is None: | |
| return "RAG engine not initialized. Check GROQ_API_KEY secret." | |
| rag_engine.clear_index() | |
| return "Index cleared. Upload new documents to rebuild the knowledge base." | |
| def respond(message, chat_history): | |
| """Top-level respond function (must be picklable for ZeroGPU).""" | |
| if not message or not message.strip(): | |
| return "", chat_history | |
| bot_message = chat_function(message, chat_history) | |
| chat_history = chat_history or [] | |
| chat_history.append({"role": "user", "content": message}) | |
| chat_history.append({"role": "assistant", "content": bot_message}) | |
| return "", chat_history | |
| def clear_chat(): | |
| """Top-level clear chat function (must be picklable for ZeroGPU).""" | |
| return [] | |
| # ════════════════════════════════════════════════════════ | |
| # BUILD GRADIO UI — demo must be a top-level variable | |
| # ════════════════════════════════════════════════════════ | |
| demo = gr.Blocks( | |
| title="Universal RAG AI Knowledge Assistant", | |
| theme=gr.themes.Soft( | |
| primary_hue="indigo", | |
| secondary_hue="slate", | |
| neutral_hue="slate", | |
| ), | |
| ) | |
| with demo: | |
| gr.HTML(f"<style>{CUSTOM_CSS}</style>") | |
| gr.HTML(""" | |
| <div class="hero-banner"> | |
| <div class="hero-badge"> | |
| <span>Powered by Groq AI</span> | <span>Hybrid RRF Vector Search</span> | <span>Enterprise Governance</span> | |
| </div> | |
| <div class="hero-title">Universal RAG AI Assistant</div> | |
| <div class="hero-desc"> | |
| Upload research papers, resumes, technical reports, or corporate files to build a | |
| context-grounded AI knowledge base with strict departmental upload governance. | |
| </div> | |
| </div> | |
| """) | |
| with gr.Tabs(): | |
| with gr.TabItem("Chat", id="chat_tab"): | |
| chatbot = gr.Chatbot( | |
| label="Knowledge Assistant", | |
| height=480, | |
| elem_classes=["chatbot-box"], | |
| ) | |
| gr.Markdown("###### Suggested Queries") | |
| with gr.Row(): | |
| chip1 = gr.Button("Summarize all documents", elem_classes=["quick-chip"], size="sm") | |
| chip2 = gr.Button("Extract projects and skills", elem_classes=["quick-chip"], size="sm") | |
| chip3 = gr.Button("Explain methodology", elem_classes=["quick-chip"], size="sm") | |
| chip4 = gr.Button("Key metrics and benchmarks", elem_classes=["quick-chip"], size="sm") | |
| with gr.Row(): | |
| msg = gr.Textbox( | |
| label="Your Query", | |
| placeholder="Ask anything about your documents...", | |
| lines=2, | |
| scale=5, | |
| ) | |
| submit_btn = gr.Button("Send", variant="primary", scale=1, elem_classes=["send-btn"]) | |
| with gr.Row(): | |
| clear_btn = gr.Button("Clear Chat", variant="secondary", size="sm") | |
| chip1.click(lambda: "Summarize the key findings, candidates, and objectives across all uploaded documents.", outputs=msg) | |
| chip2.click(lambda: "What are the main projects, technical skills, and tools mentioned in the documents?", outputs=msg) | |
| chip3.click(lambda: "Explain the dataset, experimental methodology, and pipeline used in the research papers.", outputs=msg) | |
| chip4.click(lambda: "What are the key performance metrics, evaluation benchmarks, and quantitative results?", outputs=msg) | |
| submit_btn.click(respond, inputs=[msg, chatbot], outputs=[msg, chatbot]) | |
| msg.submit(respond, inputs=[msg, chatbot], outputs=[msg, chatbot]) | |
| clear_btn.click(clear_chat, None, chatbot, queue=False) | |
| with gr.TabItem("Document Ingestion", id="ingest_tab"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Enterprise Multi-Document Upload") | |
| gr.Markdown(""" | |
| 🔒 **Departmental Upload Access Control**: | |
| * **Admin (Full Access)**: Authorized to upload all documents (Admin, HR, Finance, Public). | |
| * **HR Department Head**: HR and Public documents only — **Finance & Admin blocked**. | |
| * **Finance Department Head**: Finance and Public documents only — **HR & Admin blocked**. | |
| * **Staff / General User**: Public/General documents only — **all sensitive docs blocked**. | |
| > ⚠️ You **must** declare the Document Type before uploading. The system enforces this — no bypass via filename. | |
| """) | |
| uploader_role_dropdown = gr.Dropdown( | |
| UPLOAD_ROLES, | |
| value="Admin (Full Access)", | |
| label="🔐 Your Role / Department", | |
| ) | |
| doc_type_dropdown = gr.Dropdown( | |
| DOCUMENT_TYPES, | |
| value="Public / General Document", | |
| label="📂 Document Type (declare before uploading)", | |
| ) | |
| file_input = gr.File( | |
| label="Select Files (PDF, DOCX, TXT, CSV, XLSX, XLS)", | |
| file_types=[".pdf", ".docx", ".txt", ".csv", ".xlsx", ".xls"], | |
| file_count="multiple", | |
| ) | |
| ingest_file_btn = gr.Button("Ingest Document(s)", variant="primary", elem_classes=["send-btn"]) | |
| file_status = gr.Textbox(label="Ingestion & Security Log", interactive=False, lines=8) | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Index Management") | |
| gr.Markdown("Check indexed chunks or clear the vector store to upload new documents.") | |
| index_info = gr.Textbox(label="Vector Store Status", interactive=False, lines=6) | |
| with gr.Row(): | |
| refresh_btn = gr.Button("Refresh Stats", variant="secondary") | |
| clear_index_btn = gr.Button("Clear Index", variant="stop") | |
| ingest_file_btn.click(ingest_file, inputs=[file_input, uploader_role_dropdown, doc_type_dropdown], outputs=file_status) | |
| refresh_btn.click(get_index_info, outputs=index_info) | |
| clear_index_btn.click(clear_index, outputs=index_info) | |
| with gr.TabItem("Governance Dashboard", id="governance_tab"): | |
| gr.Markdown("### 🛡️ Enterprise Governance & Trust Layer") | |
| gr.Markdown("Real-time monitoring of document departmental distribution, staleness, upload policies, and query audit trails.") | |
| dashboard_output = gr.Markdown("Click 'Refresh Dashboard' to load live governance statistics.") | |
| refresh_dash_btn = gr.Button("Refresh Dashboard", variant="primary", elem_classes=["send-btn"]) | |
| refresh_dash_btn.click(get_governance_dashboard, outputs=dashboard_output) | |
| with gr.TabItem("Architecture", id="about_tab"): | |
| gr.Markdown(""" | |
| ### Technical Architecture and Pipeline | |
| ``` | |
| [Uploaded Files: PDF / DOCX / TXT / CSV / XLSX] --(Role-Validated by Department Head)--> | |
| │ | |
| ▼ | |
| [Page & Sheet-Aware Chunking (pypdf + pandas/openpyxl + RecursiveSplitter)] | |
| │ | |
| ▼ | |
| [SentenceTransformers: all-MiniLM-L6-v2 (384-dim)] | |
| │ | |
| ▼ | |
| [FAISS Vector Store FlatL2] + [Source-Aware Keyword Matcher] | |
| │ | |
| ▼ | |
| [Reciprocal Rank Fusion (RRF): Score = 1/(60+R_sem) + 1/(60+R_kw)] | |
| │ | |
| ▼ | |
| [Governance Layer: Department Upload RBAC + Staleness Flag + Live Audit Log] | |
| │ | |
| ▼ | |
| [Groq Qwen 3.6 27B LLM Inference with Structured Prompting and Citations] | |
| │ | |
| ▼ | |
| [Direct Polished User Answer with Source, Sheet, and Page Citations] | |
| ``` | |
| | Layer | Component | Detail | | |
| | :--- | :--- | :--- | | |
| | **Upload Security** | Department Head Validation | Enforces upload authorization (Admin / HR / Finance / Staff) | | |
| | **LLM Inference** | Groq API (`qwen/qwen3.6-27b`) | Ultra-fast token generation and deep synthesis | | |
| | **Vector Database** | FAISS | In-memory L2 similarity search | | |
| | **Embeddings** | `all-MiniLM-L6-v2` | Dense sentence representations | | |
| | **Hybrid Ranker** | Reciprocal Rank Fusion (RRF) | Balances semantic similarity and keyword precision | | |
| | **Governance** | Trust & Audit Layer | Real-time query audit trail & staleness tracking (>180d) | | |
| | **Parsing** | Page-Aware Native pypdf | Preserves page numbers and table context | | |
| """) | |
| gr.HTML('<div class="footer-text">Universal RAG AI Knowledge Assistant | Powered by Groq and FAISS</div>') | |
| if __name__ == "__main__": | |
| is_hf = bool(os.environ.get("SPACE_ID")) | |
| port = int(os.environ.get("GRADIO_SERVER_PORT", 7860)) | |
| server_name = os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0" if is_hf else "127.0.0.1") | |
| is_cloud = is_hf or server_name == "0.0.0.0" | |
| demo.launch(ssr_mode=False, server_name=server_name, server_port=port, share=not is_cloud) | |
| else: | |
| pass | |