From 21d063a46cf9d9ca90d28c0a116f9208b00f94e5 Mon Sep 17 00:00:00 2001 From: Lucas Gomide Date: Fri, 6 Jun 2025 16:28:09 -0300 Subject: [PATCH] Support multi org in CLI (#2969) * feat: support to list, switch and see your current organization * feat: store the current org after logged in * feat: filtering agents, tools and their actions by organization_uuid if present * fix linter offenses * refactor: propagate the current org thought Header instead of params * refactor: rename org column name to ID instead of Handle --------- Co-authored-by: Tony Kipkemboi --- src/crewai/cli/cli.py | 29 ++++ src/crewai/cli/config.py | 6 + src/crewai/cli/organization/__init__.py | 1 + src/crewai/cli/organization/main.py | 63 ++++++++ src/crewai/cli/plus_api.py | 9 ++ src/crewai/cli/tools/main.py | 6 + tests/cli/organization/__init__.py | 1 + tests/cli/organization/test_main.py | 206 ++++++++++++++++++++++++ tests/cli/test_plus_api.py | 125 +++++++++++++- 9 files changed, 444 insertions(+), 2 deletions(-) create mode 100644 src/crewai/cli/organization/__init__.py create mode 100644 src/crewai/cli/organization/main.py create mode 100644 tests/cli/organization/__init__.py create mode 100644 tests/cli/organization/test_main.py diff --git a/src/crewai/cli/cli.py b/src/crewai/cli/cli.py index c0eff594c..ab53c8bb3 100644 --- a/src/crewai/cli/cli.py +++ b/src/crewai/cli/cli.py @@ -16,6 +16,7 @@ from .deploy.main import DeployCommand from .evaluate_crew import evaluate_crew from .install_crew import install_crew from .kickoff_flow import kickoff_flow +from .organization.main import OrganizationCommand from .plot_flow import plot_flow from .replay_from_task import replay_task_command from .reset_memories_command import reset_memories_command @@ -353,5 +354,33 @@ def chat(): run_chat() +@crewai.group(invoke_without_command=True) +def org(): + """Organization management commands.""" + pass + + +@org.command() +def list(): + """List available organizations.""" + org_command = OrganizationCommand() + org_command.list() + + +@org.command() +@click.argument("id") +def switch(id): + """Switch to a specific organization.""" + org_command = OrganizationCommand() + org_command.switch(id) + + +@org.command() +def current(): + """Show current organization when 'crewai org' is called without subcommands.""" + org_command = OrganizationCommand() + org_command.current() + + if __name__ == "__main__": crewai() diff --git a/src/crewai/cli/config.py b/src/crewai/cli/config.py index 8e30767ca..5d363f16c 100644 --- a/src/crewai/cli/config.py +++ b/src/crewai/cli/config.py @@ -14,6 +14,12 @@ class Settings(BaseModel): tool_repository_password: Optional[str] = Field( None, description="Password for interacting with the Tool Repository" ) + org_name: Optional[str] = Field( + None, description="Name of the currently active organization" + ) + org_uuid: Optional[str] = Field( + None, description="UUID of the currently active organization" + ) config_path: Path = Field(default=DEFAULT_CONFIG_PATH, exclude=True) def __init__(self, config_path: Path = DEFAULT_CONFIG_PATH, **data): diff --git a/src/crewai/cli/organization/__init__.py b/src/crewai/cli/organization/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/src/crewai/cli/organization/__init__.py @@ -0,0 +1 @@ + diff --git a/src/crewai/cli/organization/main.py b/src/crewai/cli/organization/main.py new file mode 100644 index 000000000..233f85714 --- /dev/null +++ b/src/crewai/cli/organization/main.py @@ -0,0 +1,63 @@ +from rich.console import Console +from rich.table import Table + +from crewai.cli.command import BaseCommand, PlusAPIMixin +from crewai.cli.config import Settings + +console = Console() + +class OrganizationCommand(BaseCommand, PlusAPIMixin): + def __init__(self): + BaseCommand.__init__(self) + PlusAPIMixin.__init__(self, telemetry=self._telemetry) + + def list(self): + try: + response = self.plus_api_client.get_organizations() + response.raise_for_status() + orgs = response.json() + + if not orgs: + console.print("You don't belong to any organizations yet.", style="yellow") + return + + table = Table(title="Your Organizations") + table.add_column("Name", style="cyan") + table.add_column("ID", style="green") + for org in orgs: + table.add_row(org["name"], org["uuid"]) + + console.print(table) + except Exception as e: + console.print(f"Failed to retrieve organization list: {str(e)}", style="bold red") + raise SystemExit(1) + + def switch(self, org_id): + try: + response = self.plus_api_client.get_organizations() + response.raise_for_status() + orgs = response.json() + + org = next((o for o in orgs if o["uuid"] == org_id), None) + if not org: + console.print(f"Organization with id '{org_id}' not found.", style="bold red") + return + + settings = Settings() + settings.org_name = org["name"] + settings.org_uuid = org["uuid"] + settings.dump() + + console.print(f"Successfully switched to {org['name']} ({org['uuid']})", style="bold green") + except Exception as e: + console.print(f"Failed to switch organization: {str(e)}", style="bold red") + raise SystemExit(1) + + def current(self): + settings = Settings() + if settings.org_uuid: + console.print(f"Currently logged in to organization {settings.org_name} ({settings.org_uuid})", style="bold green") + else: + console.print("You're not currently logged in to any organization.", style="yellow") + console.print("Use 'crewai org list' to see available organizations.", style="yellow") + console.print("Use 'crewai org switch ' to switch to an organization.", style="yellow") diff --git a/src/crewai/cli/plus_api.py b/src/crewai/cli/plus_api.py index 6961f886e..e34c26b1b 100644 --- a/src/crewai/cli/plus_api.py +++ b/src/crewai/cli/plus_api.py @@ -4,6 +4,7 @@ from urllib.parse import urljoin import requests +from crewai.cli.config import Settings from crewai.cli.version import get_crewai_version @@ -13,6 +14,7 @@ class PlusAPI: """ TOOLS_RESOURCE = "/crewai_plus/api/v1/tools" + ORGANIZATIONS_RESOURCE = "/crewai_plus/api/v1/me/organizations" CREWS_RESOURCE = "/crewai_plus/api/v1/crews" AGENTS_RESOURCE = "/crewai_plus/api/v1/agents" @@ -24,6 +26,9 @@ class PlusAPI: "User-Agent": f"CrewAI-CLI/{get_crewai_version()}", "X-Crewai-Version": get_crewai_version(), } + settings = Settings() + if settings.org_uuid: + self.headers["X-Crewai-Organization-Id"] = settings.org_uuid self.base_url = getenv("CREWAI_BASE_URL", "https://app.crewai.com") def _make_request(self, method: str, endpoint: str, **kwargs) -> requests.Response: @@ -103,3 +108,7 @@ class PlusAPI: def create_crew(self, payload) -> requests.Response: return self._make_request("POST", self.CREWS_RESOURCE, json=payload) + + def get_organizations(self) -> requests.Response: + return self._make_request("GET", self.ORGANIZATIONS_RESOURCE) + \ No newline at end of file diff --git a/src/crewai/cli/tools/main.py b/src/crewai/cli/tools/main.py index cbd99b1eb..92feec3a0 100644 --- a/src/crewai/cli/tools/main.py +++ b/src/crewai/cli/tools/main.py @@ -173,6 +173,12 @@ class ToolCommand(BaseCommand, PlusAPIMixin): settings.tool_repository_password = login_response_json["credential"][ "password" ] + settings.org_uuid = login_response_json["current_organization"][ + "uuid" + ] + settings.org_name = login_response_json["current_organization"][ + "name" + ] settings.dump() console.print( diff --git a/tests/cli/organization/__init__.py b/tests/cli/organization/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/tests/cli/organization/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/cli/organization/test_main.py b/tests/cli/organization/test_main.py new file mode 100644 index 000000000..5f92284ae --- /dev/null +++ b/tests/cli/organization/test_main.py @@ -0,0 +1,206 @@ +import unittest +from unittest.mock import MagicMock, patch, call + +import pytest +from click.testing import CliRunner +import requests + +from crewai.cli.organization.main import OrganizationCommand +from crewai.cli.cli import list, switch, current + + +@pytest.fixture +def runner(): + return CliRunner() + + +@pytest.fixture +def org_command(): + with patch.object(OrganizationCommand, '__init__', return_value=None): + command = OrganizationCommand() + yield command + + +@pytest.fixture +def mock_settings(): + with patch('crewai.cli.organization.main.Settings') as mock_settings_class: + mock_settings_instance = MagicMock() + mock_settings_class.return_value = mock_settings_instance + yield mock_settings_instance + + +@patch('crewai.cli.cli.OrganizationCommand') +def test_org_list_command(mock_org_command_class, runner): + mock_org_instance = MagicMock() + mock_org_command_class.return_value = mock_org_instance + + result = runner.invoke(list) + + assert result.exit_code == 0 + mock_org_command_class.assert_called_once() + mock_org_instance.list.assert_called_once() + + +@patch('crewai.cli.cli.OrganizationCommand') +def test_org_switch_command(mock_org_command_class, runner): + mock_org_instance = MagicMock() + mock_org_command_class.return_value = mock_org_instance + + result = runner.invoke(switch, ['test-id']) + + assert result.exit_code == 0 + mock_org_command_class.assert_called_once() + mock_org_instance.switch.assert_called_once_with('test-id') + + +@patch('crewai.cli.cli.OrganizationCommand') +def test_org_current_command(mock_org_command_class, runner): + mock_org_instance = MagicMock() + mock_org_command_class.return_value = mock_org_instance + + result = runner.invoke(current) + + assert result.exit_code == 0 + mock_org_command_class.assert_called_once() + mock_org_instance.current.assert_called_once() + + +class TestOrganizationCommand(unittest.TestCase): + def setUp(self): + with patch.object(OrganizationCommand, '__init__', return_value=None): + self.org_command = OrganizationCommand() + self.org_command.plus_api_client = MagicMock() + + @patch('crewai.cli.organization.main.console') + @patch('crewai.cli.organization.main.Table') + def test_list_organizations_success(self, mock_table, mock_console): + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = [ + {"name": "Org 1", "uuid": "org-123"}, + {"name": "Org 2", "uuid": "org-456"} + ] + self.org_command.plus_api_client = MagicMock() + self.org_command.plus_api_client.get_organizations.return_value = mock_response + + mock_console.print = MagicMock() + + self.org_command.list() + + self.org_command.plus_api_client.get_organizations.assert_called_once() + mock_table.assert_called_once_with(title="Your Organizations") + mock_table.return_value.add_column.assert_has_calls([ + call("Name", style="cyan"), + call("ID", style="green") + ]) + mock_table.return_value.add_row.assert_has_calls([ + call("Org 1", "org-123"), + call("Org 2", "org-456") + ]) + + @patch('crewai.cli.organization.main.console') + def test_list_organizations_empty(self, mock_console): + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = [] + self.org_command.plus_api_client = MagicMock() + self.org_command.plus_api_client.get_organizations.return_value = mock_response + + self.org_command.list() + + self.org_command.plus_api_client.get_organizations.assert_called_once() + mock_console.print.assert_called_once_with( + "You don't belong to any organizations yet.", + style="yellow" + ) + + @patch('crewai.cli.organization.main.console') + def test_list_organizations_api_error(self, mock_console): + self.org_command.plus_api_client = MagicMock() + self.org_command.plus_api_client.get_organizations.side_effect = requests.exceptions.RequestException("API Error") + + with pytest.raises(SystemExit): + self.org_command.list() + + + self.org_command.plus_api_client.get_organizations.assert_called_once() + mock_console.print.assert_called_once_with( + "Failed to retrieve organization list: API Error", + style="bold red" + ) + + @patch('crewai.cli.organization.main.console') + @patch('crewai.cli.organization.main.Settings') + def test_switch_organization_success(self, mock_settings_class, mock_console): + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = [ + {"name": "Org 1", "uuid": "org-123"}, + {"name": "Test Org", "uuid": "test-id"} + ] + self.org_command.plus_api_client = MagicMock() + self.org_command.plus_api_client.get_organizations.return_value = mock_response + + mock_settings_instance = MagicMock() + mock_settings_class.return_value = mock_settings_instance + + self.org_command.switch("test-id") + + self.org_command.plus_api_client.get_organizations.assert_called_once() + mock_settings_instance.dump.assert_called_once() + assert mock_settings_instance.org_name == "Test Org" + assert mock_settings_instance.org_uuid == "test-id" + mock_console.print.assert_called_once_with( + "Successfully switched to Test Org (test-id)", + style="bold green" + ) + + @patch('crewai.cli.organization.main.console') + def test_switch_organization_not_found(self, mock_console): + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_response.json.return_value = [ + {"name": "Org 1", "uuid": "org-123"}, + {"name": "Org 2", "uuid": "org-456"} + ] + self.org_command.plus_api_client = MagicMock() + self.org_command.plus_api_client.get_organizations.return_value = mock_response + + self.org_command.switch("non-existent-id") + + self.org_command.plus_api_client.get_organizations.assert_called_once() + mock_console.print.assert_called_once_with( + "Organization with id 'non-existent-id' not found.", + style="bold red" + ) + + @patch('crewai.cli.organization.main.console') + @patch('crewai.cli.organization.main.Settings') + def test_current_organization_with_org(self, mock_settings_class, mock_console): + mock_settings_instance = MagicMock() + mock_settings_instance.org_name = "Test Org" + mock_settings_instance.org_uuid = "test-id" + mock_settings_class.return_value = mock_settings_instance + + self.org_command.current() + + self.org_command.plus_api_client.get_organizations.assert_not_called() + mock_console.print.assert_called_once_with( + "Currently logged in to organization Test Org (test-id)", + style="bold green" + ) + + @patch('crewai.cli.organization.main.console') + @patch('crewai.cli.organization.main.Settings') + def test_current_organization_without_org(self, mock_settings_class, mock_console): + mock_settings_instance = MagicMock() + mock_settings_instance.org_uuid = None + mock_settings_class.return_value = mock_settings_instance + + self.org_command.current() + + assert mock_console.print.call_count == 3 + mock_console.print.assert_any_call( + "You're not currently logged in to any organization.", + style="yellow" + ) diff --git a/tests/cli/test_plus_api.py b/tests/cli/test_plus_api.py index da26ba35f..eff57e1a5 100644 --- a/tests/cli/test_plus_api.py +++ b/tests/cli/test_plus_api.py @@ -1,6 +1,6 @@ import os import unittest -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, patch, ANY from crewai.cli.plus_api import PlusAPI @@ -9,6 +9,7 @@ class TestPlusAPI(unittest.TestCase): def setUp(self): self.api_key = "test_api_key" self.api = PlusAPI(self.api_key) + self.org_uuid = "test-org-uuid" def test_init(self): self.assertEqual(self.api.api_key, self.api_key) @@ -29,17 +30,96 @@ class TestPlusAPI(unittest.TestCase): ) self.assertEqual(response, mock_response) + def assert_request_with_org_id(self, mock_make_request, method: str, endpoint: str, **kwargs): + mock_make_request.assert_called_once_with( + method, f"https://app.crewai.com{endpoint}", headers={'Authorization': ANY, 'Content-Type': ANY, 'User-Agent': ANY, 'X-Crewai-Version': ANY, 'X-Crewai-Organization-Id': self.org_uuid}, **kwargs + ) + + @patch("crewai.cli.plus_api.Settings") + @patch("requests.Session.request") + def test_login_to_tool_repository_with_org_uuid(self, mock_make_request, mock_settings_class): + mock_settings = MagicMock() + mock_settings.org_uuid = self.org_uuid + mock_settings_class.return_value = mock_settings + # re-initialize Client + self.api = PlusAPI(self.api_key) + + mock_response = MagicMock() + mock_make_request.return_value = mock_response + + response = self.api.login_to_tool_repository() + + self.assert_request_with_org_id( + mock_make_request, + 'POST', + '/crewai_plus/api/v1/tools/login' + ) + self.assertEqual(response, mock_response) + + @patch("crewai.cli.plus_api.PlusAPI._make_request") + def test_get_agent(self, mock_make_request): + mock_response = MagicMock() + mock_make_request.return_value = mock_response + + response = self.api.get_agent("test_agent_handle") + mock_make_request.assert_called_once_with( + "GET", "/crewai_plus/api/v1/agents/test_agent_handle" + ) + self.assertEqual(response, mock_response) + + @patch("crewai.cli.plus_api.Settings") + @patch("requests.Session.request") + def test_get_agent_with_org_uuid(self, mock_make_request, mock_settings_class): + mock_settings = MagicMock() + mock_settings.org_uuid = self.org_uuid + mock_settings_class.return_value = mock_settings + # re-initialize Client + self.api = PlusAPI(self.api_key) + + mock_response = MagicMock() + mock_make_request.return_value = mock_response + + response = self.api.get_agent("test_agent_handle") + + self.assert_request_with_org_id( + mock_make_request, + "GET", + "/crewai_plus/api/v1/agents/test_agent_handle" + ) + self.assertEqual(response, mock_response) + @patch("crewai.cli.plus_api.PlusAPI._make_request") def test_get_tool(self, mock_make_request): mock_response = MagicMock() mock_make_request.return_value = mock_response response = self.api.get_tool("test_tool_handle") - mock_make_request.assert_called_once_with( "GET", "/crewai_plus/api/v1/tools/test_tool_handle" ) self.assertEqual(response, mock_response) + + @patch("crewai.cli.plus_api.Settings") + @patch("requests.Session.request") + def test_get_tool_with_org_uuid(self, mock_make_request, mock_settings_class): + mock_settings = MagicMock() + mock_settings.org_uuid = self.org_uuid + mock_settings_class.return_value = mock_settings + # re-initialize Client + self.api = PlusAPI(self.api_key) + + # Set up mock response + mock_response = MagicMock() + mock_make_request.return_value = mock_response + + response = self.api.get_tool("test_tool_handle") + + self.assert_request_with_org_id( + mock_make_request, + "GET", + "/crewai_plus/api/v1/tools/test_tool_handle" + ) + self.assertEqual(response, mock_response) @patch("crewai.cli.plus_api.PlusAPI._make_request") def test_publish_tool(self, mock_make_request): @@ -67,6 +147,47 @@ class TestPlusAPI(unittest.TestCase): "POST", "/crewai_plus/api/v1/tools", json=params ) self.assertEqual(response, mock_response) + + @patch("crewai.cli.plus_api.Settings") + @patch("requests.Session.request") + def test_publish_tool_with_org_uuid(self, mock_make_request, mock_settings_class): + mock_settings = MagicMock() + mock_settings.org_uuid = self.org_uuid + mock_settings_class.return_value = mock_settings + # re-initialize Client + self.api = PlusAPI(self.api_key) + + # Set up mock response + mock_response = MagicMock() + mock_make_request.return_value = mock_response + + handle = "test_tool_handle" + public = True + version = "1.0.0" + description = "Test tool description" + encoded_file = "encoded_test_file" + + response = self.api.publish_tool( + handle, public, version, description, encoded_file + ) + + # Expected params including organization_uuid + expected_params = { + "handle": handle, + "public": public, + "version": version, + "file": encoded_file, + "description": description, + "available_exports": None, + } + + self.assert_request_with_org_id( + mock_make_request, + "POST", + "/crewai_plus/api/v1/tools", + json=expected_params + ) + self.assertEqual(response, mock_response) @patch("crewai.cli.plus_api.PlusAPI._make_request") def test_publish_tool_without_description(self, mock_make_request):