mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-05-02 15:52:34 +00:00
fixing RPM controlelr being set unencessarily
This commit is contained in:
@@ -124,7 +124,8 @@ class Crew(BaseModel):
|
|||||||
if self.agents:
|
if self.agents:
|
||||||
for agent in self.agents:
|
for agent in self.agents:
|
||||||
agent.set_cache_handler(self._cache_handler)
|
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
|
return self
|
||||||
|
|
||||||
def _setup_from_config(self):
|
def _setup_from_config(self):
|
||||||
|
|||||||
@@ -14,12 +14,14 @@ class RPMController(BaseModel):
|
|||||||
_current_rpm: int = PrivateAttr(default=0)
|
_current_rpm: int = PrivateAttr(default=0)
|
||||||
_timer: threading.Timer | None = PrivateAttr(default=None)
|
_timer: threading.Timer | None = PrivateAttr(default=None)
|
||||||
_lock: threading.Lock = PrivateAttr(default=None)
|
_lock: threading.Lock = PrivateAttr(default=None)
|
||||||
|
_shutdown_flag = False
|
||||||
|
|
||||||
@model_validator(mode="after")
|
@model_validator(mode="after")
|
||||||
def reset_counter(self):
|
def reset_counter(self):
|
||||||
if self.max_rpm:
|
if self.max_rpm:
|
||||||
self._lock = threading.Lock()
|
if not self._shutdown_flag:
|
||||||
self._reset_request_count()
|
self._lock = threading.Lock()
|
||||||
|
self._reset_request_count()
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def check_or_wait(self):
|
def check_or_wait(self):
|
||||||
@@ -51,6 +53,7 @@ class RPMController(BaseModel):
|
|||||||
with self._lock:
|
with self._lock:
|
||||||
self._current_rpm = 0
|
self._current_rpm = 0
|
||||||
if self._timer:
|
if self._timer:
|
||||||
|
self._shutdown_flag = True
|
||||||
self._timer.cancel()
|
self._timer.cancel()
|
||||||
self._timer = threading.Timer(60.0, self._reset_request_count)
|
self._timer = threading.Timer(60.0, self._reset_request_count)
|
||||||
self._timer.start()
|
self._timer.start()
|
||||||
|
|||||||
@@ -356,6 +356,25 @@ def test_api_calls_throttling(capsys):
|
|||||||
moveon.assert_called()
|
moveon.assert_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_agents_rpm_is_never_set_if_crew_max_RPM_is_not_set():
|
||||||
|
agent = Agent(
|
||||||
|
role="test role",
|
||||||
|
goal="test goal",
|
||||||
|
backstory="test backstory",
|
||||||
|
allow_delegation=False,
|
||||||
|
verbose=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
task = Task(
|
||||||
|
description="just say hi!",
|
||||||
|
agent=agent,
|
||||||
|
)
|
||||||
|
|
||||||
|
Crew(agents=[agent], tasks=[task], verbose=2)
|
||||||
|
|
||||||
|
assert agent._rpm_controller is None
|
||||||
|
|
||||||
|
|
||||||
def test_async_task_execution():
|
def test_async_task_execution():
|
||||||
import threading
|
import threading
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|||||||
Reference in New Issue
Block a user