Spaces:
Sleeping
Sleeping
File size: 1,786 Bytes
94ea28c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import threading
import time
from typing import Union
from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator
from crewai.utilities.logger import Logger
class RPMController(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
max_rpm: Union[int, None] = Field(default=None)
logger: Logger = Field(default=None)
_current_rpm: int = PrivateAttr(default=0)
_timer: threading.Timer | None = PrivateAttr(default=None)
_lock: threading.Lock = PrivateAttr(default=None)
_shutdown_flag = False
@model_validator(mode="after")
def reset_counter(self):
if self.max_rpm:
if not self._shutdown_flag:
self._lock = threading.Lock()
self._reset_request_count()
return self
def check_or_wait(self):
if not self.max_rpm:
return True
with self._lock:
if self._current_rpm < self.max_rpm:
self._current_rpm += 1
return True
else:
self.logger.log(
"info", "Max RPM reached, waiting for next minute to start."
)
self._wait_for_next_minute()
self._current_rpm = 1
return True
def stop_rpm_counter(self):
if self._timer:
self._timer.cancel()
self._timer = None
def _wait_for_next_minute(self):
time.sleep(60)
self._current_rpm = 0
def _reset_request_count(self):
with self._lock:
self._current_rpm = 0
if self._timer:
self._shutdown_flag = True
self._timer.cancel()
self._timer = threading.Timer(60.0, self._reset_request_count)
self._timer.start()
|