app.py

"""
Gradio chat interface for end users
Uses Gradio 5.49 ChatInterface API
"""

import gradio as gr
import os
import time
from collections import defaultdict
from dotenv import load_dotenv
from src.chatbot import create_rag_chain, ask_question
from src.config import load_settings
import re
import uuid

load_dotenv()

# Rate limiting: max requests per window (per IP)
RATE_LIMIT_REQUESTS = int(os.getenv("RATE_LIMIT_REQUESTS", "10"))
RATE_LIMIT_WINDOW_SEC = int(os.getenv("RATE_LIMIT_WINDOW_SEC", "60"))
_rate_limit_store = defaultdict(list)  # {client_key: [timestamp, ...]}


_rate_limit_last_cleanup = 0.0


def _check_rate_limit(client_key: str) -> bool:
    """
    Sliding window rate limit. Returns True if allowed, False if exceeded.
    Periodically purges stale keys to prevent unbounded memory growth.
    """
    global _rate_limit_last_cleanup
    now = time.time()
    window_start = now - RATE_LIMIT_WINDOW_SEC

    # Purge stale keys every 5 minutes
    if now - _rate_limit_last_cleanup > 300:
        stale = [k for k, ts in _rate_limit_store.items() if not ts or ts[-1] <= window_start]
        for k in stale:
            del _rate_limit_store[k]
        _rate_limit_last_cleanup = now

    timestamps = _rate_limit_store[client_key]
    timestamps[:] = [t for t in timestamps if t > window_start]
    if len(timestamps) >= RATE_LIMIT_REQUESTS:
        return False
    timestamps.append(now)
    return True


# Initialize chatbot
print("🤖 Initializing chatbot...")
rag_chain, retriever, llm = create_rag_chain()
print("✅ Chatbot ready!")

def check_pii(text: str) -> bool:
    """
    Simple PII detection - checks for potential names
    
    Args:
        text: Input text to check
        
    Returns:
        True if PII detected
    """
    # Check for capitalized words that might be names
    name_pattern = r'\b[A-Z][a-z]+ [A-Z][a-z]+\b'
    if re.search(name_pattern, text):
        return True
    return False


def chat_response(message: str, history: list, session_id: str, request: gr.Request) -> str:
    """
    Handle chat messages (Gradio 5.x format)
    
    Args:
        message: User's message
        history: Conversation history
        session_id: Unique session ID per user (from gr.State)
        request: Injected by Gradio for IP/session access
        
    Returns:
        Bot's response
    """
    # Rate limit by IP (fallback to session_id if no client info)
    client_key = "unknown"
    if request and hasattr(request, "client") and request.client:
        client_key = getattr(request.client, "host", session_id) or session_id
    else:
        client_key = session_id

    if not _check_rate_limit(client_key):
        return "⏳ **Rate limit reached.** Please wait a minute before sending more messages. This helps us keep the service available for everyone."

    # Check for PII
    warning = ""
    if check_pii(message):
        warning = "⚠️ **Warning**: Please avoid sharing personal information about specific individuals.\n\n"
    
    # Get answer from chatbot
    try:
        answer, sources = ask_question(rag_chain, retriever, llm, message, session_id)
        
        # Format response with sources
        response = warning + answer
        
        if sources:
            response += "\n\n📚 **Sources:**\n"
            for i, doc in enumerate(sources[:3], 1):
                source = doc.metadata.get("source", "Unknown")
                response += f"{i}. {source}\n"
        
        return response
    
    except Exception as e:
        return f"❌ Error: {str(e)}\n\nPlease make sure documents have been uploaded to the system."


# Load configurable texts from config/chatbot_settings.json
_settings = load_settings()
DISCLAIMER_TEXT = _settings["disclaimer"]
WELCOME_MESSAGE = _settings["welcome_message"]
BOT_AVATAR_URL = _settings["bot_avatar_url"]
PRIMARY_COLOR = _settings["primary_color"]
SECONDARY_COLOR = _settings["secondary_color"]
FONT_FAMILY = _settings["font_family"]

_custom_css = f"""
* {{ font-family: {FONT_FAMILY} !important; }}
.gradio-container button.primary {{ background-color: {PRIMARY_COLOR} !important; border-color: {PRIMARY_COLOR} !important; }}
.gradio-container {{ background-color: {SECONDARY_COLOR} !important; }}
"""

# Create Gradio interface (Gradio 5.49 API)
with gr.Blocks(
    title="HR Intervals AI Assistant",
    theme=gr.themes.Soft(),
    css=_custom_css
) as demo:
    
    gr.Markdown("""
    # 💼 HR Intervals AI Assistant
    
    Get instant answers to your HR questions based on our knowledge base.
    """)
    
    # Disclaimer (text loaded from config)
    with gr.Accordion("⚠️ Important Disclaimer - Please Read", open=False):
        gr.Markdown(DISCLAIMER_TEXT)
    
    # Per-user session ID (each browser tab gets a unique ID)
    session_state = gr.State(value=lambda: str(uuid.uuid4()))
    
    # Chat interface (Gradio 5.x ChatInterface)
    chat_interface = gr.ChatInterface(
        fn=chat_response,
        additional_inputs=[session_state],
        chatbot=gr.Chatbot(
            height=500,
            show_label=False,
            type='messages',
            avatar_images=(None, BOT_AVATAR_URL),
            value=[{"role": "assistant", "content": WELCOME_MESSAGE}]
        ),
        textbox=gr.Textbox(
            placeholder="Ask your HR question here...",
            container=False,
            scale=7
        ),
        title="",
        description="",
        theme=gr.themes.Soft()
    )
    
    # Footer
    gr.Markdown("""
    ---
    💡 **Tip**: Be specific in your questions for better answers. Remember to consult professionals for legal matters.
    """)


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )