From 3bb27d69003dc94e83b8329439cdae1e80594141 Mon Sep 17 00:00:00 2001
From: Cyberes <cyberes@evulid.cc>
Date: Tue, 22 Aug 2023 23:37:39 -0600
Subject: [PATCH] track IPs for last min proompters

---
 llm_server/routes/stats.py       | 13 ++++---------
 llm_server/routes/v1/generate.py |  7 ++++++-
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/llm_server/routes/stats.py b/llm_server/routes/stats.py
index 91c358a..47077e0 100644
--- a/llm_server/routes/stats.py
+++ b/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)
diff --git a/llm_server/routes/v1/generate.py b/llm_server/routes/v1/generate.py
index b1af72d..dc50c04 100644
--- a/llm_server/routes/v1/generate.py
+++ b/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()