Add IP-based throttle

app.py

@@ -3,6 +3,7 @@ import datetime
 import json
 import os
 import requests
+import time
 from constants import *
 
 API_IPADDR = os.environ.get('API_IPADDR', None)
@@ -11,17 +12,31 @@ max_size = os.environ.get('max_size', 100)
 max_threads = os.environ.get('max_threads', 40)
 debug = (os.environ.get('debug', 'False') != 'False')
 
-def process(corpus_desc, query_desc, query):
+last_query_time_by_ip = {}
+
+def process(corpus_desc, query_desc, query, ret_num, request: gr.Request):
+    global last_query_time_by_ip
+    ip = request.client.host if request else ''
+    timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
+    t = time.time()
+    last_query_time = 0 if ip == '' else last_query_time_by_ip.get(ip, 0)
+    blocked = (t - last_query_time < MIN_QUERY_INTERVAL_SECONDS)
+
     corpus = CORPUS_BY_DESC[corpus_desc]
     query_type = QUERY_TYPE_BY_DESC[query_desc]
-    timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
     data = {
         'timestamp': timestamp,
+        'ip': ip,
+        'blocked': blocked,
         'corpus': corpus,
         'query_type': query_type,
         'query': query,
     }
     print(json.dumps(data))
+    if blocked:
+        return tuple([f'You queried too frequently. Please try again in {MIN_QUERY_INTERVAL_SECONDS} seconds.'] + [''] * (ret_num - 1))
+    if ip != '':
+        last_query_time_by_ip[ip] = t
     if API_IPADDR is None:
         raise ValueError(f'API_IPADDR envvar is not set!')
     response = requests.post(f'http://{API_IPADDR}:5000/', json=data)
@@ -33,18 +48,38 @@ def process(corpus_desc, query_desc, query):
         print(result)
     return result
 
-def process_ard_cnf_multi(corpus_desc, query_desc, query, maxnum):
+def process_1(corpus_desc, query_desc, query, request: gr.Request):
+    return process(corpus_desc, query_desc, query, 1, request)
+def process_2(corpus_desc, query_desc, query, request: gr.Request):
+    return process(corpus_desc, query_desc, query, 2, request)
+def process_3(corpus_desc, query_desc, query, request: gr.Request):
+    return process(corpus_desc, query_desc, query, 3, request)
+
+def process_ard_cnf_multi(corpus_desc, query_desc, query, maxnum, request: gr.Request):
+    global last_query_time_by_ip
+    ip = request.client.host if request else ''
+    timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
+    t = time.time()
+    last_query_time = 0 if ip == '' else last_query_time_by_ip.get(ip, 0)
+    blocked = (t - last_query_time < MIN_QUERY_INTERVAL_SECONDS)
+
     corpus = CORPUS_BY_DESC[corpus_desc]
     query_type = QUERY_TYPE_BY_DESC[query_desc]
     timestamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
     data = {
         'timestamp': timestamp,
+        'ip': ip,
+        'blocked': blocked,
         'corpus': corpus,
         'query_type': query_type,
         'query': query,
         'maxnum': maxnum,
     }
     print(json.dumps(data))
+    if blocked:
+        return tuple([f'You queried too frequently. Please try again in {MIN_QUERY_INTERVAL_SECONDS} seconds.'] + [''] * 11)
+    if ip != '':
+        last_query_time_by_ip[ip] = t
     if API_IPADDR is None:
         raise ValueError(f'API_IPADDR envvar is not set!')
     response = requests.post(f'http://{API_IPADDR}:5000/', json=data)
@@ -60,7 +95,7 @@ def process_ard_cnf_multi(corpus_desc, query_desc, query, maxnum):
     outputs = outputs[:maxnum]
     while len(outputs) < 10:
         outputs.append([])
-    return output_tokens, message, outputs[0], outputs[1], outputs[2], outputs[3], outputs[4], outputs[5], outputs[6], outputs[7], outputs[8], outputs[9]
+    return message, output_tokens, outputs[0], outputs[1], outputs[2], outputs[3], outputs[4], outputs[5], outputs[6], outputs[7], outputs[8], outputs[9]
 
 with gr.Blocks() as demo:
     with gr.Column():
@@ -281,14 +316,14 @@ If you find this tool useful, please kindly cite our paper:
     ard_cnf_multi_clear.add([ard_cnf_multi_input, ard_cnf_multi_output_tokens, ard_cnf_multi_output_message, ard_cnf_multi_output_0, ard_cnf_multi_output_1, ard_cnf_multi_output_2, ard_cnf_multi_output_3, ard_cnf_multi_output_4, ard_cnf_multi_output_5, ard_cnf_multi_output_6, ard_cnf_multi_output_7, ard_cnf_multi_output_8, ard_cnf_multi_output_9])
     doc_analysis_clear.add([doc_analysis_input, doc_analysis_output])
 
