Update app.py

app.py

@@ -1,76 +1,66 @@
 import time
+from collections import defaultdict
 import gradio as gr
 from datasets import load_dataset, Dataset
-from huggingface_hub import hf_hub_download
 from sentence_transformers import SentenceTransformer, util
-import torch
+import requests
+import os
 
-# ---------------------------
-# CONFIGURATION
-# ---------------------------
-HF_TOKEN = "<YOUR_HF_TOKEN>"  # set your HF token
+# =======================
+# Configuration
+# =======================
+HF_TOKEN = os.environ.get("HF_TOKEN")
 DATASET_NAME = "guardian-ai-qna"
-MAX_QUESTIONS = 5          # max questions per TIME_WINDOW
-TIME_WINDOW = 3600         # 1 hour in seconds
-EMBED_MODEL = "all-MiniLM-L6-v2"  # small but effective embedding model
+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 OR CREATE DATASET
-# ---------------------------
+# =======================
+# Load dataset
+# =======================
 try:
-    dataset = load_dataset(DATASET_NAME, use_auth_token=HF_TOKEN)
-    dataset = dataset["train"]
+    dataset = load_dataset(DATASET_NAME, use_auth_token=HF_TOKEN)["train"]
 except:
     dataset = Dataset.from_dict({"question": [], "answer": []})
 
-# ---------------------------
-# EMBEDDING MODEL
-# ---------------------------
-embedder = SentenceTransformer(EMBED_MODEL)
-
-# Precompute embeddings for existing Q&A
+# Initialize embeddings
+embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 if len(dataset) > 0:
-    dataset_embeddings = embedder.encode(dataset["question"], convert_to_tensor=True)
+    dataset_embeddings = embed_model.encode(dataset["question"], convert_to_tensor=True)
 else:
-    dataset_embeddings = torch.empty((0, embedder.get_sentence_embedding_dimension()))
+    dataset_embeddings = None
 
-# ---------------------------
-# USER RATE LIMITING
-# ---------------------------
-user_limits = {}
+# =======================
+# Rate limiting
+# =======================
+user_queries = defaultdict(list)  # {session_id: [timestamps]}
 
 def check_rate_limit(session_id):
-    current_time = time.time()
-    if session_id not in user_limits:
-        user_limits[session_id] = {"count": 0, "start_time": current_time}
-    
-    user_data = user_limits[session_id]
-    if current_time - user_data["start_time"] > TIME_WINDOW:
-        user_data["count"] = 0
-        user_data["start_time"] = current_time
-    
-    if user_data["count"] >= MAX_QUESTIONS:
-        return False, f"You have reached the max of {MAX_QUESTIONS} questions. Please wait before asking more."
-    
-    user_data["count"] += 1
-    return True, None
-
-# ---------------------------
-# HELPER FUNCTIONS
-# ---------------------------
-def find_similar_answer(user_input):
-    if len(dataset) == 0:
+    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)
+    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:
         return None
-    
-    query_emb = embedder.encode(user_input, convert_to_tensor=True)
-    scores = util.cos_sim(query_emb, dataset_embeddings)
-    top_idx = torch.argmax(scores)
-    top_score = scores[0][top_idx].item()
-    
-    if top_score > 0.6:  # threshold for similarity
-        return dataset["answer"][top_idx]
-    return None
+    return dataset["answer"][top_idx]
 
+# =======================
+# Save Q&A to dataset
+# =======================
 def save_qna(question, answer):
     global dataset, dataset_embeddings
     new_entry = Dataset.from_dict({"question": [question], "answer": [answer]})
@@ -78,53 +68,53 @@ def save_qna(question, answer):
         "question": dataset["question"] + new_entry["question"],
         "answer": dataset["answer"] + new_entry["answer"]
     })
-    
-    # update embeddings incrementally
-    new_emb = embedder.encode([question], convert_to_tensor=True)
-    if len(dataset_embeddings) == 0:
-        dataset_embeddings = new_emb
-    else:
-        dataset_embeddings = torch.vstack([dataset_embeddings, new_emb])
-    
-    # save to HF dataset (push to hub)
+    dataset_embeddings = embed_model.encode(dataset["question"], convert_to_tensor=True)
     dataset.push_to_hub(DATASET_NAME, token=HF_TOKEN)
 
-# ---------------------------
-# MAIN CHAT FUNCTION
-# ---------------------------
-def chat(history, user_input, session_id="default"):
-    # Rate limit check
-    allowed, message = check_rate_limit(session_id)
+# =======================
+# 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)
     if not allowed:
-        history.append(("System", message))
-        return history, history
-    
-    # Check existing similar Q&A
-    response = find_similar_answer(user_input)
-    
-    if not response:
-        # Fallback / simple generative response
-        response = f"Guardian AI: Sorry, I don’t know the answer yet. I’m learning!"
-    
-    # Save new Q&A for incremental learning
-    save_qna(user_input, response)
-    
-    # Update chat history
-    history.append((user_input, response))
-    return history, history
-
-# ---------------------------
-# GRADIO INTERFACE
-# ---------------------------
+        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()
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Your question")
-    session_state = gr.State("default")  # default session
-    
-    def user_submit(message, history, session_id):
-        return chat(history, message, session_id)
-    
-    msg.submit(user_submit, inputs=[msg, chatbot, session_state], outputs=[chatbot, chatbot])
-
-# Launch app
-app.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    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()
+
+    msg.submit(chat, inputs=[chatbot, msg, session_id], outputs=[chatbot])
+    clear.click(lambda: [], None, chatbot)
+
+app.launch()