IP limit testing (set 1 per day)

app.py

@@ -4,12 +4,22 @@ import time
 import html
 import re
 import traceback
+import datetime
+import threading
+from collections import defaultdict
+from flask import request
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
 client = InferenceClient("Trinoid/Data_Management")
 
+# Rate limiting settings
+MAX_REQUESTS_PER_DAY = 1  # Maximum number of requests per IP per day (set to 1 for testing)
+ip_request_counters = defaultdict(int)  # Tracks request count per IP
+ip_last_reset = {}  # Tracks when counters were last reset for each IP
+rate_limit_lock = threading.Lock()  # Lock for thread-safe counter access
+
 # Comprehensive patterns to filter out thinking and meta-commentary
 THINKING_PATTERNS = [
     r"Okay, so I('m| am) (trying to|going to|attempting to)",
@@ -32,6 +42,39 @@ THINKING_PATTERNS = [
     r"I'll approach this by",
 ]
 
+def get_client_ip():
+    """Get the client's IP address from the request"""
+    try:
+        # Try to get IP from request context in a Flask environment
+        if request:
+            return request.remote_addr
+    except:
+        pass
+    
+    # Fallback if we can't get a real IP
+    return "127.0.0.1"
+
+def check_rate_limit():
+    """Check if the current IP has exceeded its daily limit"""
+    current_ip = get_client_ip()
+    current_date = datetime.datetime.now().date()
+    
+    with rate_limit_lock:
+        # Reset counter if it's a new day
+        if current_ip in ip_last_reset and ip_last_reset[current_ip] != current_date:
+            ip_request_counters[current_ip] = 0
+        
+        # Update last reset date
+        ip_last_reset[current_ip] = current_date
+        
+        # Check if limit is exceeded
+        if ip_request_counters[current_ip] >= MAX_REQUESTS_PER_DAY:
+            return False
+        
+        # Increment counter
+        ip_request_counters[current_ip] += 1
+        return True
+
 def process_final_response(response_text):
     """Comprehensive processing of the final response to ensure quality"""
     
@@ -107,6 +150,23 @@ def respond(
     temperature,
     top_p,
 ):
+    # Check rate limit before processing the request
+    if not check_rate_limit():
+        current_ip = get_client_ip()
+        next_reset = (datetime.datetime.now() + datetime.timedelta(days=1)).replace(hour=0, minute=0, second=0)
+        hours_until_reset = int((next_reset - datetime.datetime.now()).total_seconds() / 3600)
+        
+        limit_message = f"""
+<div class="rate-limit-warning">
+    <h3>Daily Request Limit Reached</h3>
+    <p>You've reached the maximum of {MAX_REQUESTS_PER_DAY} requests allowed per day.</p>
+    <p>Your limit will reset in approximately {hours_until_reset} hours (at midnight your local time).</p>
+    <p>Please try again tomorrow. Thank you for your understanding!</p>
+</div>
+"""
+        yield limit_message
+        return
+    
     # Create a more effective system prompt
     enhanced_system_message = f"""You are an expert in Microsoft 365 services including SharePoint, OneDrive, Teams, and the Microsoft 365 compliance ecosystem.
 
@@ -202,6 +262,30 @@ If comparing services:
         error_msg = f"I apologize, but I encountered an error while generating a response. Error details: {str(e)}"
         yield error_msg
 
+# Function to get admin stats about rate limiting
+def get_rate_limit_stats():
+    stats = {
+        "max_requests_per_day": MAX_REQUESTS_PER_DAY,
+        "active_ips": len(ip_request_counters),
+        "ip_usage": dict(ip_request_counters),
+        "last_reset_dates": {ip: date.strftime("%Y-%m-%d") for ip, date in ip_last_reset.items()}
+    }
+    return stats
+
+# Admin interface to adjust rate limit settings
+def update_rate_limit(new_limit):
+    global MAX_REQUESTS_PER_DAY
+    
+    if new_limit < 1:
+        return {"success": False, "message": "Limit must be at least 1"}
+    
+    old_limit = MAX_REQUESTS_PER_DAY
+    MAX_REQUESTS_PER_DAY = new_limit
+    
+    return {
+        "success": True, 
+        "message": f"Rate limit updated from {old_limit} to {new_limit} requests per IP per day"
+    }
 
 # Custom CSS for Plant Wisdom.AI styling
 custom_css = """
@@ -309,12 +393,25 @@ summary:hover {
 .thinking-step:last-child {
     border-bottom: none;
 }
+
+/* Rate limit warning styling */
+.rate-limit-warning {
+    background-color: #fff3cd;
+    color: #856404;
+    border: 1px solid #ffeeba;
+    border-radius: 8px;
+    padding: 16px;
+    margin: 16px 0;
+    text-align: center;
+    font-weight: 500;
+}
 """
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
 """
-demo = gr.ChatInterface(
+# Main chat interface
+chat_interface = gr.ChatInterface(
     respond,
     title="AI Data Management Expert",
     description="Hello! I am your Data Management Expert, specialized in Microsoft 365. I'm here to help you with guidance on Data Management procedures. How can I assist you today?",
@@ -412,6 +509,39 @@ IMPORTANT: If a question has been asked before in the conversation, acknowledge
     ],
 )
 
+# Create Gradio Blocks app with both interfaces
+with gr.Blocks(theme=gr.themes.Base()) as demo:
+    # Main chat interface
+    chat_interface.render()
+    
+    # Hidden admin panel (accessible at /admin)
+    with gr.Tab("Admin", visible=False):
+        gr.Markdown("## Rate Limiting Settings")
+        
+        with gr.Row():
+            current_limit = gr.Number(label="Current Request Limit Per IP Per Day", value=MAX_REQUESTS_PER_DAY, interactive=False)
+            new_limit = gr.Number(label="New Request Limit", value=MAX_REQUESTS_PER_DAY, minimum=1, step=1)
+            update_btn = gr.Button("Update Limit")
+        
+        result_text = gr.Textbox(label="Result", interactive=False)
+        
+        # Stats display
+        with gr.Accordion("Usage Statistics", open=False):
+            refresh_btn = gr.Button("Refresh Stats")
+            stats_json = gr.JSON(label="Current Usage Stats")
+        
+        # Connect buttons to functions
+        update_btn.click(
+            fn=update_rate_limit, 
+            inputs=[new_limit],
+            outputs=[result_text, current_limit]
+        )
+        
+        refresh_btn.click(
+            fn=get_rate_limit_stats,
+            inputs=[],
+            outputs=[stats_json]
+        )
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0")