mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-05-01 23:32:39 +00:00
Add RPM control to both agents and crews (#133)
* moving file into utilities * creating Logger and RPMController * Adding support for RPM to agents and crew
This commit is contained in:
57
src/crewai/utilities/rpm_controller.py
Normal file
57
src/crewai/utilities/rpm_controller.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user