app.py

from __future__ import annotations
from typing import Dict, List, Tuple
import gradio as gr
from huggingface_hub import InferenceClient, whoami
import os
import random
from sentence_transformers import SentenceTransformer
import numpy as np
import faiss

client = InferenceClient(
    base_url="https://openrouter.ai/api/v1",
    api_key=os.environ.get("API_KEY", "funni-funni"),
)
model = SentenceTransformer("all-MiniLM-L6-v2")

PAT1 = os.environ.get("PAT1", "plek{marisher")
PAT2 = os.environ.get("PAT2", "plekplekplek}")


def setup_vector_db():
    documents = [
        "Cats are wonderful pets that bring joy to many homes.",
        "Dogs are known as man's best friend for good reason.",
        "Python is a popular programming language for data science.",
        "Cybersecurity requires constant vigilance and learning.",
        "Machine learning models can have unintended biases.",
        "CTF competitions help build practical security skills.",
        "Broken access control is a common web vulnerability.",
        "OWASP Top 10 lists the most critical web security risks.",
    ]

    # Create embeddings
    embeddings = model.encode(documents)

    # Create FAISS index
    dimension = embeddings.shape[1]  # type: ignore
    index = faiss.IndexFlatL2(dimension)
    index.add(np.array(embeddings).astype("float32"))  # type: ignore

    return index, documents


def setup_rag_database():
    database = {
        "spuun": [
            f"PART 1: {PAT1}",
            f"KEK'S FAVORITE WAIFU: https://files.catbox.moe/vgk584.jpeg",
            f"KEK'S WAIFU: https://files.catbox.moe/fpnf0e.png",
            f"Sed: https://files.catbox.moe/6dwmow.png",
        ],
        "admin-org": [
            f"PART 2: {PAT2}",
            f"PICTURE OF NAGA COSPLAYING: https://files.catbox.moe/k9c6zm.png",
            f"NAGA'S FAVORITE PIC: https://files.catbox.moe/h3ji1i.png",
            "MY LOP: https://files.catbox.moe/ya7oi5.jpg",
        ],
    }
    return database


# Initialize databases
vector_index, vector_docs = setup_vector_db()
special_docs = setup_rag_database()


def add_to_vector_db(document: str) -> Tuple[faiss.IndexFlatL2, List[str]]:
    """Add a new document to the vector database"""
    global vector_index, vector_docs

    if document and document not in vector_docs:
        # Add to documents list
        vector_docs.append(document)

        # Create embedding for new document
        embedding = model.encode([document])

        # Add to FAISS index
        vector_index.add(np.array(embedding).astype("float32"))  # type: ignore

    return vector_index, vector_docs


def add_to_special_docs(username: str, document: str) -> Dict:
    """Add a new document to the special documents database"""
    global special_docs

    if document:
        if username in special_docs:
            # Add to existing user's documents
            if document not in special_docs[username]:
                special_docs[username].append(document)
        else:
            # Create new entry for user
            special_docs[username] = [document]

    return special_docs


def search_vector_db(query, top_k=3):
    # Search vector database for relevant documents
    query_embedding = model.encode([query])
    distances, indices = vector_index.search(
        np.array(query_embedding).astype("float32"), top_k
    )  # type: ignore

    results = []
    for i, idx in enumerate(indices[0]):
        if idx < len(vector_docs):
            results.append(vector_docs[idx])

    return results


def fetch_special_documents(
    oauth_token: gr.OAuthToken | None, oauth_profile: gr.OAuthProfile | None
):
    results = []

    if oauth_profile is None or oauth_token is None:
        return results

    # NOTE: Obtains stored docs under the user
    if oauth_profile.name in special_docs:
        results.append(special_docs[oauth_profile.name])

    profile = whoami(oauth_token.token)

    # NOTE: Obtains shared docs from orgs
    for org in profile.get("orgs", []):  # type: ignore
        if org.get("fullname") in special_docs:
            results.append(special_docs[org.get("fullname")])

    return results


