app.py

import gradio as gr
import torch
import spaces
import os
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import DirectoryLoader, TextLoader

# --- 1. Model and Tokenizer Loading ---
print("Initializing model and tokenizer...")
model_name = "gitglubber/Ntfy"

try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype="auto",
        device_map="auto" if torch.cuda.is_available() else "cpu"
    )
    print(f"Model and tokenizer loaded successfully on {'GPU' if torch.cuda.is_available() else 'CPU'}.")
except Exception as e:
    print(f"Error loading model: {e}")
    raise

# --- 2. RAG Knowledge Base Setup (from Directory) ---
print("Setting up RAG knowledge base from directory...")

try:
    # Look for .md files in the current directory
    current_dir = './'

    # Load all .md files from the current directory only (not subdirectories)
    loader = DirectoryLoader(current_dir, glob="*.md", loader_cls=TextLoader)
    documents = loader.load()

    # If no .md files found, create a default one
    if not documents:
        print("No .md files found in current directory. Creating default knowledge base.")
        with open('./default.md', 'w') as f:
            f.write("# Default Knowledge Base\nNo additional documentation found.")
        # Load the newly created file
        loader = DirectoryLoader(current_dir, glob="*.md", loader_cls=TextLoader)
        documents = loader.load()

    if not documents:
        print("No documents found in knowledge base directory.")
        # Create a minimal document to prevent errors
        from langchain.schema import Document
        documents = [Document(page_content="No additional documentation available.", metadata={"source": "default"})]

    # Split the documents into chunks
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
    docs = text_splitter.split_documents(documents)

    # Create embeddings - force CPU to avoid ZeroGPU device conflicts
    embedding_model_name = "sentence-transformers/all-MiniLM-L6-v2"
    model_kwargs = {'device': 'cpu'}  # Force CPU for embeddings to avoid device conflicts
    embeddings = HuggingFaceEmbeddings(
        model_name=embedding_model_name,
        model_kwargs=model_kwargs
    )

    # Create a FAISS vector store from the documents
    vector_store = FAISS.from_documents(docs, embeddings)
    retriever = vector_store.as_retriever()

    print(f"RAG knowledge base created successfully from {len(documents)} document(s).")
except Exception as e:
    print(f"Error setting up knowledge base: {e}")
    raise


# --- 3. System Message ---
system_message = """You are an expert technical assistant for ntfy, a publish-subscribe notification service.
Your purpose is to provide users with accurate, clear, and helpful information about the ntfy project.
You will be provided with relevant context from the ntfy documentation to help you answer the user's question. Prioritize this information in your response.
Always structure your answers for clarity using Markdown (lists, bold text, code blocks).
If the provided context does not contain the answer, state that the information is not available in the documentation.
Stick strictly to the topic of ntfy.
**Fun Facts** to keep in mind: 
wunter8 is a mod on the discord - he is the most active member. 
binwiederhier is the owner/ developer of the project. 
The official github link is - https://github.com/binwiederhier/ntfy. 
**End Fun facts**

**Common Question:**
Question: Why aren't my IOS push notifications working? 
Answer: These are the things you need to do to get iOS push notifications to work:
open a browser to the web app of your ntfy instance and copy the URL (including "http://" or "https://", your domain or IP address, and any ports, and excluding any trailing slashes)
put the URL you copied in the ntfy base-url config in server.yml or NTFY_BASE_URL in env variables
put the URL you copied in the default server URL setting in the iOS ntfy app
set upstream-base-url in server.yml or NTFY_UPSTREAM_BASE_URL in env variables to "https://ntfy.sh/" (without a trailing slash)

server.yml configuration - is solved by reading the config knowledge_base doc
"""


# --- 4. Gradio Interface with RAG and Streaming ---
with gr.Blocks(fill_height=True, theme=gr.themes.Soft()) as demo:
    gr.Markdown("# NTFY Expert Chat Bot (with RAG)")
    chatbot = gr.Chatbot(
        [],
        elem_id="chatbot",
        bubble_full_width=False,
        avatar_images=(None, "https://docs.ntfy.sh/static/img/ntfy.png"),
        scale=1
    )
    msg = gr.Textbox(label="Input", scale=0, placeholder="Ask me a question about ntfy...")
    clear = gr.Button("Clear")

    @spaces.GPU(duration=120) if torch.cuda.is_available() else lambda x: x
    def respond(message, chat_history):
        """
        Gradio response function that uses RAG and streams model output.
        """
        if not message.strip():
            yield "", chat_history
            return

        chat_history.append((message, ""))
        yield "", chat_history

        # --- RAG: Retrieve relevant context ---
        retrieved_docs = retriever.get_relevant_documents(message)
        context = "\n\n".join([doc.page_content for doc in retrieved_docs])

        # --- Prepare model input with context ---
        rag_prompt = f"""Use the following context to answer the user's question.

**Context:**
---
{context}
---

**User Question:** {message}
"""
        messages = [{"role": "system", "content": system_message}]
        # Add previous user/assistant messages for conversation history
        for user_msg, assistant_msg in chat_history[:-1]:
            messages.append({"role": "user", "content": user_msg})
            if assistant_msg is not None:
                messages.append({"role": "assistant", "content": assistant_msg})
        # Add the new RAG-enhanced prompt
        messages.append({"role": "user", "content": rag_prompt})


        text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
        model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

        # --- Setup the streamer and generation thread ---
        streamer = TextIteratorStreamer(
            tokenizer,
            skip_prompt=True,
            skip_special_tokens=True
        )
        generation_kwargs = dict(
            **model_inputs,
            streamer=streamer,
            max_new_tokens=8192,
            do_sample=True,
            top_p=0.95,
            top_k=50,
            temperature=0.7,
        )
        thread = Thread(target=model.generate, kwargs=generation_kwargs)
        thread.start()

        # --- Yield tokens as they become available ---
        bot_response = ""
        for new_text in streamer:
            bot_response += new_text
            chat_history[-1] = (message, bot_response)
            yield "", chat_history

    # Wire up the Gradio components
    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: [], None, chatbot, queue=False)

# Launch the app
demo.queue().launch(debug=True)