From c9c1f42e44484074cc50ef20583b24f17bf5550a Mon Sep 17 00:00:00 2001
From: awwaawwa <8493196+awwaawwa@users.noreply.github.com>
Date: Tue, 24 Dec 2024 01:29:58 +0800
Subject: [PATCH] feat(translator): RateLimiter

---
 pdf2zh/translator.py    | 43 +++++++++++++++++++++++++++++++++++++++++
 test/test_translator.py | 38 +++++++++++++++++++++++++++++++++++-
 2 files changed, 80 insertions(+), 1 deletion(-)

diff --git a/pdf2zh/translator.py b/pdf2zh/translator.py
index 2556a3a9..f8558b54 100644
--- a/pdf2zh/translator.py
+++ b/pdf2zh/translator.py
@@ -2,6 +2,7 @@
 import logging
 import os
 import re
+import time
 import unicodedata
 from copy import copy
 import deepl
@@ -24,6 +25,47 @@ def remove_control_characters(s):
     return "".join(ch for ch in s if unicodedata.category(ch)[0] != "C")
 
 
+class RateLimiter:
+    def __init__(self, max_qps: int):
+        self.max_qps = max_qps
+        self.min_interval = 1.0 / max_qps
+        self.last_requests = []  # Track last N requests
+        self.window_size = max_qps  # Track requests in a sliding window
+        self.lock = asyncio.Lock()
+    
+    async def wait_async(self):
+        async with self.lock:
+            now = time.time()
+            
+            # Clean up old requests outside the 1-second window
+            while self.last_requests and now - self.last_requests[0] > 1.0:
+                self.last_requests.pop(0)
+            
+            # If we have less than max_qps requests in the last second, allow immediately
+            if len(self.last_requests) < self.max_qps:
+                self.last_requests.append(now)
+                return
+            
+            # Otherwise, wait until we can make the next request
+            next_time = self.last_requests[0] + 1.0
+            if next_time > now:
+                await asyncio.sleep(next_time - now)
+            self.last_requests.pop(0)
+            self.last_requests.append(next_time)
+
+    def set_max_qps(self, max_qps):
+        self.max_qps = max_qps
+        self.min_interval = 1.0 / max_qps
+        self.window_size = max_qps
+
+
+_translate_rate_limiter = RateLimiter(5)
+
+
+def set_translate_rate_limiter(max_qps):
+    _translate_rate_limiter.set_max_qps(max_qps)
+
+
 class BaseTranslator:
     name = "base"
     envs = {}
@@ -77,6 +119,7 @@ async def translate_async(self, text, ignore_cache=False):
             cache = self.cache.get(text)
             if cache is not None:
                 return cache
+        await _translate_rate_limiter.wait_async()
         try:
             translation = await self.do_translate_async(text)
         except NotImplementedError:
diff --git a/test/test_translator.py b/test/test_translator.py
index e0a865ec..9e30eaf3 100644
--- a/test/test_translator.py
+++ b/test/test_translator.py
@@ -1,7 +1,9 @@
 import unittest
 from pdf2zh.translator import BaseTranslator
 from pdf2zh import cache
-
+from pdf2zh.translator import RateLimiter
+import asyncio
+import time
 
 class AutoIncreaseTranslator(BaseTranslator):
     name = "auto_increase"
@@ -90,5 +92,39 @@ async def test_call_sync_from_async(self):
         self.assertEqual(await sync_translator.translate_async("Hello World"), "1")
 
 
+class TestRateLimiter(unittest.IsolatedAsyncioTestCase):
+    async def test_concurrent_rate_limit(self):
+        limiter = RateLimiter(10)  # 10 QPS
+        start_time = time.time()
+
+        async def task():
+            await limiter.wait_async()
+            return time.time()
+
+        # Run 20 concurrent tasks
+        tasks = [task() for _ in range(20)]
+        timestamps = await asyncio.gather(*tasks)
+
+        # Verify timing
+        total_time = timestamps[-1] - start_time
+        self.assertGreaterEqual(total_time, 1.0)  # Should take at least 1s for 20 requests at 10 QPS
+
+        # Check even distribution
+        intervals = [timestamps[i + 1] - timestamps[i] for i in range(len(timestamps) - 1)]
+        avg_interval = sum(intervals) / len(intervals)
+        self.assertAlmostEqual(avg_interval, 0.1, delta=0.05)  # Should be close to 0.1s (1/10 QPS)
+
+    async def test_burst_handling(self):
+        limiter = RateLimiter(10)  # 10 QPS
+
+        # First burst of 5 requests should be immediate
+        start = time.time()
+        tasks = [limiter.wait_async() for _ in range(5)]
+        await asyncio.gather(*tasks)
+        burst_time = time.time() - start
+
+        self.assertLess(burst_time, 0.1)  # Should complete quickly
+
+
 if __name__ == "__main__":
     unittest.main()