Update app.py

app.py

@@ -1,135 +1,120 @@
-import os
-import time
-from datetime import datetime, timedelta
-import gradio as gr
-from datasets import load_dataset, Dataset, DatasetDict, concatenate_datasets
-from huggingface_hub import HfFolder
-from sentence_transformers import SentenceTransformer, util
-import torch
+# app.py (Hugging Face Space)
+
 import requests
+from datasets import load_dataset, Dataset, concatenate_datasets
+from sentence_transformers import SentenceTransformer, util
+import gradio as gr
+import os
+
+# ----------------------------
+# Load datasets
+# ----------------------------
+
+# Your main dataset (where new Q/A will be saved)
+main_ds_name = "princemaxp/cybersecurity-main"
+main_ds = load_dataset(main_ds_name, split="train")
+
+# External datasets (read-only)
+external_datasets = [
+    ("Trendyol-Security/cybersecurity-defense-instruction-tuning-v2", "train"),
+    ("Rowden/CybersecurityQAA", "train"),
+    ("Nitral-AI/Cybersecurity-ShareGPT", "train"),
+]
+
+ext_ds_list = [load_dataset(name, split=split) for name, split in external_datasets]
+
+# Keywords dataset (for classification)
+keywords_ds = load_dataset("princemaxp/cybersecurity-keywords", split="train")
+
+# ----------------------------
+# Embedding model
+# ----------------------------
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
+
+# Precompute embeddings
+main_q_embeddings = embedder.encode(main_ds["question"], convert_to_tensor=True)
+
+ext_q_embeddings = []
+for ds in ext_ds_list:
+    ext_q_embeddings.append(embedder.encode(ds["question"], convert_to_tensor=True))
+
+keywords = keywords_ds["keyword"]
+keyword_embeddings = embedder.encode(keywords, convert_to_tensor=True)
+
+# ----------------------------
+# Helper functions
+# ----------------------------
+
+def is_cybersecurity_question(user_query, threshold=0.65):
+    query_embedding = embedder.encode(user_query, convert_to_tensor=True)
+    cos_sim = util.cos_sim(query_embedding, keyword_embeddings)
+    max_score = cos_sim.max().item()
+    return max_score >= threshold
+
+def search_dataset(user_query, dataset, dataset_embeddings, top_k=1, threshold=0.7):
+    query_embedding = embedder.encode(user_query, convert_to_tensor=True)
+    cos_sim = util.cos_sim(query_embedding, dataset_embeddings)[0]
+    best_idx = cos_sim.argmax().item()
+    best_score = cos_sim[best_idx].item()
 
-# ================================
-# CONFIG
-# ================================
-MODEL_TOKEN = os.environ.get("HF_TOKEN")            # for model usage
-DATASET_TOKEN = os.environ.get("dataset_HF_TOKEN")  # for dataset updates
-DATASET_NAME = "guardian-ai-qna"
-
-MAX_QUERIES = 5      # max queries per user per window
-WINDOW_HOURS = 1     # time window for rate limiting
-
-# Rate limiter store
-user_queries = {}
-
-# Save dataset token for pushes
-HfFolder.save_token(DATASET_TOKEN)
-
-# Load or create dataset
-try:
-    dataset = load_dataset(DATASET_NAME, use_auth_token=DATASET_TOKEN)
-except:
-    dataset = DatasetDict({"train": Dataset.from_dict({"question": [], "answer": [], "embedding": []})})
-
-# Load embedding model
-embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-
-# ================================
-# HELPER FUNCTIONS
-# ================================
-
-def check_rate_limit(user_id):
-    now = datetime.now()
-    queries = user_queries.get(user_id, [])
-    # Remove expired queries
-    queries = [q for q in queries if q > now - timedelta(hours=WINDOW_HOURS)]
-    user_queries[user_id] = queries
-
-    if len(queries) >= MAX_QUERIES:
-        next_allowed = min(queries) + timedelta(hours=WINDOW_HOURS)
-        wait_seconds = int((next_allowed - now).total_seconds())
-        return False, wait_seconds
-    return True, 0
-
-def log_query(user_id):
-    now = datetime.now()
-    user_queries.setdefault(user_id, []).append(now)
-
-def find_in_dataset(question, threshold=0.75):
-    if len(dataset["train"]) == 0:
-        return None
-    # Compute embedding for input
-    question_emb = embed_model.encode(question, convert_to_tensor=True)
-    # Load existing embeddings
-    existing_embs = torch.tensor(dataset["train"]["embedding"]) if dataset["train"]["embedding"] else None
-    if existing_embs is None or len(existing_embs) == 0:
-        return None
-    # Compute cosine similarities
-    similarities = util.cos_sim(question_emb, existing_embs)[0]
-    max_score, idx = torch.max(similarities, dim=0)
-    if max_score >= threshold:
-        return dataset["train"]["answer"][idx.item()]
+    if best_score >= threshold:
+        return dataset[best_idx]["answer"]
     return None
 