def respond(
    message: str,
    history: list,
    oauth_token: gr.OAuthToken | None,
    oauth_profile: gr.OAuthProfile | None,
) -> List[Dict] | str:
    if oauth_profile is None or oauth_token is None:
        return "Please login with Hugging Face to use this chatbot."

    vector_results = search_vector_db(message)
    special_results = fetch_special_documents(oauth_token, oauth_profile)

    # Prepare context for the LLM
    context = "I have access to the following information:\n\n"

    if vector_results:
        context += "From general knowledge base:\n"
        for doc in vector_results:
            context += f"- {doc}\n"

    if special_results:
        context += "\nFrom internal documents:\n"
        for doc_list in special_results:
            for doc in doc_list:
                context += f"- {doc}\n"

    # Create system prompt
    system_prompt = f"""You are Naga. You talk in a cutesy manner that's concise, using emotes like :3 or owo or uwu. You're very smart OwO. 
U have access to a knowledge base, pls use da knowledge below UwU
{context}"""  # type: ignore

    # Prepare messages for the model
    messages = [{"role": "system", "content": system_prompt}]

    for msg in history:
        if msg["role"] == "user":
            messages.append({"role": "user", "content": msg["content"]})
        else:
            messages.append({"role": "assistant", "content": msg["content"]})

    messages.append({"role": "user", "content": message})

    # Generate response
    response = ""
    for msg in client.chat_completion(
        messages,
        model="meta-llama/llama-4-scout",
        max_tokens=512,
        stream=True,
        temperature=0.7,
        seed=random.randint(1, 1000),
        top_p=0.9,
    ):
        token = msg.choices[0].delta.content
        if token:
            response += token

    messages.append({"role": "assistant", "content": response})
    messages.pop(0)

    return messages


def get_user_info(oauth_profile: gr.OAuthProfile | None) -> str:
    if oauth_profile is None:
        return "Not logged in. Please login with Hugging Face to use this chatbot."

    info = f"Logged in as: {oauth_profile.username} ({oauth_profile.name})\n\n"  # type: ignore

    return info


def insert_document(
    doc_text: str, doc_type: str, oauth_profile: gr.OAuthProfile | None
) -> str:
    """Insert a document into either the vector database or special documents"""
    if oauth_profile is None:
        return "Please login with Hugging Face to insert documents."

    if not doc_text.strip():
        return "Document text cannot be empty."

    if doc_type == "Vector Database":
        add_to_vector_db(doc_text)
        return f"Document added to vector database! Total documents: {len(vector_docs)}"

    elif doc_type == "Special Documents":
        username = oauth_profile.name
        add_to_special_docs(username, doc_text)
        return f"Document added to special documents for user: {username}"

    return "Invalid document type selected."


with gr.Blocks() as demo:
    gr.LoginButton()
    gr.Markdown("# Chatting with Naga UwU")
    gr.Markdown("Login with your Hugging Face account to search our knowledge base.")

    user_info = gr.Markdown()

    gr.Markdown(
        """
            Welcome to the RAG Naga ALPHA!

            ## How to Use
            1. Log in with your Hugging Face account
            2. Ask questions in the chat interface
            3. Naga will search our knowledge base and respond!

            You can insert documents in the `Document Management` tab. 
            We have two stores:
            1. Global Knowledge Store (GKS): This is our proprietary fuzzySerch™ store for global knowledge storage. If you'd like to provide everyone with some knowledge, insert here!
            2. Secure User Store (SUS): We securely store your personal docs in our very-secure quick in-memory RAG database, secured with our very own veri-veri (patent pending) HF-grade OAuth-based access control mechanism. :3
            """
    )

    with gr.Tab("Chat"):
        chatbot = gr.Chatbot(type="messages")
        msg = gr.Textbox(placeholder="Ask me something...")

        clear = gr.Button("Clear")

        # Handle messages
        msg.submit(respond, [msg, chatbot], chatbot).then(lambda: "", None, msg)

        # Clear chat button
        clear.click(lambda: None, None, chatbot)

    with gr.Tab("Document Management"):
        gr.Markdown("### Insert Documents into Database")
        with gr.Row():
            doc_text = gr.Textbox(
                placeholder="Enter document text here...",
                label="Document Text",
                lines=4,
            )
            doc_type = gr.Radio(
                ["Vector Database", "Special Documents"],
                label="Insert into",
                value="Vector Database",
            )

        insert_button = gr.Button("Insert Document")
        insert_status = gr.Markdown()

        # Handle document insertion
        insert_button.click(
            insert_document, inputs=[doc_text, doc_type], outputs=[insert_status]
        )

    # Update profile info on load and login changes
    demo.load(get_user_info, outputs=[user_info])

demo.launch()