Repository: crewAIInc/crewAI
Pull request #2816 (2025-05-12 13:51:15)

feat: support to load an Agent from a repository
Docs will be properly submitted soon!

joaomdmoura: Disclaimer: This review was made by a crew of AI Agents.

## Code Review Comment for PR #2816 - Support to load an Agent from a repository

### Overview
This pull request introduces the capability to load Agent configurations directly from a repository, enhancing modularity and flexibility. It includes additions across several files, complete with test cases that validate the new behavior.

### File-by-File Suggestions

#### 1. `src/crewai/agent.py`
- **Improvements:**
  - **Error Handling**: Enhance the error management in `_load_agent_from_repository` to handle empty inputs and import errors:
    ```python
    @classmethod
    def _load_agent_from_repository(cls, from_repository: str) -> Dict[str, Any]:
        if not from_repository:
            raise ValueError("Repository name cannot be empty")
    
        try:
            import importlib
            # Load agent details..
        except ImportError as e:
            raise AgentRepositoryError(f"Failed to import required modules: {str(e)}")
    ```
  - **Dynamic Import Security**: Introduce an allowed tools list to mitigate risks associated with dynamic imports:
    ```python
    ALLOWED_TOOLS = {"SerperDevTool", "OtherAllowedTool"}
    if tool_name not in ALLOWED_TOOLS:
        raise AgentRepositoryError(f"Tool {tool_name} is not in allowed tools list")
    ```

#### 2. `src/crewai/cli/authentication/token.py`
- **Improvements:**
  - Utilize custom exceptions for clearer error handling when tokens are missing:
    ```python
    from .errors import AuthenticationError

    def get_auth_token() -> str:
        access_token: Optional[str] = TokenManager().get_token()
        if not access_token:
            raise AuthenticationError("No token found, make sure you are logged in")
        return access_token
    ```

#### 3. `src/crewai/cli/plus_api.py`
- **Improvements:**
  - Consider adding a timeout parameter to the `_make_request` function to avoid long delays:
    ```python
    def _make_request(self, method: str, path: str, timeout: int = 30, **kwargs):
        return requests.request(method, f"{self.API_URL}{path}", headers=self._get_headers(), timeout=timeout, **kwargs)
    ```

#### 4. `src/crewai/utilities/errors.py`
- **Improvements:**
  - Extend the `AgentRepositoryError` class to include a status code, enhancing the level of detail in error handling:
    ```python
    class AgentRepositoryError(Exception):
        def __init__(self, message: str, status_code: Optional[int] = None):
            self.status_code = status_code
            super().__init__(message)
    ```

#### 5. `tests/agent_test.py`
- **Improvements:**
  - Organize tests for clarity and better maintainability. Utilize pytest markers to group tests:
    ```python
    @pytest.mark.repository
    class TestAgentRepository:
        ...
    ```

### General Recommendations
1. **Documentation**: Ensure all new methods and classes are well-documented, including examples for clarity.
2. **Error Handling**: Continue refining exception types and messages for better context during failures.
3. **Testing**: Expand test coverage for edge cases, especially involving network issues and integration tests for complete workflows.
4. **Security**: Adhere to best practices regarding dynamic imports and include validation checks where applicable.

### Conclusion
The PR significantly enhances the load functionality for agents. While the changes are well-implemented, following the suggestions for improving error handling, security, and documentation will bolster the overall robustness and maintainability of the codebase. Happy to discuss further or assist with implementation details!

Patch:
diff --git a/src/crewai/agent.py b/src/crewai/agent.py
index dc637967..9a16b704 100644
--- a/src/crewai/agent.py
+++ b/src/crewai/agent.py
@@ -25,6 +25,7 @@ from crewai.utilities.agent_utils import (
 )
 from crewai.utilities.constants import TRAINED_AGENTS_DATA_FILE, TRAINING_DATA_FILE
 from crewai.utilities.converter import generate_model_description
+from crewai.utilities.errors import AgentRepositoryError
 from crewai.utilities.events.agent_events import (
     AgentExecutionCompletedEvent,
     AgentExecutionErrorEvent,
@@ -134,6 +135,49 @@ class Agent(BaseAgent):
         default=None,
         description="Knowledge search query for the agent dynamically generated by the agent.",
     )