-    count_submit.click(process, inputs=[corpus_desc, query_desc, count_input], outputs=[count_output, count_output_tokens], api_name=False)
-    ngram_submit.click(process, inputs=[corpus_desc, query_desc, ngram_input], outputs=[ngram_output, ngram_output_tokens], api_name=False)
-    ntd_submit.click(process, inputs=[corpus_desc, query_desc, ntd_input], outputs=[ntd_output, ntd_output_tokens], api_name=False)
-    infgram_submit.click(process, inputs=[corpus_desc, query_desc, infgram_input], outputs=[infgram_output, infgram_output_tokens, infgram_longest_suffix], api_name=False)
-    infntd_submit.click(process, inputs=[corpus_desc, query_desc, infntd_input], outputs=[infntd_output, infntd_output_tokens, infntd_longest_suffix], api_name=False)
+    count_submit.click(process_2, inputs=[corpus_desc, query_desc, count_input], outputs=[count_output, count_output_tokens], api_name=False)
+    ngram_submit.click(process_2, inputs=[corpus_desc, query_desc, ngram_input], outputs=[ngram_output, ngram_output_tokens], api_name=False)
+    ntd_submit.click(process_2, inputs=[corpus_desc, query_desc, ntd_input], outputs=[ntd_output, ntd_output_tokens], api_name=False)
+    infgram_submit.click(process_3, inputs=[corpus_desc, query_desc, infgram_input], outputs=[infgram_output, infgram_output_tokens, infgram_longest_suffix], api_name=False)
+    infntd_submit.click(process_3, inputs=[corpus_desc, query_desc, infntd_input], outputs=[infntd_output, infntd_output_tokens, infntd_longest_suffix], api_name=False)
     # ard_cnf_submit.click(process, inputs=[corpus_desc, query_desc, ard_cnf_input], outputs=[ard_cnf_output, ard_cnf_output_tokens, ard_cnf_output_message], api_name=False)
-    ard_cnf_multi_submit.click(process_ard_cnf_multi, inputs=[corpus_desc, query_desc, ard_cnf_multi_input, ard_cnf_multi_maxnum], outputs=[ard_cnf_multi_output_tokens, ard_cnf_multi_output_message, ard_cnf_multi_output_0, ard_cnf_multi_output_1, ard_cnf_multi_output_2, ard_cnf_multi_output_3, ard_cnf_multi_output_4, ard_cnf_multi_output_5, ard_cnf_multi_output_6, ard_cnf_multi_output_7, ard_cnf_multi_output_8, ard_cnf_multi_output_9], api_name=False)
-    doc_analysis_submit.click(process, inputs=[corpus_desc, query_desc, doc_analysis_input], outputs=[doc_analysis_output], api_name=False)
+    ard_cnf_multi_submit.click(process_ard_cnf_multi, inputs=[corpus_desc, query_desc, ard_cnf_multi_input, ard_cnf_multi_maxnum], outputs=[ard_cnf_multi_output_message, ard_cnf_multi_output_tokens, ard_cnf_multi_output_0, ard_cnf_multi_output_1, ard_cnf_multi_output_2, ard_cnf_multi_output_3, ard_cnf_multi_output_4, ard_cnf_multi_output_5, ard_cnf_multi_output_6, ard_cnf_multi_output_7, ard_cnf_multi_output_8, ard_cnf_multi_output_9], api_name=False)
+    doc_analysis_submit.click(process_1, inputs=[corpus_desc, query_desc, doc_analysis_input], outputs=[doc_analysis_output], api_name=False)
 
     def update_query_desc(selection):
         return {

constants.py

@@ -30,3 +30,5 @@ MAX_DIFF_TOKENS = int(os.environ.get('MAX_DIFF_TOKENS', 100))
 MAX_DIFF_BYTES = 2 * MAX_DIFF_TOKENS
 MAX_CLAUSES_IN_CNF = int(os.environ.get('MAX_CLAUSES_IN_CNF', 4))
 MAX_TERMS_IN_DISJ_CLAUSE = int(os.environ.get('MAX_TERMS_IN_DISJ_CLAUSE', 4))
+
+MIN_QUERY_INTERVAL_SECONDS = int(os.environ.get('MIN_QUERY_INTERVAL_SECONDS', 5))