andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-23 07:08:14 +00:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

fixing RPM controlelr being set unencessarily

Browse Source
This commit is contained in:
João Moura
2024-02-07 23:09:36 -08:00
parent 98133a4eb6
commit a45c82c5f7
3 changed files with 26 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/crewai/crew.py
View File

@@ -124,7 +124,8 @@ class Crew(BaseModel):
if self.agents:
for agent in self.agents:
agent.set_cache_handler(self._cache_handler)
agent.set_rpm_controller(self._rpm_controller)
if self.max_rpm:
agent.set_rpm_controller(self._rpm_controller)
return self
def _setup_from_config(self):

7
src/crewai/utilities/rpm_controller.py
View File

@@ -14,12 +14,14 @@ class RPMController(BaseModel):
_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:
self._lock = threading.Lock()
self._reset_request_count()
if not self._shutdown_flag:
self._lock = threading.Lock()
self._reset_request_count()
return self
def check_or_wait(self):
@@ -51,6 +53,7 @@ class RPMController(BaseModel):
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()