Skip to content

Commit

Permalink
Update some docstrings / typehints (#144)
Browse files Browse the repository at this point in the history
  • Loading branch information
greysonlalonde authored Jan 21, 2024
1 parent 23a549a commit 852ca21
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 13 deletions.
24 changes: 19 additions & 5 deletions src/crewai/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,6 @@ class Agent(BaseModel):
)
llm: Optional[Any] = Field(
default_factory=lambda: ChatOpenAI(
temperature=0.7,
model_name="gpt-4",
),
description="Language model that will run the agent.",
Expand All @@ -109,6 +108,7 @@ def _deny_user_set_id(cls, v: Optional[UUID4]) -> None:

@model_validator(mode="after")
def set_private_attrs(self):
"""Set private attributes."""
self._logger = Logger(self.verbose)
if self.max_rpm and not self._rpm_controller:
self._rpm_controller = RPMController(
Expand All @@ -118,12 +118,16 @@ def set_private_attrs(self):

@model_validator(mode="after")
def check_agent_executor(self) -> "Agent":
"""Check if the agent executor is set."""
if not self.agent_executor:
self.set_cache_handler(self.cache_handler)
return self

def execute_task(
self, task: str, context: str = None, tools: List[Any] = None
self,
task: str,
context: Optional[str] = None,
tools: Optional[List[Any]] = None,
) -> str:
"""Execute a task with the agent.
Expand Down Expand Up @@ -157,17 +161,27 @@ def execute_task(

return result

def set_cache_handler(self, cache_handler) -> None:
def set_cache_handler(self, cache_handler: CacheHandler) -> None:
"""Set the cache handler for the agent.
Args:
cache_handler: An instance of the CacheHandler class.
"""
self.cache_handler = cache_handler
self.tools_handler = ToolsHandler(cache=self.cache_handler)
self.__create_agent_executor()

def set_rpm_controller(self, rpm_controller) -> None:
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.
"""
if not self._rpm_controller:
self._rpm_controller = rpm_controller
self.__create_agent_executor()

def __create_agent_executor(self) -> CrewAgentExecutor:
def __create_agent_executor(self) -> None:
"""Create an agent executor for the agent.
Returns:
Expand Down
36 changes: 29 additions & 7 deletions src/crewai/crew.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ class Crew(BaseModel):
process: The process flow that the crew will follow (e.g., sequential).
verbose: Indicates the verbosity level for logging during execution.
config: Configuration settings for the crew.
cache_handler: Handles caching for the crew's operations.
_cache_handler: Handles caching for the crew's operations.
max_rpm: Maximum number of requests per minute for the crew execution to be respected.
id: A unique identifier for the crew instance.
"""
Expand Down Expand Up @@ -71,11 +71,22 @@ def _deny_user_set_id(cls, v: Optional[UUID4]) -> None:

@classmethod
@field_validator("config", mode="before")
def check_config_type(cls, v: Union[Json, Dict[str, Any]]):
def check_config_type(
cls, v: Union[Json, Dict[str, Any]]
) -> Union[Json, Dict[str, Any]]:
"""Validates that the config is a valid type.
Args:
v: The config to be validated.
Returns:
The config if it is valid.
"""
return json.loads(v) if isinstance(v, Json) else v

@model_validator(mode="after")
def set_private_attrs(self):
def set_private_attrs(self) -> "Crew":
"""Set private attributes."""
self._cache_handler = CacheHandler()
self._logger = Logger(self.verbose)
self._rpm_controller = RPMController(max_rpm=self.max_rpm, logger=self._logger)
Expand Down Expand Up @@ -110,8 +121,15 @@ def _setup_from_config(self):
self.agents = [Agent(**agent) for agent in self.config["agents"]]
self.tasks = [self._create_task(task) for task in self.config["tasks"]]

def _create_task(self, task_config):
"""Creates a task instance from its configuration."""
def _create_task(self, task_config: Dict[str, Any]) -> Task:
"""Creates a task instance from its configuration.
Args:
task_config: The configuration of the task.
Returns:
A task instance.
"""
task_agent = next(
agt for agt in self.agents if agt.role == task_config["agent"]
)
Expand Down Expand Up @@ -140,8 +158,12 @@ def _sequential_loop(self) -> str:
self._rpm_controller.stop_rpm_counter()
return task_output

def _prepare_and_execute_task(self, task):
"""Prepares and logs information about the task being executed."""
def _prepare_and_execute_task(self, task: Task) -> None:
"""Prepares and logs information about the task being executed.
Args:
task: The task to be executed.
"""
if task.agent.allow_delegation:
task.tools += AgentTools(agents=self.agents).tools()

Expand Down
3 changes: 2 additions & 1 deletion src/crewai/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,12 @@ def _deny_user_set_id(cls, v: Optional[UUID4]) -> None:

@model_validator(mode="after")
def check_tools(self):
"""Check if the tools are set."""
if not self.tools and (self.agent and self.agent.tools):
self.tools.extend(self.agent.tools)
return self

def execute(self, context: str = None) -> str:
def execute(self, context: Optional[str] = None) -> str:
"""Execute the task.
Returns:
Expand Down

0 comments on commit 852ca21

Please sign in to comment.