vectorvein 0.1.87__py3-none-any.whl → 0.1.89__py3-none-any.whl

vectorvein/utilities/rate_limiter.py

@@ -0,0 +1,312 @@
+import time
+import asyncio
+from typing import Tuple
+from collections import defaultdict
+from abc import ABC, abstractmethod
+
+
+class AsyncRateLimiterBackend(ABC):
+    """Rate Limiter Backend Abstract Base Class"""
+
+    @abstractmethod
+    async def check_limit(self, key: str, rpm: int, tpm: int, request_cost: int = 1) -> Tuple[bool, float]:
+        """Returns (allowed, wait_time)"""
+        pass
+
+
+class SyncRateLimiterBackend(ABC):
+    """Rate Limiter Backend Abstract Base Class"""
+
+    @abstractmethod
+    def check_limit(self, key: str, rpm: int, tpm: int, request_cost: int = 1) -> Tuple[bool, float]:
+        """Returns (allowed, wait_time)"""
+        pass
+
+
+class AsyncMemoryRateLimiter(AsyncRateLimiterBackend):
+    """Async Memory Rate Limiter"""
+
+    def __init__(self):
+        self.windows = defaultdict(list)
+        self.tokens = defaultdict(int)
+        self.lock = asyncio.Lock()
+
+    def _get_last_reset(self, key):
+        return self.windows[key][0] if self.windows[key] else time.time()
+
+    async def check_limit(self, key: str, rpm: int, tpm: int, request_cost: int = 1):
+        async with self.lock:
+            now = time.time()
+
+            # RPM 检查
+            window = self.windows[key]
+            window = [t for t in window if t > now - 60]
+            if len(window) >= rpm:
+                return False, 60 - (now - window[0])
+
+            # TPM 检查
+            if self.tokens[key] + request_cost > tpm:
+                return False, 60 - (now - self._get_last_reset(key))
+
+            window.append(now)
+            self.tokens[key] += request_cost
+            self.windows[key] = window[-rpm:]
+            return True, 0
+
+
+class SyncMemoryRateLimiter(SyncRateLimiterBackend):
+    """Sync Memory Rate Limiter"""
+
+    def __init__(self):
+        self.windows = defaultdict(list)
+        self.tokens = defaultdict(int)
+        self.lock = asyncio.Lock()
+
+    def _get_last_reset(self, key):
+        return self.windows[key][0] if self.windows[key] else time.time()
+
+    def check_limit(self, key: str, rpm: int, tpm: int, request_cost: int = 1) -> Tuple[bool, float]:
+        """Sync Rate Limiter Check
+
+        Args:
+            key: Rate Limiter Key
+            rpm: Requests per minute limit
+            tpm: Tokens per minute limit
+            request_cost: The number of tokens consumed by this request
+
+
+
+        Returns:
+            Tuple[bool, float]: (allowed, wait_time)
+        """
+        now = time.time()
+
+        # RPM 检查
+        window = self.windows[key]
+        window = [t for t in window if t > now - 60]
+        if len(window) >= rpm:
+            return False, 60 - (now - window[0])
+
+        # TPM 检查
+        if self.tokens[key] + request_cost > tpm:
+            return False, 60 - (now - self._get_last_reset(key))
+
+        window.append(now)
+        self.tokens[key] += request_cost
+        self.windows[key] = window[-rpm:]
+        return True, 0
+
+
+REDIS_SCRIPT = """
+local key = KEYS[1]
+local rpm = tonumber(ARGV[1])
+local tpm = tonumber(ARGV[2])
+local cost = tonumber(ARGV[3])
+
+-- 使用Redis服务器时间（精确到微秒）
+local server_time = redis.call('TIME')
+local now = tonumber(server_time[1]) + tonumber(server_time[2])/1000000
+
+-- RPM限制检查
+local rpm_key = key..'_rpm'
+local elements = redis.call('LRANGE', rpm_key, 0, -1)
+local valid_elements = {}
+local min_valid_time = now - 60
+
+-- 过滤过期时间戳
+for _, ts in ipairs(elements) do
+    local timestamp = tonumber(ts)
+    if timestamp > min_valid_time then
+        table.insert(valid_elements, timestamp)
+    end
+end
+local valid_count = #valid_elements
+
+-- 新增：自动清理过期时间戳
+if valid_count > 0 then
+    redis.call('DEL', rpm_key)
+    for _, ts in ipairs(valid_elements) do
+        redis.call('RPUSH', rpm_key, ts)
+    end
+    redis.call('EXPIRE', rpm_key, 60)
+end
+
+if valid_count >= rpm then
+    local oldest = valid_elements[valid_count]  -- 最旧的有效时间戳
+    local remaining = math.max(0.001, 60 - (now - oldest))  -- 保证最小等待时间
+    return {0, math.ceil(remaining * 1000)/1000}  -- 保留3位小数
+end
+
+-- TPM限制检查（保持不变）
+local tpm_key = key..'_tpm'
+local current_tokens = tonumber(redis.call('GET', tpm_key) or 0)
+if current_tokens + cost > tpm then
+    local ttl = redis.call('TTL', tpm_key)
+    if ttl < 0 then
+        redis.call('SETEX', tpm_key, 60, current_tokens)
+        ttl = 60
+    end
+    return {0, ttl}
+end
+
+-- 更新计数（增加时间戳覆盖写入）
+redis.call('LPUSH', rpm_key, now)
+redis.call('LTRIM', rpm_key, 0, rpm-1)
+redis.call('EXPIRE', rpm_key, 60)
+
+redis.call('INCRBY', tpm_key, cost)
+redis.call('EXPIRE', tpm_key, 60)
+
+return {1, 0}
+"""
+
+
+class AsyncRedisRateLimiter(AsyncRateLimiterBackend):
+    """Async Redis Rate Limiter"""
+
+    def __init__(self, host: str = "localhost", port: int = 6379, db: int = 0):
+        import redis.asyncio as redis
+
+        self.redis = redis.Redis(host=host, port=port, db=db)
+        self.script = self.redis.register_script(REDIS_SCRIPT)
+
+    async def check_limit(self, key: str, rpm: int, tpm: int, request_cost: int = 1):
+        result = await self.script(keys=[key], args=[rpm, tpm, request_cost])
+        allowed, wait_time = result
+        return bool(allowed), max(1, float(wait_time))
+
+
+class SyncRedisRateLimiter(SyncRateLimiterBackend):
+    """Sync Redis Rate Limiter"""
+
+    def __init__(self, host: str = "localhost", port: int = 6379, db: int = 0):
+        import redis
+
+        self.redis = redis.Redis(host=host, port=port, db=db)
+        self.script = self.redis.register_script(REDIS_SCRIPT)
+
+    def check_limit(self, key: str, rpm: int, tpm: int, request_cost: int = 1):
+        result = self.script(keys=[key], args=[rpm, tpm, request_cost])
+        allowed, wait_time = result
+        return bool(allowed), max(1, float(wait_time))
+
+
+class AsyncDiskCacheRateLimiter(AsyncRateLimiterBackend):
+    """基于 diskcache 的异步限流器实现"""
+
+    def __init__(self, cache_dir: str = ".rate_limit_cache"):
+        """初始化 diskcache 限流器
+
+        Args:
+            cache_dir: 缓存目录路径
+        """
+        from diskcache import Cache
+
+        self.cache = Cache(cache_dir)
+        self._lock = asyncio.Lock()
+
+    def _get_rpm_key(self, key: str) -> str:
+        return f"{key}_rpm"
+
+    def _get_tpm_key(self, key: str) -> str:
+        return f"{key}_tpm"
+
+    async def check_limit(self, key: str, rpm: int, tpm: int, request_cost: int = 1) -> Tuple[bool, float]:
+        """检查是否超出限流阈值
+
+        Args:
+            key: 限流键
+            rpm: 每分钟请求数限制
+            tpm: 每分钟令牌数限制
+            request_cost: 本次请求消耗的令牌数
+
+        Returns:
+            Tuple[bool, float]: (是否允许请求, 需要等待的时间)
+        """
+        async with self._lock:
+            now = time.time()
+            rpm_key = self._get_rpm_key(key)
+            tpm_key = self._get_tpm_key(key)
+
+            # RPM 检查
+            window = self.cache.get(rpm_key, []) or []
+            window = [t for t in window if t > now - 60]  # type: ignore  清理过期时间戳
+
+            if len(window) >= rpm:
+                return False, 60 - (now - window[0])  # type: ignore
+
+            # TPM 检查
+            current_tokens = self.cache.get(tpm_key, 0)
+            if current_tokens + request_cost > tpm:  # type: ignore
+                # 获取最早的请求时间
+                oldest_time = window[0] if window else now
+                return False, 60 - (now - oldest_time)  # type: ignore
+
+            # 更新状态
+            window.append(now)  # type: ignore
+            window = window[-rpm:]  # type: ignore  # 只保留最近的 rpm 个时间戳
+            self.cache.set(rpm_key, window, expire=60)
+            self.cache.set(tpm_key, current_tokens + request_cost, expire=60)  # type: ignore
+
+            return True, 0
+
+
+class SyncDiskCacheRateLimiter(SyncRateLimiterBackend):
+    """基于 diskcache 的同步限流器实现"""
+
+    def __init__(self, cache_dir: str = ".rate_limit_cache"):
+        """初始化 diskcache 限流器
+
+        Args:
+            cache_dir: 缓存目录路径
+        """
+        from diskcache import Cache
+        import threading
+
+        self.cache = Cache(cache_dir)
+        self._lock = threading.Lock()
+
+    def _get_rpm_key(self, key: str) -> str:
+        return f"{key}_rpm"
+
+    def _get_tpm_key(self, key: str) -> str:
+        return f"{key}_tpm"
+
+    def check_limit(self, key: str, rpm: int, tpm: int, request_cost: int = 1) -> Tuple[bool, float]:
+        """检查是否超出限流阈值
+
+        Args:
+            key: 限流键
+            rpm: 每分钟请求数限制
+            tpm: 每分钟令牌数限制
+            request_cost: 本次请求消耗的令牌数
+
+        Returns:
+            Tuple[bool, float]: (是否允许请求, 需要等待的时间)
+        """
+        with self._lock:
+            now = time.time()
+            rpm_key = self._get_rpm_key(key)
+            tpm_key = self._get_tpm_key(key)
+
+            # RPM 检查
+            window = self.cache.get(rpm_key, []) or []
+            window = [t for t in window if t > now - 60]  # type: ignore   清理过期时间戳
+
+            if len(window) >= rpm:
+                return False, 60 - (now - window[0])  # type: ignore
+
+            # TPM 检查
+            current_tokens = self.cache.get(tpm_key, 0)
+            if current_tokens + request_cost > tpm:  # type: ignore
+                # 获取最早的请求时间
+                oldest_time = window[0] if window else now
+                return False, 60 - (now - oldest_time)  # type: ignore
+
+            # 更新状态
+            window.append(now)
+            window = window[-rpm:]  # 只保留最近的 rpm 个时间戳
+            self.cache.set(rpm_key, window, expire=60)
+            self.cache.set(tpm_key, current_tokens + request_cost, expire=60)  # type: ignore
+
+            return True, 0