+    from_repository: Optional[str] = Field(
+        default=None,
+        description="The Agent's role to be used from your repository.",
+    )
+
+    @model_validator(mode="before")
+    def validate_from_repository(cls, v):
+        if v is not None and (from_repository := v.get("from_repository")):
+            return cls._load_agent_from_repository(from_repository) | v
+        return v
+
+    @classmethod
+    def _load_agent_from_repository(cls, from_repository: str) -> Dict[str, Any]:
+        attributes: Dict[str, Any] = {}
+        if from_repository:
+            import importlib
+
+            from crewai.cli.authentication.token import get_auth_token
+            from crewai.cli.plus_api import PlusAPI
+
+            client = PlusAPI(api_key=get_auth_token())
+            response = client.get_agent(from_repository)
+            if response.status_code != 200:
+                raise AgentRepositoryError(
+                    f"Agent {from_repository} could not be loaded: {response.text}"
+                )
+
+            agent = response.json()
+            for key, value in agent.items():
+                if key == "tools":
+                    attributes[key] = []
+                    for tool_name in value:
+                        try:
+                            module = importlib.import_module("crewai_tools")
+                            tool_class = getattr(module, tool_name)
+                            attributes[key].append(tool_class())
+                        except Exception as e:
+                            raise AgentRepositoryError(
+                                f"Tool {tool_name} could not be loaded: {e}"
+                            ) from e
+                else:
+                    attributes[key] = value
+        return attributes
 
     @model_validator(mode="after")
     def post_init_setup(self):
diff --git a/src/crewai/cli/authentication/token.py b/src/crewai/cli/authentication/token.py
index 30a33b4b..46ff75e8 100644
--- a/src/crewai/cli/authentication/token.py
+++ b/src/crewai/cli/authentication/token.py
@@ -5,5 +5,5 @@ def get_auth_token() -> str:
     """Get the authentication token."""
     access_token = TokenManager().get_token()
     if not access_token:
-        raise Exception()
+        raise Exception("No token found, make sure you are logged in")
     return access_token
diff --git a/src/crewai/cli/plus_api.py b/src/crewai/cli/plus_api.py
index 23032ca8..93e5750c 100644
--- a/src/crewai/cli/plus_api.py
+++ b/src/crewai/cli/plus_api.py
@@ -14,6 +14,7 @@ class PlusAPI:
 
     TOOLS_RESOURCE = "/crewai_plus/api/v1/tools"
     CREWS_RESOURCE = "/crewai_plus/api/v1/crews"
+    AGENTS_RESOURCE = "/crewai_plus/api/v1/agents"
 
     def __init__(self, api_key: str) -> None:
         self.api_key = api_key
@@ -37,6 +38,9 @@ class PlusAPI:
     def get_tool(self, handle: str):
         return self._make_request("GET", f"{self.TOOLS_RESOURCE}/{handle}")
 
