feat: add rpm & tpm limit

ChenZiHong-Gavin · ChenZiHong-Gavin · commit 22f4693e9e34 · 2025-04-18T23:34:27.000+08:00
diff --git a/graphgen/models/llm/limitter.py b/graphgen/models/llm/limitter.py
@@ -0,0 +1,88 @@
+import time
+from datetime import datetime, timedelta
+import asyncio
+
+from graphgen.utils import logger
+
+
+class RPM:
+
+    def __init__(self, rpm: int = 1000):
+        self.rpm = rpm
+        self.record = {'rpm_slot': self.get_minute_slot(), 'counter': 0}
+
+    def get_minute_slot(self):
+        current_time = time.time()
+        dt_object = datetime.fromtimestamp(current_time)
+        total_minutes_since_midnight = dt_object.hour * 60 + dt_object.minute
+        return total_minutes_since_midnight
+
+    async def wait(self, silent=False):
+        current = time.time()
+        dt_object = datetime.fromtimestamp(current)
+        minute_slot = self.get_minute_slot()
+
+        if self.record['rpm_slot'] == minute_slot:
+            # check RPM exceed
+            if self.record['counter'] >= self.rpm:
+                # wait until next minute
+                next_minute = dt_object.replace(
+                    second=0, microsecond=0) + timedelta(minutes=1)
+                _next = next_minute.timestamp()
+                sleep_time = abs(_next - current)
+                if not silent:
+                    logger.info('RPM sleep %s', sleep_time)
+                await asyncio.sleep(sleep_time)
+
+                self.record = {
+                    'rpm_slot': self.get_minute_slot(),
+                    'counter': 0
+                }
+        else:
+            self.record = {'rpm_slot': self.get_minute_slot(), 'counter': 0}
+        self.record['counter'] += 1
+
+        if not silent:
+            logger.debug(self.record)
+
+
+class TPM:
+
+    def __init__(self, tpm: int = 20000):
+        self.tpm = tpm
+        self.record = {'tpm_slot': self.get_minute_slot(), 'counter': 0}
+
+    def get_minute_slot(self):
+        current_time = time.time()
+        dt_object = datetime.fromtimestamp(current_time)
+        total_minutes_since_midnight = dt_object.hour * 60 + dt_object.minute
+        return total_minutes_since_midnight
+
+    async def wait(self, token_count, silent=False):
+        current = time.time()
+        dt_object = datetime.fromtimestamp(current)
+        minute_slot = self.get_minute_slot()
+
+        # get next slot, skip
+        if self.record['tpm_slot'] != minute_slot:
+            self.record = {'tpm_slot': minute_slot, 'counter': token_count}
+            return
+
+        # check RPM exceed
+        self.record['counter'] += token_count
+        if self.record['counter'] > self.tpm:
+            # wait until next minute
+            next_minute = dt_object.replace(
+                second=0, microsecond=0) + timedelta(minutes=1)
+            _next = next_minute.timestamp()
+            sleep_time = abs(_next - current)
+            logger.info('TPM sleep %s', sleep_time)
+            await asyncio.sleep(sleep_time)
+
+            self.record = {
+                'tpm_slot': self.get_minute_slot(),
+                'counter': token_count
+            }
+
+        if not silent:
+            logger.debug(self.record)
diff --git a/graphgen/models/llm/openai_model.py b/graphgen/models/llm/openai_model.py
@@ -11,7 +11,8 @@
 )
 
 from graphgen.models.llm.topk_token_model import TopkTokenModel, Token
-
+from graphgen.models.llm.tokenizer import Tokenizer
+from graphgen.models.llm.limitter import RPM, TPM
 
 def get_top_response_tokens(response: openai.ChatCompletion) -> List[Token]:
     token_logprobs = response.choices[0].logprobs.content
@@ -31,10 +32,16 @@ class OpenAIModel(TopkTokenModel):
     model_name: str = "gpt-4o-mini"
     api_key: str = None
     base_url: str = None
+
     system_prompt: str = ""
     json_mode: bool = False
     seed: int = None
+
     token_usage: list = field(default_factory=list)
+    request_limit: bool = False
+    rpm: RPM = field(default_factory=lambda: RPM(rpm=1000))
+    tpm: TPM = field(default_factory=lambda: TPM(tpm=50000))
+
 
     def __post_init__(self):
         assert self.api_key is not None, "Please provide api key to access openai api."
@@ -63,6 +70,7 @@ def _pre_generate(self, text: str, history: List[str]) -> Dict:
         kwargs['messages']= messages
         return kwargs
 