{vectorvein-0.1.87.dist-info → vectorvein-0.1.89.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vectorvein
-Version: 0.1.87
+Version: 0.1.89
 Summary: VectorVein python SDK
 Author-Email: Anderson <andersonby@163.com>
 License: MIT
@@ -14,9 +14,14 @@ Requires-Dist: Pillow>=10.4.0
 Requires-Dist: deepseek-tokenizer>=0.1.0
 Requires-Dist: qwen-tokenizer>=0.2.0
 Requires-Dist: google-auth>=2.35.0
+Requires-Dist: diskcache>=5.0.0
 Provides-Extra: server
 Requires-Dist: fastapi; extra == "server"
 Requires-Dist: uvicorn; extra == "server"
+Provides-Extra: redis
+Requires-Dist: redis; extra == "redis"
+Provides-Extra: diskcache
+Requires-Dist: diskcache; extra == "diskcache"
 Description-Content-Type: text/markdown
 
 # vectorvein

{vectorvein-0.1.87.dist-info → vectorvein-0.1.89.dist-info}/RECORD

@@ -1,37 +1,38 @@
-vectorvein-0.1.87.dist-info/METADATA,sha256=qaCzwLhxB8o0buhZ-_94iqCMM8tITVVbZQMGt_xi_Lk,641
-vectorvein-0.1.87.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-vectorvein-0.1.87.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+vectorvein-0.1.89.dist-info/METADATA,sha256=qO2cLUOWAPGVGMH_ufs-2-fosZiJNkEM8fIo9npYEaY,807
+vectorvein-0.1.89.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+vectorvein-0.1.89.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
 vectorvein/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectorvein/chat_clients/__init__.py,sha256=omQuG4PRRPNflSAgtdU--rwsWG6vMpwMEyIGZyFVHVQ,18596
-vectorvein/chat_clients/anthropic_client.py,sha256=PGIKldH4FnGrqozoY_FZ6LqhDHC-jY7NF5J1F1zT2Ok,38257
+vectorvein/chat_clients/anthropic_client.py,sha256=Zk6X1feIvv7Az5dgyipJXbm9TkgWgpFghSTxLiXKKA8,38405
 vectorvein/chat_clients/baichuan_client.py,sha256=CVMvpgjdrZGv0BWnTOBD-f2ufZ3wq3496wqukumsAr4,526
-vectorvein/chat_clients/base_client.py,sha256=Rw-BYFxy86ohZQH8KABiXP6xuNQBjWC3JkeN6WsXfLw,13638
+vectorvein/chat_clients/base_client.py,sha256=QLvcGhjravPbvha6-spU-w6ugHU1LrsbdFUcs6NwMgE,18842
 vectorvein/chat_clients/deepseek_client.py,sha256=3qWu01NlJAP2N-Ff62d5-CZXZitlizE1fzb20LNetig,526
-vectorvein/chat_clients/gemini_client.py,sha256=qqqjQ9X8sIgJaT8xgvtG_cY-lmNGzA_f9V4tUcGRcBo,20853
+vectorvein/chat_clients/gemini_client.py,sha256=ufovIZrmAE3RLEe8h5WXombf7bARAZxnkj6ydNK2FQM,475
 vectorvein/chat_clients/groq_client.py,sha256=Uow4pgdmFi93ZQSoOol2-0PhhqkW-S0XuSldvppz5U4,498
 vectorvein/chat_clients/local_client.py,sha256=55nOsxzqUf79q3Y14MKROA71zxhsT7p7FsDZ89rts2M,422
 vectorvein/chat_clients/minimax_client.py,sha256=ooJU92UCACC4TVWKJ-uo8vqQ8qF3K14ziAuSFm8Wj3M,20025
 vectorvein/chat_clients/mistral_client.py,sha256=1aKSylzBDaLYcFnaBIL4-sXSzWmXfBeON9Q0rq-ziWw,534
 vectorvein/chat_clients/moonshot_client.py,sha256=gbu-6nGxx8uM_U2WlI4Wus881rFRotzHtMSoYOcruGU,526
 vectorvein/chat_clients/openai_client.py,sha256=Nz6tV45pWcsOupxjnsRsGTicbQNJWIZyxuJoJ5DGMpg,527
-vectorvein/chat_clients/openai_compatible_client.py,sha256=HqdECEcm0JzENjmqclwTXFlnoiguYLfAZ7Z8tqYIsNE,28894
+vectorvein/chat_clients/openai_compatible_client.py,sha256=F_kHsoCtrqJ7jLsgyIZ2mJSNQ_YnDp9SRNW4ydFDtic,28950
 vectorvein/chat_clients/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectorvein/chat_clients/qwen_client.py,sha256=-ryh-m9PgsO0fc4ulcCmPTy1155J8YUy15uPoJQOHA0,513
 vectorvein/chat_clients/stepfun_client.py,sha256=zsD2W5ahmR4DD9cqQTXmJr3txrGuvxbRWhFlRdwNijI,519
-vectorvein/chat_clients/utils.py,sha256=FurSHSSpdPjeL8ktZ0AuRKtV4pcAJJYas0PHkw5WBw4,28070
+vectorvein/chat_clients/utils.py,sha256=Nf7EKtKCuWkIi1zkoU-sSjsTT271OvWJsKzxo0WKJX4,24791
 vectorvein/chat_clients/xai_client.py,sha256=eLFJJrNRJ-ni3DpshODcr3S1EJQLbhVwxyO1E54LaqM,491
 vectorvein/chat_clients/yi_client.py,sha256=RNf4CRuPJfixrwLZ3-DEc3t25QDe1mvZeb9sku2f8Bc,484
 vectorvein/chat_clients/zhipuai_client.py,sha256=Ys5DSeLCuedaDXr3PfG1EW2zKXopt-awO2IylWSwY0s,519
 vectorvein/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectorvein/server/token_server.py,sha256=36F9PKSNOX8ZtYBXY_l-76GQTpUSmQ2Y8EMy1H7wtdQ,1353
-vectorvein/settings/__init__.py,sha256=g01y74x0k2JEAqNpRGG0PDs0NTULjOAZV6HRhydPX1c,3874
+vectorvein/settings/__init__.py,sha256=ecGyrE_6YfX9z6Igb1rDCu1Q-qMTcVozWF3WEl_hiKA,4871
 vectorvein/settings/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vectorvein/types/defaults.py,sha256=rM5qCUmaVx3-hBJ3x7ClHF4mT3qU7mrcsS2FrumLdUA,26530
+vectorvein/types/defaults.py,sha256=MAoxFhtvzPWBl1Zroz6hhKl7AshyHcO90hcE3pXzePQ,27384
 vectorvein/types/enums.py,sha256=7KTJSVtQueImmbr1fSwv3rQVtc0RyMWXJmoE2tDOaso,1667
 vectorvein/types/exception.py,sha256=gnW4GnJ76jND6UGnodk9xmqkcbeS7Cz2rvncA2HpD5E,69
-vectorvein/types/llm_parameters.py,sha256=CLhDSp9KI_zzjIXUvjiTuGxfYXpubTNBCVcJ6RgH2iY,5879
+vectorvein/types/llm_parameters.py,sha256=jXHGR9aORkBrUaG4oQef3zorFzvVX2oTn2lMu57IOQs,5989
 vectorvein/types/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vectorvein/utilities/media_processing.py,sha256=CTRq-lGlFkFgP_FSRhNwF_qUgmOrXPf2_1Ok9HY42_g,5887
+vectorvein/utilities/rate_limiter.py,sha256=dwolIUVw2wP83Odqpx0AAaE77de1GzxkYDGH4tM_u_4,10300
 vectorvein/utilities/retry.py,sha256=6KFS9R2HdhqM3_9jkjD4F36ZSpEx2YNFGOVlpOsUetM,2208
 vectorvein/workflow/graph/edge.py,sha256=xLZEJmBjAfVB53cd7CuRcKhgE6QqXv9nz32wJn8cmyk,1064
 vectorvein/workflow/graph/node.py,sha256=A3M_GghrSju1D3xc_HtPdGyr-7XSkplGPKJveOUiIF4,3256
@@ -54,4 +55,4 @@ vectorvein/workflow/nodes/vector_db.py,sha256=t6I17q6iR3yQreiDHpRrksMdWDPIvgqJs0
 vectorvein/workflow/nodes/video_generation.py,sha256=qmdg-t_idpxq1veukd-jv_ChICMOoInKxprV9Z4Vi2w,4118
 vectorvein/workflow/nodes/web_crawlers.py,sha256=LsqomfXfqrXfHJDO1cl0Ox48f4St7X_SL12DSbAMSOw,5415
 vectorvein/workflow/utils/json_to_code.py,sha256=F7dhDy8kGc8ndOeihGLRLGFGlquoxVlb02ENtxnQ0C8,5914
-vectorvein-0.1.87.dist-info/RECORD,,
+vectorvein-0.1.89.dist-info/RECORD,,