Merge pull request #574 from pierotofy/charlimit

Granular API key char limit support

README.md

@@ -247,7 +247,7 @@ helm install libretranslate libretranslate/libretranslate --namespace libretrans
 
 ## Manage API Keys
 
-LibreTranslate supports per-user limit quotas, e.g. you can issue API keys to users so that they can enjoy higher requests limits per minute (if you also set `--req-limit`). By default all users are rate-limited based on `--req-limit`, but passing an optional `api_key` parameter to the REST endpoints allows a user to enjoy higher request limits.
+LibreTranslate supports per-user limit quotas, e.g. you can issue API keys to users so that they can enjoy higher requests limits per minute (if you also set `--req-limit`). By default all users are rate-limited based on `--req-limit`, but passing an optional `api_key` parameter to the REST endpoints allows a user to enjoy higher request limits. You can also specify different character limits that bypass the default `--char-limit` value on a per-key basis.
 
 To use API keys simply start LibreTranslate with the `--api-keys` option. If you modified the API keys database path with the option `--api-keys-db-path`, you must specify the path with the same argument flag when using the `ltmanage keys` command.
 
@@ -259,6 +259,12 @@ To issue a new API key with 120 requests per minute limits:
 ltmanage keys add 120
 ```
 
+To issue a new API key with 120 requests per minute and a maximum of 5,000 characters per request:
+
+```bash
+ltmanage keys add 120 --char-limit 5000
+```
+
 If you changed the API keys database path:
 
 ```bash

VERSION

@@ -1 +1 @@
-1.5.4
+1.5.5

libretranslate/api_keys.py

@@ -32,41 +32,49 @@ class Database:
             """CREATE TABLE IF NOT EXISTS api_keys (
             "api_key"	TEXT NOT NULL,
             "req_limit"	INTEGER NOT NULL,
+            "char_limit" INTEGER DEFAULT NULL,
             PRIMARY KEY("api_key")
         );"""
         )
 
+        # Schema/upgrade checks
+        schema = self.c.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='api_keys';").fetchone()[0]
+        if '"char_limit" INTEGER DEFAULT NULL' not in schema:
+            self.c.execute('ALTER TABLE api_keys ADD COLUMN "char_limit" INTEGER DEFAULT NULL;')
+
     def lookup(self, api_key):
-        req_limit = self.cache.get(api_key)
+        val = self.cache.get(api_key)
-        if req_limit is None:
+        if val is None:
             # DB Lookup
             stmt = self.c.execute(
-                "SELECT req_limit FROM api_keys WHERE api_key = ?", (api_key,)
+                "SELECT req_limit, char_limit FROM api_keys WHERE api_key = ?", (api_key,)
             )
             row = stmt.fetchone()
             if row is not None:
-                self.cache[api_key] = row[0]
+                self.cache[api_key] = row
-                req_limit = row[0]
+                val = row
             else:
                 self.cache[api_key] = False
-                req_limit = False
+                val = False
 
-        if isinstance(req_limit, bool):
+        if isinstance(val, bool):
-            req_limit = None
+            val = None
 
-        return req_limit
+        return val
 
-    def add(self, req_limit, api_key="auto"):
+    def add(self, req_limit, api_key="auto", char_limit=None):
         if api_key == "auto":
             api_key = str(uuid.uuid4())
+        if char_limit == 0:
+            char_limit = None
 
         self.remove(api_key)
         self.c.execute(
-            "INSERT INTO api_keys (api_key, req_limit) VALUES (?, ?)",
+            "INSERT INTO api_keys (api_key, req_limit, char_limit) VALUES (?, ?, ?)",
-            (api_key, req_limit),
+            (api_key, req_limit, char_limit),
         )
         self.c.commit()
-        return (api_key, req_limit)
+        return (api_key, req_limit, char_limit)
 
     def remove(self, api_key):
         self.c.execute("DELETE FROM api_keys WHERE api_key = ?", (api_key,))
@@ -74,7 +82,7 @@ class Database:
         return api_key
 
     def all(self):
-        row = self.c.execute("SELECT api_key, req_limit FROM api_keys")
+        row = self.c.execute("SELECT api_key, req_limit, char_limit FROM api_keys")
         return row.fetchall()
 
 
@@ -84,8 +92,8 @@ class RemoteDatabase:
         self.cache = ExpiringDict(max_len=max_cache_len, max_age_seconds=max_cache_age)
 
     def lookup(self, api_key):
