Update some docstrings / typehints (#144)

This commit is contained in:
Greyson LaLonde
2024-01-21 14:55:17 -05:00
committed by GitHub
parent 7a93124cdb
commit e8a31da05d
3 changed files with 50 additions and 13 deletions

View File

@@ -93,7 +93,6 @@ class Agent(BaseModel):
) )
llm: Optional[Any] = Field( llm: Optional[Any] = Field(
default_factory=lambda: ChatOpenAI( default_factory=lambda: ChatOpenAI(
temperature=0.7,
model_name="gpt-4", model_name="gpt-4",
), ),
description="Language model that will run the agent.", description="Language model that will run the agent.",
@@ -109,6 +108,7 @@ class Agent(BaseModel):
@model_validator(mode="after") @model_validator(mode="after")
def set_private_attrs(self): def set_private_attrs(self):
"""Set private attributes."""
self._logger = Logger(self.verbose) self._logger = Logger(self.verbose)
if self.max_rpm and not self._rpm_controller: if self.max_rpm and not self._rpm_controller:
self._rpm_controller = RPMController( self._rpm_controller = RPMController(
@@ -118,12 +118,16 @@ class Agent(BaseModel):
@model_validator(mode="after") @model_validator(mode="after")
def check_agent_executor(self) -> "Agent": def check_agent_executor(self) -> "Agent":
"""Check if the agent executor is set."""
if not self.agent_executor: if not self.agent_executor:
self.set_cache_handler(self.cache_handler) self.set_cache_handler(self.cache_handler)
return self return self
def execute_task( 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: ) -> str:
"""Execute a task with the agent. """Execute a task with the agent.
@@ -157,17 +161,27 @@ class Agent(BaseModel):
return result 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.cache_handler = cache_handler
self.tools_handler = ToolsHandler(cache=self.cache_handler) self.tools_handler = ToolsHandler(cache=self.cache_handler)
self.__create_agent_executor() 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: if not self._rpm_controller:
self._rpm_controller = rpm_controller self._rpm_controller = rpm_controller
self.__create_agent_executor() self.__create_agent_executor()
def __create_agent_executor(self) -> CrewAgentExecutor: def __create_agent_executor(self) -> None:
"""Create an agent executor for the agent. """Create an agent executor for the agent.
Returns: Returns:

View File

@@ -33,7 +33,7 @@ class Crew(BaseModel):
process: The process flow that the crew will follow (e.g., sequential). process: The process flow that the crew will follow (e.g., sequential).
verbose: Indicates the verbosity level for logging during execution. verbose: Indicates the verbosity level for logging during execution.
config: Configuration settings for the crew. 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. max_rpm: Maximum number of requests per minute for the crew execution to be respected.
id: A unique identifier for the crew instance. id: A unique identifier for the crew instance.
""" """
@@ -71,11 +71,22 @@ class Crew(BaseModel):
@classmethod @classmethod
@field_validator("config", mode="before") @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 return json.loads(v) if isinstance(v, Json) else v
@model_validator(mode="after") @model_validator(mode="after")
def set_private_attrs(self): def set_private_attrs(self) -> "Crew":
"""Set private attributes."""
self._cache_handler = CacheHandler() self._cache_handler = CacheHandler()
self._logger = Logger(self.verbose) self._logger = Logger(self.verbose)
self._rpm_controller = RPMController(max_rpm=self.max_rpm, logger=self._logger) self._rpm_controller = RPMController(max_rpm=self.max_rpm, logger=self._logger)
@@ -110,8 +121,15 @@ class Crew(BaseModel):
self.agents = [Agent(**agent) for agent in self.config["agents"]] self.agents = [Agent(**agent) for agent in self.config["agents"]]
self.tasks = [self._create_task(task) for task in self.config["tasks"]] self.tasks = [self._create_task(task) for task in self.config["tasks"]]
def _create_task(self, task_config): def _create_task(self, task_config: Dict[str, Any]) -> Task:
"""Creates a task instance from its configuration.""" """Creates a task instance from its configuration.
Args:
task_config: The configuration of the task.
Returns:
A task instance.
"""
task_agent = next( task_agent = next(
agt for agt in self.agents if agt.role == task_config["agent"] agt for agt in self.agents if agt.role == task_config["agent"]
) )
@@ -140,8 +158,12 @@ class Crew(BaseModel):
self._rpm_controller.stop_rpm_counter() self._rpm_controller.stop_rpm_counter()
return task_output return task_output
def _prepare_and_execute_task(self, task): def _prepare_and_execute_task(self, task: Task) -> None:
"""Prepares and logs information about the task being executed.""" """Prepares and logs information about the task being executed.
Args:
task: The task to be executed.
"""
if task.agent.allow_delegation: if task.agent.allow_delegation:
task.tools += AgentTools(agents=self.agents).tools() task.tools += AgentTools(agents=self.agents).tools()

View File

@@ -39,11 +39,12 @@ class Task(BaseModel):
@model_validator(mode="after") @model_validator(mode="after")
def check_tools(self): def check_tools(self):
"""Check if the tools are set."""
if not self.tools and (self.agent and self.agent.tools): if not self.tools and (self.agent and self.agent.tools):
self.tools.extend(self.agent.tools) self.tools.extend(self.agent.tools)
return self return self
def execute(self, context: str = None) -> str: def execute(self, context: Optional[str] = None) -> str:
"""Execute the task. """Execute the task.
Returns: Returns: