From 185556b7e35163cc830214f49bc1231f93085dd4 Mon Sep 17 00:00:00 2001 From: Brandon Hancock Date: Fri, 14 Feb 2025 11:40:23 -0500 Subject: [PATCH] WIP --- src/crewai/agent.py | 6 ++-- src/crewai/agents/agent_builder/base_agent.py | 31 +++++++++++++------ 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/src/crewai/agent.py b/src/crewai/agent.py index 93a653c7b..0f107c307 100644 --- a/src/crewai/agent.py +++ b/src/crewai/agent.py @@ -175,8 +175,9 @@ class Agent(BaseAgent): Returns: Output of the agent """ + # The RPM controller is now managed by the Crew, so no need to set it here. if self.tools_handler: - self.tools_handler.last_used_tool = {} # type: ignore # Incompatible types in assignment (expression has type "dict[Never, Never]", variable has type "ToolCalling") + self.tools_handler.last_used_tool = {} # type: ignore # Incompatible types in assignment (expression has type "dict[Never, Never]", variable has type "ToolCalli task_prompt = task.prompt() @@ -250,9 +251,6 @@ class Agent(BaseAgent): raise e result = self.execute_task(task, context, tools) - if self.max_rpm and self._rpm_controller: - self._rpm_controller.stop_rpm_counter() - # If there was any tool in self.tools_results that had result_as_answer # set to True, return the results of the last tool that had # result_as_answer set to True diff --git a/src/crewai/agents/agent_builder/base_agent.py b/src/crewai/agents/agent_builder/base_agent.py index 32ce6d8d5..9910031d8 100644 --- a/src/crewai/agents/agent_builder/base_agent.py +++ b/src/crewai/agents/agent_builder/base_agent.py @@ -74,7 +74,7 @@ class BaseAgent(ABC, BaseModel): Increment formatting errors. copy() -> "BaseAgent": Create a copy of the agent. - set_rpm_controller(rpm_controller: RPMController) -> None: + set_rpm_controller(rpm_controller: Optional[RPMController] = None) -> None: Set the rpm controller for the agent. set_private_attrs() -> "BaseAgent": Set private attributes. @@ -347,15 +347,26 @@ class BaseAgent(ABC, BaseModel): def increment_formatting_errors(self) -> None: self.formatting_errors += 1 - def set_rpm_controller(self, rpm_controller: RPMController) -> None: - """Set the rpm controller for the agent. - - Args: - rpm_controller: An instance of the RPMController class. + def set_rpm_controller( + self, rpm_controller: Optional[RPMController] = None + ) -> None: """ - if not self._rpm_controller: - self._rpm_controller = rpm_controller - # Only create the executor if it hasn't been created yet. + Set the RPM controller for the agent. If no rpm_controller is provided, then: + - use self.max_rpm if set, or + - if self.crew exists and has max_rpm, use that. + """ + if self._rpm_controller is None: + if rpm_controller is not None: + self._rpm_controller = rpm_controller + elif self.max_rpm: + self._rpm_controller = RPMController( + max_rpm=self.max_rpm, logger=self._logger + ) + elif self.crew and getattr(self.crew, "max_rpm", None): + self._rpm_controller = RPMController( + max_rpm=self.crew.max_rpm, logger=self._logger + ) + # else: no rpm limit provided – leave the controller None if self.agent_executor is None: self.create_agent_executor() @@ -370,4 +381,4 @@ class BaseAgent(ABC, BaseModel): if self.cache: self.set_cache_handler(cache_handler) if self.max_rpm: - self.set_rpm_controller(rpm_controller) + self.set_rpm_controller()