diff --git a/src/crewai/agent.py b/src/crewai/agent.py index 999d1d800..59226c94e 100644 --- a/src/crewai/agent.py +++ b/src/crewai/agent.py @@ -23,6 +23,7 @@ from crewai.utilities.constants import TRAINED_AGENTS_DATA_FILE, TRAINING_DATA_F from crewai.utilities.converter import generate_model_description from crewai.utilities.token_counter_callback import TokenCalcHandler from crewai.utilities.training_handler import CrewTrainingHandler +from crewai.utilities.typing import AgentConfig agentops = None @@ -88,6 +89,7 @@ class Agent(BaseAgent): function_calling_llm: Optional[Any] = Field( description="Language model that will run the agent.", default=None ) + config: Optional[Union[Dict[str, Any], AgentConfig]] = Field(default=None) system_template: Optional[str] = Field( default=None, description="System format for the agent." ) diff --git a/src/crewai/project/annotations.py b/src/crewai/project/annotations.py index bf0051c4d..e2b6280a0 100644 --- a/src/crewai/project/annotations.py +++ b/src/crewai/project/annotations.py @@ -16,6 +16,12 @@ def after_kickoff(func): def task(func): + """Decorator to mark a method as a task creator. + + When applied to a method in a class decorated with @CrewBase, + this makes the method's return value accessible as an element + of the self.tasks list. + """ func.is_task = True @wraps(func) @@ -29,6 +35,12 @@ def task(func): def agent(func): + """Decorator to mark a method as an agent creator. + + When applied to a method in a class decorated with @CrewBase, + this makes the method's return value accessible as an element + of the self.agents list. + """ func.is_agent = True func = memoize(func) return func diff --git a/src/crewai/project/crew_base.py b/src/crewai/project/crew_base.py index 0b43882f2..161610f81 100644 --- a/src/crewai/project/crew_base.py +++ b/src/crewai/project/crew_base.py @@ -1,6 +1,6 @@ import inspect from pathlib import Path -from typing import Any, Callable, Dict, TypeVar, cast +from typing import Any, Callable, Dict, List, TypeVar, cast import yaml from dotenv import load_dotenv @@ -66,6 +66,9 @@ def CrewBase(cls: T) -> T: self._kickoff = self._filter_functions( self._original_functions, "is_kickoff" ) + + self.agents = [] # type: List[Any] + self.tasks = [] # type: List[Any] @staticmethod def load_yaml(config_path: Path): diff --git a/src/crewai/task.py b/src/crewai/task.py index 30ab79c00..5ecc63640 100644 --- a/src/crewai/task.py +++ b/src/crewai/task.py @@ -41,6 +41,7 @@ from crewai.tools.base_tool import BaseTool from crewai.utilities.config import process_config from crewai.utilities.converter import Converter, convert_to_model from crewai.utilities.i18n import I18N +from crewai.utilities.typing import TaskConfig class Task(BaseModel): @@ -74,7 +75,7 @@ class Task(BaseModel): expected_output: str = Field( description="Clear definition of expected output for the task." ) - config: Optional[Dict[str, Any]] = Field( + config: Optional[Union[Dict[str, Any], TaskConfig]] = Field( description="Configuration for the agent", default=None, ) diff --git a/src/crewai/utilities/typing/__init__.py b/src/crewai/utilities/typing/__init__.py new file mode 100644 index 000000000..5d5fd44af --- /dev/null +++ b/src/crewai/utilities/typing/__init__.py @@ -0,0 +1,14 @@ +from typing import Dict, List, Optional, Any, TypedDict, Union + +class AgentConfig(TypedDict, total=False): + """TypedDict for agent configuration loaded from YAML.""" + role: str + goal: str + backstory: str + verbose: bool + +class TaskConfig(TypedDict, total=False): + """TypedDict for task configuration loaded from YAML.""" + description: str + expected_output: str + agent: str # Role of the agent to execute this task diff --git a/tests/typing_test.py b/tests/typing_test.py new file mode 100644 index 000000000..5e6690213 --- /dev/null +++ b/tests/typing_test.py @@ -0,0 +1,55 @@ +from typing import Dict, Any + +import pytest + +from crewai.agent import Agent +from crewai.task import Task +from crewai.utilities.typing import AgentConfig, TaskConfig + + +def test_agent_with_config_dict(): + config: AgentConfig = { + "role": "Test Agent", + "goal": "Test Goal", + "backstory": "Test Backstory", + "verbose": True + } + + agent = Agent(config=config) + + assert agent.role == "Test Agent" + assert agent.goal == "Test Goal" + assert agent.backstory == "Test Backstory" + assert agent.verbose is True + + +def test_agent_with_yaml_config(): + config: Dict[str, Any] = { + "researcher": { + "role": "Researcher", + "goal": "Research Goal", + "backstory": "Researcher Backstory", + "verbose": True + } + } + + agent = Agent(config=config["researcher"]) + + assert agent.role == "Researcher" + assert agent.goal == "Research Goal" + assert agent.backstory == "Researcher Backstory" + + +def test_task_with_config_dict(): + config: TaskConfig = { + "description": "Test Task", + "expected_output": "Test Output", + "agent": "researcher" + } + + agent = Agent(role="Researcher", goal="Goal", backstory="Backstory") + task = Task(config=config, agent=agent) + + assert task.description == "Test Task" + assert task.expected_output == "Test Output" + assert task.agent == agent