fixing RPM controlelr being set unencessarily

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

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):

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()

View File

@@ -356,6 +356,25 @@ def test_api_calls_throttling(capsys):
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():
import threading
from unittest.mock import patch