Update app.py

app.py

@@ -2,13 +2,64 @@ import gradio as gr
 import requests
 import json
 import os
+import time
+from collections import defaultdict
+from PIL import Image
+import io
 
 BASE_URL = "https://api.jigsawstack.com/v1"
 headers = {
     "x-api-key": os.getenv("JIGSAWSTACK_API_KEY")
 }
 
-def web_ai_search(query, ai_overview, safe_search, spell_check, deep_research, max_depth, max_breadth, max_output_tokens, target_output_tokens):
+# Rate limiting configuration
+request_times = defaultdict(list)
+MAX_REQUESTS = 1  # Maximum requests per time window
+TIME_WINDOW = 3600   # Time window in seconds (1 hour)
+
+def get_real_ip(request: gr.Request):
+    """Extract real IP address using x-forwarded-for header or fallback"""
+    if not request:
+        return "unknown"
+    
+    forwarded = request.headers.get("x-forwarded-for")
+    if forwarded:
+        ip = forwarded.split(",")[0].strip()  # First IP in the list is the client's
+    else:
+        ip = request.client.host  # fallback
+    return ip
+
+def check_rate_limit(request: gr.Request):
+    """Check if the current request exceeds rate limits"""
+    if not request:
+        return True, "Rate limit check failed - no request info"
+    
+    ip = get_real_ip(request)
+    now = time.time()
+
+    # Clean up old timestamps outside the time window
+    request_times[ip] = [t for t in request_times[ip] if now - t < TIME_WINDOW]
+    
+
+    # Check if rate limit exceeded
+    if len(request_times[ip]) >= MAX_REQUESTS:
+        time_remaining = int(TIME_WINDOW - (now - request_times[ip][0]))
+        time_remaining_minutes = round(time_remaining / 60, 1)      
+        time_window_minutes = round(TIME_WINDOW / 60, 1)
+        
+        return False, f"Rate limit exceeded. You can make {MAX_REQUESTS} requests per {time_window_minutes} minutes. Try again in {time_remaining_minutes} minutes."
+    
+    # Add current request timestamp
+    request_times[ip].append(now)
+    return True, ""
+
+# ----------------- JigsawStack API Wrappers ------------------
+
+def web_ai_search(query, ai_overview, safe_search, spell_check, deep_research, max_depth, max_breadth, max_output_tokens, target_output_tokens, request: gr.Request):
+    rate_limit_ok, rate_limit_msg = check_rate_limit(request)
+    if not rate_limit_ok:
+        return f"❌ {rate_limit_msg}", "", [], ""
+    
     if not query or not query.strip():
         return "❌ Please enter a search query.", "", [], ""
     payload = {
@@ -81,8 +132,8 @@ with gr.Blocks() as demo:
     def toggle_deep_research(checked):
             return {deep_research_group: gr.update(visible=checked)}
     deep_research.change(fn=toggle_deep_research, inputs=deep_research, outputs=deep_research_group)
-    def on_search(query, overview, safe, spell, deep, d_depth, d_breadth, d_tokens, d_target):
-            return web_ai_search(query, overview, safe, spell, deep, d_depth, d_breadth, d_tokens, d_target)
+    def on_search(query, overview, safe, spell, deep, d_depth, d_breadth, d_tokens, d_target, request: gr.Request):
+        return web_ai_search(query, overview, safe, spell, deep, d_depth, d_breadth, d_tokens, d_target, request)
     search_btn.click(fn=on_search, inputs=[search_query, ai_overview, safe_search, spell_check, deep_research, max_depth, max_breadth, max_output_tokens, target_output_tokens],
                         outputs=[search_status, overview_box, results_box, search_json_output])
     def clear_search():