-        req_limit = self.cache.get(api_key)
+        val = self.cache.get(api_key)
-        if req_limit is None:
+        if val is None:
             try:
                 r = requests.post(self.url, data={'api_key': api_key}, timeout=60)
                 res = r.json()
@@ -94,6 +102,8 @@ class RemoteDatabase:
                 return None
 
             req_limit = res.get('req_limit', None) if res.get('error', None) is None else None
-            self.cache[api_key] = req_limit
+            char_limit = res.get('char_limit', None) if res.get('error', None) is None else None
 
-        return req_limit
+            self.cache[api_key] = (req_limit, char_limit)
+
+        return val

libretranslate/app.py

@@ -96,13 +96,28 @@ def get_req_limits(default_limit, api_keys_db, db_multiplier=1, multiplier=1):
         api_key = get_req_api_key()
 
         if api_key:
-            db_req_limit = api_keys_db.lookup(api_key)
+            api_key_limits = api_keys_db.lookup(api_key)
-            if db_req_limit is not None:
+            if api_key_limits is not None:
-                req_limit = db_req_limit * db_multiplier
+                req_limit = api_key_limits[0] * db_multiplier
 
     return int(req_limit * multiplier)
 
 
+def get_char_limit(default_limit, api_keys_db):
+    char_limit = default_limit
+
+    if api_keys_db:
+        api_key = get_req_api_key()
+
+        if api_key:
+            api_key_limits = api_keys_db.lookup(api_key)
+            if api_key_limits is not None:
+                if api_key_limits[1] is not None:
+                    char_limit = api_key_limits[1]
+
+    return char_limit
+
+
 def get_routes_limits(args, api_keys_db):
     default_req_limit = args.req_limit
     if default_req_limit == -1:
@@ -547,6 +562,8 @@ def create_app(args):
             # https://www.rfc-editor.org/rfc/rfc2046#section-4.1.1
             q = "\n".join(q.splitlines())
 
+        char_limit = get_char_limit(args.char_limit, api_keys_db)
+
         batch = isinstance(q, list)
 
         if batch and args.batch_limit != -1:
@@ -559,12 +576,12 @@ def create_app(args):
 
         src_texts = q if batch else [q]
 
-        if args.char_limit != -1:
+        if char_limit != -1:
             for text in src_texts:
-                if len(text) > args.char_limit:
+                if len(text) > char_limit:
                     abort(
                         400,
-                        description=_("Invalid request: request (%(size)s) exceeds text limit (%(limit)s)", size=len(text), limit=args.char_limit),
+                        description=_("Invalid request: request (%(size)s) exceeds text limit (%(limit)s)", size=len(text), limit=char_limit),
                     )
 
         if batch:
@@ -736,6 +753,7 @@ def create_app(args):
         source_lang = request.form.get("source")
         target_lang = request.form.get("target")
         file = request.files['file']
+        char_limit = get_char_limit(args.char_limit, api_keys_db)
 
         if not file:
             abort(400, description=_("Invalid request: missing %(name)s parameter", name='file'))
@@ -771,8 +789,8 @@ def create_app(args):
             # set the cost of the request to N = bytes / char_limit, which is
             # roughly equivalent to a batch process of N batches assuming
             # each batch uses all available limits
-            if args.char_limit > 0:
+            if char_limit > 0:
-                request.req_cost = max(1, int(os.path.getsize(filepath) / args.char_limit))
+                request.req_cost = max(1, int(os.path.getsize(filepath) / char_limit))
 
             translated_file_path = argostranslatefiles.translate_file(src_lang.get_translation(tgt_lang), filepath)
             translated_filename = os.path.basename(translated_file_path)

libretranslate/manage.py

@@ -29,6 +29,9 @@ def manage():
     keys_add_parser.add_argument(
         "--key", type=str, default="auto", required=False, help="API Key"
     )
+    keys_add_parser.add_argument(
+        "--char-limit", type=int, default=0, required=False, help="Character limit"
+    )
 
     keys_remove_parser = keys_subparser.add_parser(
         "remove", help="Remove API keys to database"
@@ -52,7 +55,7 @@ def manage():
                     print("{}: {}".format(*item))
 
         elif args.sub_command == "add":
-            print(db.add(args.req_limit, args.key)[0])
+            print(db.add(args.req_limit, args.key, args.char_limit)[0])
         elif args.sub_command == "remove":
             print(db.remove(args.key))
     else: