Adding support for Crew throttling using RPM (#124)

* Add translations
* fixing translations
* Adding support for Crew throttling with RPM
This commit is contained in:
João Moura
2024-01-13 11:20:30 -03:00
committed by GitHub
parent 3c74fbf9ab
commit 1fc806161f
5 changed files with 784 additions and 26 deletions

View File

@@ -1,6 +1,8 @@
import json
import threading
import time
import uuid
from typing import Any, Dict, List, Optional, Union
from typing import Any, ClassVar, Dict, List, Optional, Union
from pydantic import (
UUID4,
@@ -9,6 +11,7 @@ from pydantic import (
Field,
InstanceOf,
Json,
PrivateAttr,
field_validator,
model_validator,
)
@@ -33,10 +36,16 @@ class Crew(BaseModel):
verbose: Indicates the verbosity level for logging during execution.
config: Configuration settings for the crew.
cache_handler: Handles caching for the crew's operations.
max_rpm: Maximum number of requests per minute for the crew execution to be respected.
rpm: Current number of requests per minute for the crew execution.
id: A unique identifier for the crew instance.
"""
__hash__ = object.__hash__
_timer: Optional[threading.Timer] = PrivateAttr(default=None)
lock: ClassVar[threading.Lock] = threading.Lock()
rpm: ClassVar[int] = 0
max_rpm: Optional[int] = Field(default=None)
model_config = ConfigDict(arbitrary_types_allowed=True)
tasks: List[Task] = Field(default_factory=list)
agents: List[Agent] = Field(default_factory=list)
@@ -64,6 +73,12 @@ class Crew(BaseModel):
def check_config_type(cls, v: Union[Json, Dict[str, Any]]):
return json.loads(v) if isinstance(v, Json) else v
@model_validator(mode="after")
def set_reset_counter(self):
if self.max_rpm:
self._reset_request_count()
return self
@model_validator(mode="after")
def check_config(self):
"""Validates that the crew is properly configured with agents and tasks."""
@@ -80,6 +95,7 @@ class Crew(BaseModel):
if self.agents:
for agent in self.agents:
agent.set_cache_handler(self.cache_handler)
agent.set_request_within_rpm_limit(self.ensure_request_within_rpm_limit)
return self
def _setup_from_config(self):
@@ -100,6 +116,24 @@ class Crew(BaseModel):
del task_config["agent"]
return Task(**task_config, agent=task_agent)
def ensure_request_within_rpm_limit(self):
if not self.max_rpm:
return True
with Crew.lock:
if Crew.rpm < self.max_rpm:
Crew.rpm += 1
return True
self._log("info", "Max RPM reached, waiting for next minute to start.")
return self._wait_for_next_minute()
def _wait_for_next_minute(self):
time.sleep(60)
with Crew.lock:
Crew.rpm = 0
return True
def kickoff(self) -> str:
"""Starts the crew to work on its assigned tasks."""
for agent in self.agents:
@@ -115,9 +149,8 @@ class Crew(BaseModel):
for task in self.tasks:
self._prepare_and_execute_task(task)
task_output = task.execute(task_output)
self._log(
"debug", f"\n\n[{task.agent.role}] Task output: {task_output}\n\n"
)
self._log("debug", f"\n[{task.agent.role}] Task output: {task_output}\n\n")
self._stop_timer()
return task_output
def _prepare_and_execute_task(self, task):
@@ -135,4 +168,14 @@ class Crew(BaseModel):
2 if isinstance(self.verbose, bool) and self.verbose else self.verbose
)
if verbose_level and level_map[level] <= verbose_level:
print(message)
print(f"\n{message}")
def _stop_timer(self):
if self._timer:
self._timer.cancel()
def _reset_request_count(self):
self._stop_timer()
self._timer = threading.Timer(60.0, self._reset_request_count)
self._timer.start()
Crew.rpm = 0