+
     @retry(
         stop=stop_after_attempt(5),
         wait=wait_exponential(multiplier=1, min=4, max=10),
@@ -95,6 +103,15 @@ async def generate_answer(self, text: str, history: Optional[List[str]] = None,
         kwargs = self._pre_generate(text, history)
         kwargs["temperature"] = temperature
 
+        prompt_tokens = 0
+        for message in kwargs['messages']:
+            prompt_tokens += len(Tokenizer().encode_string(message['content']))
+        estimated_tokens = prompt_tokens + kwargs['max_tokens']
+
+        if self.request_limit:
+            await self.rpm.wait(silent=True)
+            await self.tpm.wait(estimated_tokens, silent=True)
+
         completion = await self.client.chat.completions.create( # pylint: disable=E1125
             model=self.model_name,
             **kwargs
diff --git a/webui/app.py b/webui/app.py
@@ -17,6 +17,7 @@
 
 from graphgen.graphgen import GraphGen
 from graphgen.models import OpenAIModel, Tokenizer, TraverseStrategy
+from graphgen.models.llm.limitter import RPM, TPM
 
 css = """
 .center-row {
@@ -38,12 +39,20 @@ def init_graph_gen(config: dict, env: dict) -> GraphGen:
     graph_gen.synthesizer_llm_client = OpenAIModel(
         model_name=env.get("SYNTHESIZER_MODEL", ""),
         base_url=env.get("SYNTHESIZER_BASE_URL", ""),
-        api_key=env.get("SYNTHESIZER_API_KEY", ""))
+        api_key=env.get("SYNTHESIZER_API_KEY", ""),
+        request_limit=True,
+        rpm= RPM(env.get("RPM", 1000)),
+        tpm= TPM(env.get("TPM", 50000)),
+    )
 
     graph_gen.trainee_llm_client = OpenAIModel(
         model_name=env.get("TRAINEE_MODEL", ""),
         base_url=env.get("TRAINEE_BASE_URL", ""),
-        api_key=env.get("TRAINEE_API_KEY", ""))
+        api_key=env.get("TRAINEE_API_KEY", ""),
+        request_limit=True,
+        rpm= RPM(env.get("RPM", 1000)),
+        tpm= TPM(env.get("TPM", 50000)),
+    )
 
     graph_gen.tokenizer_instance = Tokenizer(
         config.get("tokenizer", "cl100k_base"))
@@ -97,7 +106,9 @@ def sum_tokens(client):
         "TRAINEE_BASE_URL": arguments[12],
         "TRAINEE_MODEL": arguments[14],
         "SYNTHESIZER_API_KEY": arguments[15],
-        "TRAINEE_API_KEY": arguments[15]
+        "TRAINEE_API_KEY": arguments[15],
+        "RPM": arguments[17],
+        "TPM": arguments[18],
     }
 
     # Test API connection
@@ -362,6 +373,28 @@ def sum_tokens(client):
             with gr.Column(scale=1):
                 test_connection_btn = gr.Button("Test Connection")
 
+        with gr.Blocks():
+            with gr.Row(equal_height=True):
+                with gr.Column():
+                    rpm = gr.Slider(
+                        label="RPM",
+                        minimum=500,
+                        maximum=10000,
+                        value=1000,
+                        step=100,
+                        interactive=True,
+                        visible=True)
+                with gr.Column():
+                    tpm = gr.Slider(
+                        label="TPM",
+                        minimum=5000,
+                        maximum=100000,
+                        value=50000,
+                        step=1000,
+                        interactive=True,
+                        visible=True)
+
+
         with gr.Blocks():
             with gr.Row(equal_height=True):
                 with gr.Column(scale=1):
@@ -442,7 +475,7 @@ def sum_tokens(client):
                 bidirectional, expand_method, max_extra_edges, max_tokens,
                 max_depth, edge_sampling, isolated_node_strategy,
                 loss_strategy, base_url, synthesizer_model, trainee_model,
-                api_key, chunk_size, token_counter
+                api_key, chunk_size, rpm, tpm, token_counter
             ],
             outputs=[output, token_counter],
         )
diff --git a/webui/count_tokens.py b/webui/count_tokens.py
@@ -9,6 +9,9 @@
 from graphgen.models import Tokenizer
 
 def count_tokens(file, tokenizer_name, data_frame):
+    if not file or not os.path.exists(file):
+        return data_frame
+
     if file.endswith(".jsonl"):
         with open(file, "r", encoding='utf-8') as f:
             data = [json.loads(line) for line in f]
