Update app.py

app.py

@@ -1,120 +1,117 @@
+import os
 import time
-from collections import defaultdict
+from datetime import datetime, timedelta
 import gradio as gr
-from datasets import load_dataset, Dataset
-from sentence_transformers import SentenceTransformer, util
-import requests
-import os
-
-# =======================
-# Configuration
-# =======================
-HF_TOKEN = os.environ.get("HF_TOKEN")
+from datasets import load_dataset, Dataset, DatasetDict
+from huggingface_hub import HfFolder
+
+# ================================
+# 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"
-RENDER_API_URL = "https://your-render-api.com/get_answer"  # Replace with your Render API
-MAX_QUERIES_PER_HOUR = 5
-SIMILARITY_THRESHOLD = 0.75
 
-# =======================
-# Load dataset
-# =======================
+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=HF_TOKEN)["train"]
+    dataset = load_dataset(DATASET_NAME, use_auth_token=DATASET_TOKEN)
 except:
-    dataset = Dataset.from_dict({"question": [], "answer": []})
-
-# Initialize embeddings
-embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-if len(dataset) > 0:
-    dataset_embeddings = embed_model.encode(dataset["question"], convert_to_tensor=True)
-else:
-    dataset_embeddings = None
-
-# =======================
-# Rate limiting
-# =======================
-user_queries = defaultdict(list)  # {session_id: [timestamps]}
-
-def check_rate_limit(session_id):
-    now = time.time()
-    # Keep only queries in the last hour
-    user_queries[session_id] = [t for t in user_queries[session_id] if now - t < 3600]
-    if len(user_queries[session_id]) >= MAX_QUERIES_PER_HOUR:
-        return False, 3600 - (now - user_queries[session_id][0])
-    user_queries[session_id].append(now)
+    dataset = DatasetDict({"train": Dataset.from_dict({"question": [], "answer": []})})
+
+# ================================
+# 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
 
-# =======================
-# Dataset search
-# =======================
-def find_in_dataset(user_input):
-    global dataset_embeddings
-    if dataset_embeddings is None or len(dataset_embeddings) == 0:
-        return None
-    user_emb = embed_model.encode(user_input, convert_to_tensor=True)
-    cos_scores = util.cos_sim(user_emb, dataset_embeddings)[0]
-    top_idx = cos_scores.argmax().item()
-    if cos_scores[top_idx] < SIMILARITY_THRESHOLD:
+def log_query(user_id):
+    now = datetime.now()
+    user_queries.setdefault(user_id, []).append(now)
+
+def find_in_dataset(question):
+    if len(dataset["train"]) == 0:
         return None
-    return dataset["answer"][top_idx]
+    for entry in dataset["train"]:
+        if question.strip().lower() == entry["question"].strip().lower():
+            return entry["answer"]
+    return None
 
-# =======================
-# Save Q&A to dataset
-# =======================
 def save_qna(question, answer):
-    global dataset, dataset_embeddings
-    new_entry = Dataset.from_dict({"question": [question], "answer": [answer]})
-    dataset = Dataset.from_dict({
-        "question": dataset["question"] + new_entry["question"],
-        "answer": dataset["answer"] + new_entry["answer"]
-    })
-    dataset_embeddings = embed_model.encode(dataset["question"], convert_to_tensor=True)
-    dataset.push_to_hub(DATASET_NAME, token=HF_TOKEN)
-
-# =======================
-# Render API fallback
-# =======================
-def call_render_api(question):
-    try:
-        response = requests.post(RENDER_API_URL, json={"question": question}, timeout=10)
-        if response.status_code == 200:
-            return response.json().get("answer", "Sorry, no answer found.")
-    except Exception as e:
-        print("Render API error:", e)
-    return "Sorry, no answer found."
-
-# =======================
-# Chat function
-# =======================
-def chat(history, user_input, session_id):
-    allowed, wait_time = check_rate_limit(session_id)
+    global dataset
+    new_entry = {"question": [question], "answer": [answer]}
+    new_ds = Dataset.from_dict(new_entry)
+    dataset["train"] = dataset["train"].concatenate(new_ds)
+    dataset["train"].push_to_hub(DATASET_NAME, token=DATASET_TOKEN)
+
+def call_render(question):
+    """
+    Replace this with your actual Render API call logic
+    that fetches the answer from the internet.
+    """
+    import requests
+    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"Rate limit reached. Please wait {int(wait_time//60)} minutes.", "")]
-
-    answer = find_in_dataset(user_input)
-    if not answer:
-        answer = call_render_api(user_input)
-        save_qna(user_input, answer)
-
-    history.append((user_input, answer))
-    return history
-
-# =======================
-# Gradio App
-# =======================
-with gr.Blocks() as app:
-    session_id = gr.State()
+        return history + [(f"System", f"Rate limit reached. Try again in {wait_seconds//60} minutes.")], ""
+
+    log_query(session_id)
+
+    # Check dataset first
+    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()
-    msg = gr.Textbox(label="Ask Guardian AI")
-    with gr.Row():
-        clear = gr.Button("Clear Chat")
-
-    def start_session():
-        return str(time.time())  # simple session id
-
-    session_id.value = start_session()
+    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")
 
-    msg.submit(chat, inputs=[chatbot, msg, session_id], outputs=[chatbot])
-    clear.click(lambda: [], None, chatbot)
+    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])
 
-app.launch()
+demo.launch(server_name="0.0.0.0", server_port=7860)