app.py

import sys
import os
import html
import asyncio
import logging
import warnings

# ── Model cache setup ────────────────────────────────────────────────
# Try persistent /data/.cache first (if the volume is read-write).
# Fall back to /tmp/.cache (writable but ephemeral) if /data is read-only.
_cache_set = False
for _candidate in ["/data/.cache", "/tmp/.cache"]:
    try:
        os.makedirs(_candidate, exist_ok=True)
        # Verify it's actually writable
        _probe = os.path.join(_candidate, ".write_test")
        with open(_probe, "w") as f:
            f.write("ok")
        os.remove(_probe)
        os.environ.setdefault("HF_HOME", os.path.join(_candidate, "huggingface"))
        os.environ.setdefault("HUGGINGFACE_HUB_CACHE", os.path.join(_candidate, "huggingface", "hub"))
        os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", os.path.join(_candidate, "sentence_transformers"))
        os.environ.setdefault("TORCH_HOME", os.path.join(_candidate, "torch"))
        print(f"[cache] Model cache set to: {_candidate}")
        _cache_set = True
        break
    except OSError:
        continue
if not _cache_set:
    print("[cache] WARNING: No writable cache directory found, using defaults")

# ── Suppress harmless asyncio event-loop cleanup errors ──────────────
# Libraries like huggingface_hub and sentence-transformers create transient
# asyncio event loops during model downloads.  When Python's GC collects
# them later, BaseEventLoop.__del__ tries to close an already-closed
# selector, producing noisy "Invalid file descriptor: -1" tracebacks.
# The loops are already dead at that point so the error is cosmetic.
_original_loop_del = getattr(asyncio.BaseEventLoop, "__del__", None)

def _quiet_loop_del(self):
    """Patched __del__ that silences ValueError on stale file descriptors."""
    try:
        if _original_loop_del is not None:
            _original_loop_del(self)
    except (ValueError, OSError):
        pass  # fd already closed – nothing to clean up

asyncio.BaseEventLoop.__del__ = _quiet_loop_del

# Also silence the "Exception ignored in" log line that the interpreter
# emits independently of the traceback.
warnings.filterwarnings("ignore", message=".*Invalid file descriptor.*")
logging.getLogger("asyncio").setLevel(logging.CRITICAL)

import gradio as gr

# Ensure the local directories inside hf/ are in the path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from nrsc_rag.utils.config import load_config
from nrsc_rag.engine.retriever import RagRetriever

# Initialize RAG configuration
config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config", "settings.yaml")
config = load_config(config_path)

print("--- BOOTING NRSC RAG ENGINE ---")
retriever = None
rag_available = True
try:
    retriever = RagRetriever(config)
    print("Core RAG Engine initialized successfully!")
except Exception as e:
    print(f"CRITICAL: Failed to initialize core RagRetriever: {e}")
    rag_available = False

# Premium Dark-Mode Glassmorphism custom styling
css = """
body {
    background-color: #090d16;
    color: #f1f5f9;
    font-family: 'Outfit', 'Inter', sans-serif;
}
.gradio-container {
    background: radial-gradient(circle at 10% 20%, rgba(18, 24, 38, 0.95), rgba(9, 13, 22, 0.99) 80%);
    border: none;
    border-radius: 20px;
    box-shadow: 0 12px 40px 0 rgba(0, 0, 0, 0.5);
    padding: 30px;
}
.glass-card {
    background: rgba(255, 255, 255, 0.02) !important;
    backdrop-filter: blur(12px) !important;
    -webkit-backdrop-filter: blur(12px) !important;
    border: 1px solid rgba(255, 255, 255, 0.05) !important;
    border-radius: 16px !important;
    padding: 18px !important;
    margin-bottom: 12px;
}
.chatbot {
    background: rgba(255, 255, 255, 0.01) !important;
    backdrop-filter: blur(8px) !important;
    border: 1px solid rgba(255, 255, 255, 0.04) !important;
    border-radius: 12px !important;
    min-height: 480px !important;
}
.source-card {
    background: rgba(99, 102, 241, 0.05);
    border-left: 4px solid #6366f1;
    border-radius: 8px;
    padding: 12px;
    margin-bottom: 12px;
    font-size: 0.9rem;
    transition: transform 0.2s ease;
}
.source-card:hover {
    transform: translateY(-2px);
    background: rgba(99, 102, 241, 0.08);
}
.badge {
    background: #6366f1;
    color: white;
    padding: 2px 8px;
    border-radius: 12px;
    font-size: 0.75rem;
    font-weight: bold;
    display: inline-block;
    margin-bottom: 6px;
}
.source-title {
    font-weight: bold;
    color: #818cf8;
    margin-bottom: 4px;
}
.source-snippet {
    color: #cbd5e1;
    line-height: 1.4;
}
input, textarea {
    background: rgba(255, 255, 255, 0.03) !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    color: white !important;
    border-radius: 8px !important;
}
button.primary {
    background: linear-gradient(135deg, #6366f1, #3b82f6) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    transition: all 0.2s ease !important;
}
button.primary:hover {
    transform: scale(1.02);
    box-shadow: 0 4px 15px rgba(99, 102, 241, 0.4);
}
"""

theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Outfit"), "sans-serif"],
)

def extract_text_content(content):
    """Securely extracts string text from Gradio 5/6 multimodal dictionary/list structures."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        text_parts = []
        for item in content:
            if isinstance(item, dict) and "text" in item:
                text_parts.append(str(item["text"]))
            elif isinstance(item, str):
                text_parts.append(item)
        return "".join(text_parts)
    if isinstance(content, dict):
        if "text" in content:
            return str(content["text"])
    return str(content)

def chat_interface(user_message, history):
    """Gradient stream handler for conversational interface using Gradio dictionary format."""
    if not user_message.strip():
        return "", history
    if history is None:
        history = []
    # Append user question and empty placeholder for assistant
    history = history + [
        {"role": "user", "content": user_message},
        {"role": "assistant", "content": ""}
    ]
    return "", history

def bot_response(history):
    """Executes RAG search via core RagRetriever and streams LLM generation in GPT style."""
    if not history or len(history) < 2:
        yield history
        return

    # Extract user question (the second to last message in the messages dictionary history list)
    question = extract_text_content(history[-2]["content"])
    
    print(f"RAG Request received: '{question}'")
    
    if not rag_available or retriever is None:
        history[-1]["content"] = "⚠️ Core RAG Engine is not available. Check startup logs."
        yield history
        return

    # 1. Retrieve top context documents directly using the core retriever's optimized hybrid pipeline
    try:
        retrieved_chunks = retriever._retrieve_top_chunks(question)
    except Exception as e:
        print(f"Retrieval error: {e}")
        retrieved_chunks = []

    # Pre-clean history contents to strictly extract string text
    clean_history = []
    for msg in history[:-1]:
        clean_history.append({
            "role": str(msg.get("role", "")),
            "content": extract_text_content(msg.get("content", ""))
        })

    # 2. Build context and messages list for native chat completions
    context_parts = []
    for idx, row in enumerate(retrieved_chunks):
        chunk_text = row.get("chunk_text", "")
        filename = row.get("filename", "unknown")
        context_parts.append(f"{chunk_text}\n")
        print(f"  [Chunk {idx+1}] from '{filename}': {chunk_text[:120]}...")
    context_str = "\n".join(context_parts)
    print(f"  Total context length: {len(context_str)} chars from {len(retrieved_chunks)} chunks")

    system_prompt = f"Use the following pieces of information and conversation history to answer the user's question.\nIf you don't know the answer, just say that you don't know, don't try to make up an answer.\nDo not mention any source file names or document names in your answer. Provide a direct and natural response.\n\nContext: {context_str}"
    
    messages = [{"role": "system", "content": system_prompt}]
    for msg in clean_history:
        messages.append({
            "role": str(msg.get("role", "")),
            "content": str(msg.get("content", ""))
        })

    model_name = config["rag"].get("llm_model", "").lower()
    is_thinking_model = "nemotron" in model_name or "deepseek" in model_name or "r1" in model_name

    def format_display_response(raw_text):
        if "</think>" in raw_text:
            parts = raw_text.split("</think>", 1)
            return parts[1].strip()
        elif "<think>" in raw_text or raw_text.startswith("<think>") or is_thinking_model:
            # While the model is generating its thoughts, show a clean indicator
            return "Thinking..."
        else:
            return raw_text

    # 3. Stream LLM Response Generation
    has_local_llama = (
        hasattr(retriever, "llm") and 
        hasattr(retriever.llm, "model_path") and 
        hasattr(retriever.llm, "_llm_type") and 
        retriever.llm._llm_type == "local-llama-cpp"
    )
    
    has_hf_api = (
        hasattr(retriever, "llm") and
        hasattr(retriever.llm, "model_id") and
        hasattr(retriever.llm, "_llm_type") and
        retriever.llm._llm_type == "huggingface-inference-api"
    )

    if has_hf_api:
        if not hasattr(retriever.llm, "_client_instance") or retriever.llm._client_instance is None:
            from huggingface_hub import InferenceClient
            retriever.llm._client_instance = InferenceClient(model=retriever.llm.model_id, token=os.environ.get("HF_TOKEN"))
            
        if retriever.llm._client_instance is not None:
            try:
                response_stream = retriever.llm._client_instance.chat_completion(
                    messages=messages,
                    max_tokens=1024,
                    temperature=config["rag"].get("llm_temperature", 0.1),
                    stream=True
                )
                
                full_response = ""
                for chunk in response_stream:
                    choices = chunk.choices
                    if choices:
                        delta = choices[0].delta
                        if delta and hasattr(delta, "content") and delta.content:
                            full_response += delta.content
                            history[-1]["content"] = format_display_response(full_response)
                            yield history
                return
            except Exception as stream_err:
                print(f"Hugging Face Inference API streaming error: {stream_err}")

    elif has_local_llama:
        # Trigger lazy loading of the model if not already loaded
        if not hasattr(retriever.llm, "_llm_instance") or retriever.llm._llm_instance is None:
            print("Triggering lazy load of local Llama model...")
            try:
                from llama_cpp import Llama
                retriever.llm._llm_instance = Llama(
                    model_path=retriever.llm.model_path,
                    n_ctx=4096,
                    n_gpu_layers=-1,
                    verbose=True
                )
            except Exception as e:
                print(f"Lazy load failed: {e}")
                
        if hasattr(retriever.llm, "_llm_instance") and retriever.llm._llm_instance is not None:
            try:
                response_stream = retriever.llm._llm_instance.create_chat_completion(
                    messages=messages,
                    max_tokens=1024,
                    temperature=config["rag"].get("llm_temperature", 0.1),
                    stream=True
                )
                
                full_response = ""
                for chunk in response_stream:
                    choices = chunk.get("choices", [])
                    if choices:
                        delta = choices[0].get("delta", {})
                        if "content" in delta:
                            full_response += delta["content"]
                            history[-1]["content"] = format_display_response(full_response)
                            yield history
                return
            except Exception as stream_err:
                print(f"Streaming error: {stream_err}")
                
    # Fallback to synchronous query or Retrieval summary if streaming fails or model not present
    try:
        print("Fallback to standard query invocation...")
        # Format history string for query
        history_list = []
        for msg in clean_history:
            if msg["role"] == "user":
                history_list.append((msg["content"], ""))
            elif msg["role"] == "assistant" and history_list:
                history_list[-1] = (history_list[-1][0], msg["content"])
        
        answer = retriever.query(question, history=history_list)
        if "</think>" in answer:
            answer = answer.split("</think>", 1)[1].strip()
        history[-1]["content"] = answer
        yield history
    except Exception as query_err:
        print(f"Core query invocation failed: {query_err}")
        # Final retrieval summary fallback
        fallback_msg = (
            "⚠️ **Local model generation and Hugging Face API are unavailable in this environment (falls back to Retrieval Mode).**\n\n"
            "Below is the relevant context retrieved directly from the NRSC index files to answer your question:\n\n"
        )
        for idx, row in enumerate(retrieved_chunks):
            escaped_text = html.escape(str(row.get("chunk_text", "")))
            escaped_file = html.escape(str(row.get("filename", "Unknown")))
            fallback_msg += f"📄 **Snippet #{idx+1} (from *{escaped_file}*):**\n> {escaped_text}\n\n"
        
        history[-1]["content"] = fallback_msg
        yield history

# Build the ChatGPT style full-width centered layout
with gr.Blocks() as demo:
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 24px; padding-top: 10px;">
        <h1 style="font-size: 2.8rem; font-weight: 800; background: linear-gradient(to right, #818cf8, #3b82f6); -webkit-background-clip: text; -webkit-text-fill-color: transparent; margin-bottom: 8px;">
            NRSC Document Intelligence Space
        </h1>
        <p style="font-size: 1.15rem; color: #94a3b8; max-width: 700px; margin: 0 auto; line-height: 1.5;">
            Premium ChatGPT-style Conversational AI assistant for National Remote Sensing Centre documentation.
        </p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(elem_classes="chat-container"):
            chatbot = gr.Chatbot(
                label="Conversational AI",
                elem_classes="chatbot",
                height=560,
                show_label=False
            )
            with gr.Row():
                msg_input = gr.Textbox(
                    placeholder="Ask a question (e.g. What is Bhuvan? Or what is the rule for official tours?)", 
                    label="Question", 
                    show_label=False,
                    scale=9
                )
                submit_btn = gr.Button("Send", variant="primary", scale=1)
            
            with gr.Row():
                clear_btn = gr.Button("Clear Chat History", size="sm")

    # Wire up interactivity
    submit_btn.click(
        fn=chat_interface,
        inputs=[msg_input, chatbot],
        outputs=[msg_input, chatbot],
        queue=False
    ).then(
        fn=bot_response,
        inputs=[chatbot],
        outputs=[chatbot]
    )

    msg_input.submit(
        fn=chat_interface,
        inputs=[msg_input, chatbot],
        outputs=[msg_input, chatbot],
        queue=False
    ).then(
        fn=bot_response,
        inputs=[chatbot],
        outputs=[chatbot]
    )

    # Clear logs and outputs
    def clear_session():
        return []
        
    clear_btn.click(
        fn=clear_session,
        inputs=[],
        outputs=[chatbot],
        queue=False
    )

if __name__ == "__main__":
    # Listen on all interfaces so HF Space proxy can access the app
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, theme=theme, css=css)