-
Notifications
You must be signed in to change notification settings - Fork 3.4k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
4fb7014
commit 1954be4
Showing
2 changed files
with
68 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,11 @@ | ||
class Logger: | ||
def __init__(self, verbose_level=0): | ||
verbose_level = ( | ||
2 if isinstance(verbose_level, bool) and verbose_level else verbose_level | ||
) | ||
self.verbose_level = verbose_level | ||
|
||
def log(self, level, message): | ||
level_map = {"debug": 1, "info": 2} | ||
if self.verbose_level and level_map.get(level, 0) <= self.verbose_level: | ||
print(f"\n[{level.upper()}]: {message}") |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,57 @@ | ||
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 = PrivateAttr(default=None) | ||
_lock: threading.Lock = PrivateAttr(default=None) | ||
|
||
@model_validator(mode="after") | ||
def reset_counter(self): | ||
if self.max_rpm: | ||
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) | ||
with self._lock: | ||
self._current_rpm = 0 | ||
|
||
def _reset_request_count(self): | ||
with self._lock: | ||
self._current_rpm = 0 | ||
if self._timer: | ||
self._timer.cancel() | ||
self._timer = threading.Timer(60.0, self._reset_request_count) | ||
self._timer.start() |