track IPs for last min proompters

This commit is contained in:
Cyberes 2023-08-22 23:37:39 -06:00
parent 61b9e313d2
commit 3bb27d6900
2 changed files with 10 additions and 10 deletions

View File

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

View File

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