track IPs for last min proompters

llm_server/routes/stats.py

@@ -15,21 +15,16 @@ start_time = datetime.now()
 
 class SemaphoreCheckerThread(Thread):
     proompters_1_min = 0
+    recent_prompters = {}
 
     def __init__(self, semaphore):
         Thread.__init__(self)
         self.semaphore = semaphore
-        self.values = collections.deque(maxlen=60)
-        self.prev_semaphore_value = self.semaphore._value
         self.daemon = True
 
     def run(self):
         while True:
-            current_semaphore_value = self.semaphore._value
-            if current_semaphore_value < opts.concurrent_gens and current_semaphore_value != self.prev_semaphore_value:
-                self.values.append(1)
-            else:
-                self.values.append(0)
-            self.prev_semaphore_value = current_semaphore_value
-            SemaphoreCheckerThread.proompters_1_min = sum(self.values)
+            current_time = time.time()
+            SemaphoreCheckerThread.recent_prompters = {ip: timestamp for ip, timestamp in SemaphoreCheckerThread.recent_prompters.items() if current_time - timestamp <= 60}
+            SemaphoreCheckerThread.proompters_1_min = len(SemaphoreCheckerThread.recent_prompters)
             time.sleep(1)

llm_server/routes/v1/generate.py

@@ -1,6 +1,8 @@
+import time
+
 from flask import jsonify, request
 
-from llm_server.routes.stats import concurrent_semaphore, proompts
+from llm_server.routes.stats import SemaphoreCheckerThread, concurrent_semaphore, proompts
 from . import bp
 from ..helpers.client import format_sillytavern_err
 from ..helpers.http import cache_control, validate_json
@@ -34,6 +36,9 @@ def generate():
             client_ip = request.headers.get('x-forwarded-for')
         else:
             client_ip = request.remote_addr
+
+        SemaphoreCheckerThread.recent_prompters[client_ip] = time.time()
+
         token = request.headers.get('X-Api-Key')
 
         parameters = request_json_body.copy()