+    def get_agent(self, handle: str):
+        return self._make_request("GET", f"{self.AGENTS_RESOURCE}/{handle}")
+
     def publish_tool(
         self,
         handle: str,
diff --git a/src/crewai/utilities/errors.py b/src/crewai/utilities/errors.py
index f673c060..16c59321 100644
--- a/src/crewai/utilities/errors.py
+++ b/src/crewai/utilities/errors.py
@@ -1,4 +1,5 @@
 """Error message definitions for CrewAI database operations."""
+
 from typing import Optional
 
 
@@ -37,3 +38,9 @@ class DatabaseError:
             The formatted error message
         """
         return template.format(str(error))
+
+
+class AgentRepositoryError(Exception):
+    """Exception raised when an agent repository is not found."""
+
+    ...


Test patch:
diff --git a/tests/agent_test.py b/tests/agent_test.py
index faad4ca8..8cd2c075 100644
--- a/tests/agent_test.py
+++ b/tests/agent_test.py
@@ -2,7 +2,7 @@
 
 import os
 from unittest import mock
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -18,6 +18,7 @@ from crewai.tools import tool
 from crewai.tools.tool_calling import InstructorToolCalling
 from crewai.tools.tool_usage import ToolUsage
 from crewai.utilities import RPMController
+from crewai.utilities.errors import AgentRepositoryError
 from crewai.utilities.events import crewai_event_bus
 from crewai.utilities.events.tool_usage_events import ToolUsageFinishedEvent
 
@@ -308,9 +309,7 @@ def test_cache_hitting():
     def handle_tool_end(source, event):
         received_events.append(event)
 
-    with (
-        patch.object(CacheHandler, "read") as read,
-    ):
+    with (patch.object(CacheHandler, "read") as read,):
         read.return_value = "0"
         task = Task(
             description="What is 2 times 6? Ignore correctness and just return the result of the multiplication tool, you must use the tool.",
@@ -1040,7 +1039,7 @@ def test_agent_human_input():
             CrewAgentExecutor,
             "_invoke_loop",
             return_value=AgentFinish(output="Hello", thought="", text=""),
-        ) as mock_invoke_loop,
+        ),
     ):
         # Execute the task
         output = agent.execute_task(task)
@@ -2025,3 +2024,86 @@ def test_get_knowledge_search_query():
                 },
             ]
         )
+
+
+@pytest.fixture
+def mock_get_auth_token():
+    with patch(
+        "crewai.cli.authentication.token.get_auth_token", return_value="test_token"
+    ):
+        yield
+
+
+@patch("crewai.cli.plus_api.PlusAPI.get_agent")
+def test_agent_from_repository(mock_get_agent, mock_get_auth_token):
+    from crewai_tools import SerperDevTool
+
+    mock_get_response = MagicMock()
+    mock_get_response.status_code = 200
+    mock_get_response.json.return_value = {
+        "role": "test role",
+        "goal": "test goal",
+        "backstory": "test backstory",
+        "tools": ["SerperDevTool"],
+    }
+    mock_get_agent.return_value = mock_get_response
+    agent = Agent(from_repository="test_agent")
+
+    assert agent.role == "test role"
+    assert agent.goal == "test goal"
+    assert agent.backstory == "test backstory"
+    assert len(agent.tools) == 1
+    assert isinstance(agent.tools[0], SerperDevTool)
+
+
+@patch("crewai.cli.plus_api.PlusAPI.get_agent")
+def test_agent_from_repository_override_attributes(mock_get_agent, mock_get_auth_token):
+    from crewai_tools import SerperDevTool
+
+    mock_get_response = MagicMock()
+    mock_get_response.status_code = 200
+    mock_get_response.json.return_value = {
+        "role": "test role",
+        "goal": "test goal",
+        "backstory": "test backstory",
+        "tools": ["SerperDevTool"],
+    }
+    mock_get_agent.return_value = mock_get_response
+    agent = Agent(from_repository="test_agent", role="Custom Role")
+
+    assert agent.role == "Custom Role"
+    assert agent.goal == "test goal"
+    assert agent.backstory == "test backstory"
+    assert len(agent.tools) == 1
+    assert isinstance(agent.tools[0], SerperDevTool)
+
+
+@patch("crewai.cli.plus_api.PlusAPI.get_agent")
+def test_agent_from_repository_with_invalid_tools(mock_get_agent, mock_get_auth_token):
+    mock_get_response = MagicMock()
+    mock_get_response.status_code = 200
+    mock_get_response.json.return_value = {
+        "role": "test role",
+        "goal": "test goal",
+        "backstory": "test backstory",
+        "tools": ["DoesNotExist"],
+    }
+    mock_get_agent.return_value = mock_get_response
+    with pytest.raises(
+        AgentRepositoryError,
+        match="Tool DoesNotExist could not be loaded: module 'crewai_tools' has no attribute 'DoesNotExist'",
+    ):
+        Agent(from_repository="test_agent")
+
+
+@patch("crewai.cli.plus_api.PlusAPI.get_agent")
+def test_agent_from_repository_agent_not_found(mock_get_agent, mock_get_auth_token):
+    mock_get_response = MagicMock()
+    mock_get_response.status_code = 404
+    mock_get_response.text = "Agent not found"
+    mock_get_agent.return_value = mock_get_response
+    with pytest.raises(
+        AgentRepositoryError,
+        match="Agent NOT_FOUND could not be loaded: Agent not found",
+    ):
+        Agent(from_repository="NOT_FOUND")