mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-15 02:58:30 +00:00
Adding support for Crew throttling using RPM (#124)
* Add translations * fixing translations * Adding support for Crew throttling with RPM
This commit is contained in:
@@ -12,6 +12,7 @@ from pydantic import (
|
||||
ConfigDict,
|
||||
Field,
|
||||
InstanceOf,
|
||||
PrivateAttr,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
@@ -46,6 +47,8 @@ class Agent(BaseModel):
|
||||
"""
|
||||
|
||||
__hash__ = object.__hash__
|
||||
_request_within_rpm_limit: Any = PrivateAttr(default=None)
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
id: UUID4 = Field(
|
||||
default_factory=uuid.uuid4,
|
||||
@@ -139,6 +142,10 @@ class Agent(BaseModel):
|
||||
self.tools_handler = ToolsHandler(cache=self.cache_handler)
|
||||
self.__create_agent_executor()
|
||||
|
||||
def set_request_within_rpm_limit(self, ensure_function) -> None:
|
||||
self._request_within_rpm_limit = ensure_function
|
||||
self.__create_agent_executor()
|
||||
|
||||
def __create_agent_executor(self) -> CrewAgentExecutor:
|
||||
"""Create an agent executor for the agent.
|
||||
|
||||
@@ -157,11 +164,12 @@ class Agent(BaseModel):
|
||||
"verbose": self.verbose,
|
||||
"handle_parsing_errors": True,
|
||||
"max_iterations": self.max_iter,
|
||||
"request_within_rpm_limit": self._request_within_rpm_limit,
|
||||
}
|
||||
|
||||
if self.memory:
|
||||
summary_memory = ConversationSummaryMemory(
|
||||
llm=self.llm, memory_key="chat_history", input_key="input"
|
||||
llm=self.llm, input_key="input", memory_key="chat_history"
|
||||
)
|
||||
executor_args["memory"] = summary_memory
|
||||
agent_args["chat_history"] = lambda x: x["chat_history"]
|
||||
|
||||
@@ -19,6 +19,7 @@ from crewai.tools.cache_tools import CacheTools
|
||||
class CrewAgentExecutor(AgentExecutor):
|
||||
i18n: I18N = I18N()
|
||||
iterations: int = 0
|
||||
request_within_rpm_limit: Any = None
|
||||
max_iterations: Optional[int] = 15
|
||||
force_answer_max_iterations: Optional[int] = None
|
||||
|
||||
@@ -54,29 +55,30 @@ class CrewAgentExecutor(AgentExecutor):
|
||||
start_time = time.time()
|
||||
# We now enter the agent loop (until it returns something).
|
||||
while self._should_continue(self.iterations, time_elapsed):
|
||||
next_step_output = self._take_next_step(
|
||||
name_to_tool_map,
|
||||
color_mapping,
|
||||
inputs,
|
||||
intermediate_steps,
|
||||
run_manager=run_manager,
|
||||
)
|
||||
if isinstance(next_step_output, AgentFinish):
|
||||
return self._return(
|
||||
next_step_output, intermediate_steps, run_manager=run_manager
|
||||
if not self.request_within_rpm_limit or self.request_within_rpm_limit():
|
||||
next_step_output = self._take_next_step(
|
||||
name_to_tool_map,
|
||||
color_mapping,
|
||||
inputs,
|
||||
intermediate_steps,
|
||||
run_manager=run_manager,
|
||||
)
|
||||
|
||||
intermediate_steps.extend(next_step_output)
|
||||
if len(next_step_output) == 1:
|
||||
next_step_action = next_step_output[0]
|
||||
# See if tool should return directly
|
||||
tool_return = self._get_tool_return(next_step_action)
|
||||
if tool_return is not None:
|
||||
if isinstance(next_step_output, AgentFinish):
|
||||
return self._return(
|
||||
tool_return, intermediate_steps, run_manager=run_manager
|
||||
next_step_output, intermediate_steps, run_manager=run_manager
|
||||
)
|
||||
self.iterations += 1
|
||||
time_elapsed = time.time() - start_time
|
||||
|
||||
intermediate_steps.extend(next_step_output)
|
||||
if len(next_step_output) == 1:
|
||||
next_step_action = next_step_output[0]
|
||||
# See if tool should return directly
|
||||
tool_return = self._get_tool_return(next_step_action)
|
||||
if tool_return is not None:
|
||||
return self._return(
|
||||
tool_return, intermediate_steps, run_manager=run_manager
|
||||
)
|
||||
self.iterations += 1
|
||||
time_elapsed = time.time() - start_time
|
||||
output = self.agent.return_stopped_response(
|
||||
self.early_stopping_method, intermediate_steps, **inputs
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user