gradio rate limiting

src/know_lang_bot/chat_bot/chat_interface.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from know_lang_bot.config import AppConfig
 from know_lang_bot.utils.fancy_log import FancyLogger
+from know_lang_bot.utils.rate_limiter import RateLimiter
 from know_lang_bot.chat_bot.chat_graph import stream_chat_progress, ChatStatus
 import chromadb
 from typing import List, Dict, AsyncGenerator
@@ -15,6 +16,7 @@ class CodeQAChatInterface:
         self.config = config
         self._init_chroma()
         self.codebase_dir = Path(config.db.codebase_directory)
+        self.rate_limiter = RateLimiter()
         
     def _init_chroma(self):
         """Initialize ChromaDB connection"""
@@ -57,12 +59,32 @@ class CodeQAChatInterface:
     async def stream_response(
         self,
         message: str,
-        history: List[ChatMessage]
+        history: List[ChatMessage],
+        request: gr.Request, # gradio injects the request object
     ) -> AsyncGenerator[List[ChatMessage], None]:
         """Stream chat responses with progress updates"""
         # Add user message
         history.append(ChatMessage(role="user", content=message))
         yield history
+
+        # Check rate limit before processing
+        client_ip : str = request.request.client.host
+        print(f"Client IP: {client_ip}")
+        if self.rate_limiter.check_rate_limit(client_ip):
+            wait_time = self.rate_limiter.get_remaining_time(client_ip)
+            rate_limit_message = (
+                f"Rate limit exceeded. Please wait {wait_time:.0f} seconds before sending another message."
+            )
+            history.append(ChatMessage(
+                role="assistant",
+                content=rate_limit_message,
+                metadata={
+                    "title": "⚠️ Rate Limit Warning",
+                    "status": "done"
+                }
+            ))
+            yield history
+            return
         
         current_progress: ChatMessage | None = None
         code_blocks_added = False
@@ -140,8 +162,8 @@ class CodeQAChatInterface:
                 submit = gr.Button("Submit", scale=1)
                 clear = gr.ClearButton([msg, chatbot], scale=1)
 
-            async def respond(message: str, history: List[ChatMessage]) -> AsyncGenerator[List[ChatMessage], None]:
-                async for updated_history in self.stream_response(message, history):
+            async def respond(message: str, history: List[ChatMessage], request: gr.Request) -> AsyncGenerator[List[ChatMessage], None]:
+                async for updated_history in self.stream_response(message, history, request):
                     yield updated_history
                     
             # Set up event handlers

src/know_lang_bot/utils/rate_limiter.py

@@ -0,0 +1,53 @@
+from dataclasses import dataclass, field
+from collections import defaultdict
+import time
+from typing import Dict, List
+from threading import Lock
+
+
+@dataclass
+class RateLimiter:
+    """Rate limiter implementation for chat interface"""
+    requests_per_minute: int = 2
+    window_size: int = 60  # seconds
+    
+    # Track request timestamps per user IP
+    _requests: Dict[str, List[float]] = field(default_factory=lambda: defaultdict(list))
+    _lock: Lock = field(default_factory=Lock)
+
+    def _clean_old_requests(self, user_ip: str) -> None:
+        """Remove requests older than the window size"""
+        current_time = time.time()
+        with self._lock:
+            self._requests[user_ip] = [
+                timestamp for timestamp in self._requests[user_ip]
+                if current_time - timestamp < self.window_size
+            ]
+
+    def check_rate_limit(self, user_ip: str) -> bool:
+        """
+        Check if user has exceeded rate limit
+        Returns True if rate limit exceeded, False otherwise
+        """
+        self._clean_old_requests(user_ip)
+        
+        with self._lock:
+            if len(self._requests[user_ip]) >= self.requests_per_minute:
+                return True
+                
+            self._requests[user_ip].append(time.time())
+            return False
+
+    def get_remaining_time(self, user_ip: str) -> float:
+        """Get remaining time until next request is allowed"""
+        self._clean_old_requests(user_ip)
+        
+        with self._lock:
+            if not self._requests[user_ip]:
+                return 0
+                
+            oldest_request = min(self._requests[user_ip])
+            current_time = time.time()
+            time_until_reset = self.window_size - (current_time - oldest_request)
+            
+            return max(0, time_until_reset)