Add --req-time-cost

libretranslate/app.py

@@ -1,4 +1,5 @@
 import io
+import math
 import os
 import re
 import tempfile
@@ -219,6 +220,13 @@ def create_app(args):
 
         from flask_limiter import Limiter
 
+        def limits_cost():
+          req_cost = getattr(request, 'req_cost', 1)
+          if args.req_time_cost > 0:
+            return max(req_cost, int(math.ceil(getattr(request, 'duration', 0) / args.req_time_cost)))
+          else:
+            return req_cost
+
         limiter = Limiter(
             key_func=get_remote_address,
             default_limits=get_routes_limits(
@@ -226,7 +234,7 @@ def create_app(args):
             ),
             storage_uri=args.req_limit_storage,
             default_limits_deduct_when=lambda req: True, # Force cost to be called after the request
-            default_limits_cost=lambda: getattr(request, 'req_cost', 1)
+            default_limits_cost=limits_cost
         )
     else:
         from .no_limiter import Limiter
@@ -325,12 +333,19 @@ def create_app(args):
                 status = e.code
                 raise e
               finally:
-                duration = max(default_timer() - start_t, 0)
+                request.duration = max(default_timer() - start_t, 0)
-                measure_request.labels(request.path, status, ip, ak).observe(duration)
+                measure_request.labels(request.path, status, ip, ak).observe(request.duration)
                 g.dec()
           return measure_func
         else:
-          return func
+          @wraps(func)
+          def time_func(*a, **kw):
+            start_t = default_timer()
+            try:
+              return func(*a, **kw)
+            finally:
+              request.duration = max(default_timer() - start_t, 0)
+          return time_func
 
     @bp.errorhandler(400)
     def invalid_api(e):

libretranslate/default_values.py

@@ -81,6 +81,11 @@ _default_options_objects = [
         'default_value': -1,
         'value_type': 'int'
     },
+    {
+        'name': 'REQ_TIME_COST',
+        'default_value': -1,
+        'value_type': 'int'
+    },
     {
         'name': 'BATCH_LIMIT',
         'default_value': -1,

libretranslate/main.py

@@ -63,6 +63,13 @@ def get_args():
         metavar="<number>",
         help="Set the maximum number of request limit offences that a client can exceed before being banned. (%(default)s)",
     )
+    parser.add_argument(
+        "--req-time-cost",
+        default=DEFARGS['REQ_TIME_COST'],
+        type=int,
+        metavar="<number>",
+        help="Considers a time cost (in seconds) for request limiting purposes. If a request takes 10 seconds and this value is set to 5, the request cost is either 2 or the actual request cost (whichever is greater). (%(default)s)",
+    )
     parser.add_argument(
         "--batch-limit",
         default=DEFARGS['BATCH_LIMIT'],