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

vectorvein/chat_clients/anthropic_client.py

@@ -434,6 +434,8 @@ class AnthropicChatClient(BaseChatClient):
             else:
                 max_tokens = self.model_setting.context_length - token_counts
 
+        self._acquire_rate_limit(self.endpoint, self.model, messages)
+
         if self.stream:
             stream_response = raw_client.messages.create(
                 model=self.model_id,
@@ -824,6 +826,8 @@ class AsyncAnthropicChatClient(BaseAsyncChatClient):
             else:
                 max_tokens = self.model_setting.context_length - token_counts
 
+        await self._acquire_rate_limit(self.endpoint, self.model, messages)
+
         if self.stream:
             stream_response = await raw_client.messages.create(
                 model=self.model_id,

vectorvein/chat_clients/base_client.py

@@ -1,7 +1,8 @@
-# @Author: Bi Ying
-# @Date:   2024-07-26 14:48:55
+import time
 import random
+import asyncio
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from functools import cached_property
 from typing import Generator, AsyncGenerator, Any, overload, Literal, Iterable
 
@@ -29,6 +30,8 @@ from ..types.llm_parameters import (
     ChatCompletionDeltaMessage,
     ChatCompletionStreamOptionsParam,
 )
+from ..utilities.rate_limiter import SyncMemoryRateLimiter, SyncRedisRateLimiter, SyncDiskCacheRateLimiter
+from ..utilities.rate_limiter import AsyncMemoryRateLimiter, AsyncRedisRateLimiter, AsyncDiskCacheRateLimiter
 
 
 class BaseChatClient(ABC):
@@ -59,11 +62,65 @@ class BaseChatClient(ABC):
 
         self.backend_settings = settings.get_backend(self.BACKEND_NAME)
 
+        self.rate_limiter = self._init_rate_limiter()
+        self.active_requests = defaultdict(int)
+        self.rpm = None
+        self.tpm = None
+        self.concurrent_requests = None
+
         if endpoint_id:
             self.endpoint_id = endpoint_id
             self.random_endpoint = False
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
+    def _init_rate_limiter(self):
+        if not settings.rate_limit:
+            return None
+        if not settings.rate_limit.enabled:
+            return None
+
+        if settings.rate_limit.backend == "memory":
+            return SyncMemoryRateLimiter()
+        elif settings.rate_limit.backend == "redis":
+            if not settings.rate_limit.redis:
+                raise ValueError("Redis settings must be provided if Redis backend is selected.")
+            return SyncRedisRateLimiter(
+                host=settings.rate_limit.redis.host,
+                port=settings.rate_limit.redis.port,
+                db=settings.rate_limit.redis.db,
+            )
+        elif settings.rate_limit.backend == "diskcache":
+            if not settings.rate_limit.diskcache:
+                raise ValueError("Diskcache settings must be provided if Diskcache backend is selected.")
+            return SyncDiskCacheRateLimiter(
+                cache_dir=settings.rate_limit.diskcache.cache_dir,
+            )
+        return None
+
+    def _acquire_rate_limit(self, endpoint: EndpointSetting | None, model: str, messages: list):
+        if endpoint is None:
+            return
+
+        key = f"{endpoint.id}:{model}"
+
+        # Get rate limit parameters
+        # Priority: parameters in model.endpoints > parameters in endpoint > default parameters
+        rpm = self.rpm or endpoint.rpm or (settings.rate_limit.default_rpm if settings.rate_limit else 60)
+        tpm = self.tpm or endpoint.tpm or (settings.rate_limit.default_tpm if settings.rate_limit else 1000000)
+
+        while self.rate_limiter:
+            allowed, wait_time = self.rate_limiter.check_limit(key, rpm, tpm, self._estimate_request_tokens(messages))
+            if allowed:
+                break
+            time.sleep(wait_time)
+
+    def _estimate_request_tokens(self, messages: list) -> int:
+        """Roughly estimate the number of tokens in the request"""
+        tokens = 0
+        for message in messages:
+            tokens += int(len(message.get("content", "")) * 0.6)
+        return tokens
+
     def set_model_id_by_endpoint_id(self, endpoint_id: str):
         for endpoint_option in self.backend_settings.models[self.model].endpoints:
             if isinstance(endpoint_option, dict) and endpoint_id == endpoint_option["endpoint_id"]:
@@ -79,6 +136,9 @@ class BaseChatClient(ABC):
                 if isinstance(endpoint, dict):
                     self.endpoint_id = endpoint["endpoint_id"]
                     self.model_id = endpoint["model_id"]
+                    self.rpm = endpoint.get("rpm", None)
+                    self.tpm = endpoint.get("tpm", None)
+                    self.concurrent_requests = endpoint.get("concurrent_requests", None)
                 else:
                     self.endpoint_id = endpoint
                 self.endpoint = settings.get_endpoint(self.endpoint_id)
@@ -236,11 +296,67 @@ class BaseAsyncChatClient(ABC):
 
         self.backend_settings = settings.get_backend(self.BACKEND_NAME)
 
+        self.rate_limiter = self._init_rate_limiter()
+        self.active_requests = defaultdict(int)
+        self.rpm = None
+        self.tpm = None
+        self.concurrent_requests = None
+
         if endpoint_id:
             self.endpoint_id = endpoint_id
             self.random_endpoint = False
             self.endpoint = settings.get_endpoint(self.endpoint_id)
 
+    def _init_rate_limiter(self):
+        if not settings.rate_limit:
+            return None
+        if not settings.rate_limit.enabled:
+            return None
+
+        if settings.rate_limit.backend == "memory":
+            return AsyncMemoryRateLimiter()
+        elif settings.rate_limit.backend == "redis":
+            if not settings.rate_limit.redis:
+                raise ValueError("Redis settings must be provided if Redis backend is selected.")
+            return AsyncRedisRateLimiter(
+                host=settings.rate_limit.redis.host,
+                port=settings.rate_limit.redis.port,
+                db=settings.rate_limit.redis.db,
+            )
+        elif settings.rate_limit.backend == "diskcache":
+            if not settings.rate_limit.diskcache:
+                raise ValueError("Diskcache settings must be provided if Diskcache backend is selected.")
+            return AsyncDiskCacheRateLimiter(
+                cache_dir=settings.rate_limit.diskcache.cache_dir,
+            )
+        return None
+
+    async def _acquire_rate_limit(self, endpoint: EndpointSetting | None, model: str, messages: list):
+        if endpoint is None:
+            return
+
+        key = f"{endpoint.id}:{model}"
+
+        # Get rate limit parameters
+        # Priority: parameters in model.endpoints > parameters in endpoint > default parameters
+        rpm = self.rpm or endpoint.rpm or (settings.rate_limit.default_rpm if settings.rate_limit else 60)
+        tpm = self.tpm or endpoint.tpm or (settings.rate_limit.default_tpm if settings.rate_limit else 1000000)
+
+        while self.rate_limiter:
+            allowed, wait_time = await self.rate_limiter.check_limit(
+                key, rpm, tpm, self._estimate_request_tokens(messages)
+            )
+            if allowed:
+                break
+            await asyncio.sleep(wait_time)
+
+    def _estimate_request_tokens(self, messages: list) -> int:
+        """Roughly estimate the number of tokens in the request"""
+        tokens = 0
+        for message in messages:
+            tokens += int(len(message.get("content", "")) * 0.6)
+        return tokens
+
     def set_model_id_by_endpoint_id(self, endpoint_id: str):
         for endpoint_option in self.backend_settings.models[self.model].endpoints:
             if isinstance(endpoint_option, dict) and endpoint_id == endpoint_option["endpoint_id"]:
@@ -256,6 +372,9 @@ class BaseAsyncChatClient(ABC):
                 if isinstance(endpoint, dict):
                     self.endpoint_id = endpoint["endpoint_id"]
                     self.model_id = endpoint["model_id"]
+                    self.rpm = endpoint.get("rpm", None)
+                    self.tpm = endpoint.get("tpm", None)
+                    self.concurrent_requests = endpoint.get("concurrent_requests", None)
                 else:
                     self.endpoint_id = endpoint
                 self.endpoint = settings.get_endpoint(self.endpoint_id)