-def save_qna(question, answer):
-    # Load your dataset
-    dataset = load_dataset("princemaxp/guardian-ai-qna", use_auth_token=os.environ["dataset_HF_TOKEN"])
-    
-    new_ds = Dataset.from_dict({"question": [question], "answer": [answer]})
-    
-    # Correct way to add new data
-    dataset["train"] = concatenate_datasets([dataset["train"], new_ds])
-    
-    # Push updated dataset to HF
-    dataset.push_to_hub("username/guardian-ai-qna", token=os.environ["dataset_HF_TOKEN"])
-
-def call_render(question):
-    """
-    Replace this with your actual Render API call logic
-    that fetches the answer from the internet.
-    """
-    RENDER_API_URL = os.environ.get("RENDER_API_URL")
-    if not RENDER_API_URL:
-        return "Render API not configured."
-    resp = requests.post(RENDER_API_URL, json={"question": question})
-    if resp.status_code == 200:
-        return resp.json().get("answer", "No answer found.")
-    return "Error fetching answer from Render."
-
-# ================================
-# CHAT FUNCTION
-# ================================
-
-def chat(history, message, session_id):
-    # Rate limit
-    allowed, wait_seconds = check_rate_limit(session_id)
-    if not allowed:
-        return history + [(f"System", f"Rate limit reached. Try again in {wait_seconds//60} minutes.")], ""
-
-    log_query(session_id)
-
-    # Check dataset first (embedding-based)
-    response = find_in_dataset(message)
-    if response is None:
-        # Call Render API fallback
-        response = call_render(message)
-        # Save in dataset
-        save_qna(message, response)
-
-    history.append(("User", message))
-    history.append(("Guardian AI", response))
-    return history, ""
-
-# ================================
-# GRADIO UI
-# ================================
-with gr.Blocks() as demo:
-    gr.Markdown("## Guardian AI Chatbot")
-    chatbot = gr.Chatbot()
-    session_id = gr.Textbox(label="Session ID (unique per user)", value=str(time.time()), visible=False)
-    msg = gr.Textbox(label="Enter your message")
-    send_btn = gr.Button("Send")
-
-    send_btn.click(fn=chat, inputs=[chatbot, msg, session_id], outputs=[chatbot, msg])
-    msg.submit(fn=chat, inputs=[chatbot, msg, session_id], outputs=[chatbot, msg])
-
-demo.launch(server_name="0.0.0.0", server_port=7860)
+def call_render_service(question):
+    url = os.getenv("RENDER_URL", "https://render-python-app-4ty3.onrender.com/answer")
+    try:
+        response = requests.post(url, json={"question": question}, timeout=15)
+        if response.status_code == 200:
+            return response.json().get("answer", "No answer received.")
+        return "Render service error."
+    except Exception as e:
+        return f"Render service failed: {str(e)}"
+
+def save_to_main_dataset(question, answer):
+    global main_ds, main_q_embeddings
+    new_row = {"question": [question], "answer": [answer]}
+    new_ds = Dataset.from_dict(new_row)
+    main_ds = concatenate_datasets([main_ds, new_ds])
+    main_q_embeddings = embedder.encode(main_ds["question"], convert_to_tensor=True)
+    # Push back to HF Hub (requires HF token set as secret)
+    main_ds.push_to_hub(main_ds_name)
+
+def get_answer(user_query):
+    # Step 1: Check if cybersecurity-related
+    if not is_cybersecurity_question(user_query):
+        return "This doesn’t seem like a cybersecurity-related question. Please refine your query."
+
+    # Step 2: Search in main dataset
+    answer = search_dataset(user_query, main_ds, main_q_embeddings)
+    if answer:
+        return answer
+
+    # Step 3: Search in external datasets
+    for ds, emb in zip(ext_ds_list, ext_q_embeddings):
+        answer = search_dataset(user_query, ds, emb)
+        if answer:
+            save_to_main_dataset(user_query, answer)
+            return answer
+
+    # Step 4: Fallback → Render service
+    answer = call_render_service(user_query)
+    save_to_main_dataset(user_query, answer)
+    return answer
+
+# ----------------------------
+# Gradio interface
+# ----------------------------
+def chatbot_interface(user_input):
+    return get_answer(user_input)
+
+iface = gr.Interface(
+    fn=chatbot_interface,
+    inputs=gr.Textbox(lines=2, placeholder="Ask a cybersecurity question..."),
+    outputs="text",
+    title="Guardian AI Chatbot",
+    description="Ask me anything about cybersecurity!"
+)
+
+if __name__ == "__main__":
+    iface.launch()