Tool auth (#30)

Note - This Engine PR must go first: https://github.com/ArcadeAI/Engine/pull/65 In this PR: - Add `client.tool.authorize` to authorize a tool by name by @Spartee - Refactored client.auth methods to always pass around scopes (as needed by the above Engine PR) by @nbarbettini - Reduced the scopes needed in the Slack toolkit, which was blocked by this until now! @nbarbettini --------- Co-authored-by: Nate Barbettini <nate@arcade-ai.com>
2024-09-09 15:00:17 -07:00 · 2024-09-09 15:00:17 -07:00 · d12542db55
commit d12542db55
parent 408f2e6300
7 changed files with 385 additions and 79 deletions
--- a/arcade/arcade/cli/main.py
+++ b/arcade/arcade/cli/main.py
@ -215,7 +215,9 @@ def dev(
    host: str = typer.Option(
        "127.0.0.1", help="Host for the app, from settings by default.", show_default=True
    ),
-    port: int = typer.Option("8000", help="Port for the app, defaults to ", show_default=True),
+    port: int = typer.Option(
+        "8000", "-p", "--port", help="Port for the app, defaults to ", show_default=True
+    ),
 ) -> None:
    """
    Starts the actor with host, port, and reload options. Uses
--- a/arcade/arcade/client/base.py
+++ b/arcade/arcade/client/base.py
@ -5,9 +5,21 @@ from urllib.parse import urljoin
 import httpx
 from httpx import Timeout

+from arcade.client.errors import (
+    BadRequestError,
+    InternalServerError,
+    NotFoundError,
+    PermissionDeniedError,
+    RateLimitError,
+    UnauthorizedError,
+)
+
 T = TypeVar("T")
 ResponseT = TypeVar("ResponseT")

+API_VERSION = "v1"
+BASE_URL = "http://localhost:9099"
+

 class BaseResource(Generic[T]):
    """Base class for all resources."""
@ -21,7 +33,7 @@ class BaseArcadeClient:

    def __init__(
        self,
-        base_url: str,
+        base_url: str = BASE_URL,
        api_key: str | None = None,
        headers: dict[str, str] | None = None,
        timeout: float | Timeout = 10.0,
@ -51,6 +63,25 @@ class BaseArcadeClient:
        """
        return urljoin(self._base_url, path)

+    def _chat_url(self, base_url: str) -> str:
+        chat_url = str(base_url)
+        if not base_url.endswith(API_VERSION):
+            chat_url = f"{base_url}/{API_VERSION}"
+        return chat_url
+
+    def _handle_http_error(self, e: httpx.HTTPStatusError) -> None:
+        error_map = {
+            400: BadRequestError,
+            401: UnauthorizedError,
+            403: PermissionDeniedError,
+            404: NotFoundError,
+            429: RateLimitError,
+            500: InternalServerError,
+        }
+        status_code = e.response.status_code
+        error_class = error_map.get(status_code, InternalServerError)
+        raise error_class(str(e), response=e.response)
+

 class SyncArcadeClient(BaseArcadeClient):
    """Synchronous Arcade client."""
--- a/arcade/arcade/client/client.py
+++ b/arcade/arcade/client/client.py
@ -1,17 +1,14 @@
-from typing import Any, Generic, TypeVar
+from typing import Any, TypeVar, Union

 import httpx
 from openai import AsyncOpenAI, OpenAI
 from openai.resources.chat import AsyncChat, Chat

-from arcade.client.base import AsyncArcadeClient, BaseResource, SyncArcadeClient
-from arcade.client.errors import (
-    APIStatusError,
-    BadRequestError,
-    InternalServerError,
-    NotFoundError,
-    PermissionDeniedError,
-    UnauthorizedError,
+from arcade.client.base import (
+    API_VERSION,
+    AsyncArcadeClient,
+    BaseResource,
+    SyncArcadeClient,
 )
 from arcade.client.schema import (
    AuthProvider,
@ -24,9 +21,6 @@ from arcade.core.schema import ToolDefinition
 T = TypeVar("T")
 ClientT = TypeVar("ClientT", SyncArcadeClient, AsyncArcadeClient)

-API_VERSION = "v1"
-BASE_URL = "https://api.arcade-ai.com"
-

 class AuthResource(BaseResource[ClientT]):
    """Authentication resource."""
@ -54,7 +48,7 @@ class AuthResource(BaseResource[ClientT]):
        body = {
            "auth_requirement": {
                "provider": auth_provider,
-                auth_provider: AuthRequest(scope=scopes, authority=authority).dict(
+                auth_provider: AuthRequest(scope=scopes, authority=authority).model_dump(
                    exclude_none=True
                ),
            },
@ -68,16 +62,28 @@ class AuthResource(BaseResource[ClientT]):
        )
        return AuthResponse(**data)

-    def poll_authorization(self, auth_id: str) -> AuthResponse:
+    def status(
+        self, auth_id_or_response: Union[str, AuthResponse], scopes: list[str] | None = None
+    ) -> AuthResponse:
        """Poll for the status of an authorization

-        Polls using the authorization ID returned from the authorize method.
+        Polls using either the authorization ID or the data returned from the authorize method.

        Example:
-            auth_status = client.auth.poll_authorization("auth_123")
+            auth_response = client.auth.authorize(...)
+            auth_status = client.auth.poll_authorization(auth_response)
+            auth_status = client.auth.poll_authorization("auth_123", ["scope1", "scope2"])
        """
+        if isinstance(auth_id_or_response, AuthResponse):
+            auth_id = auth_id_or_response.auth_id
+            scopes = auth_id_or_response.scopes
+        else:
+            auth_id = auth_id_or_response
+
        data = self._client._execute_request(  # type: ignore[attr-defined]
-            "GET", f"{self._base_path}/status", params={"authorizationID": auth_id}
+            "GET",
+            f"{self._base_path}/status",
+            params={"authorizationID": auth_id, "scopes": " ".join(scopes) if scopes else None},
        )
        return AuthResponse(**data)

@ -125,47 +131,116 @@ class ToolResource(BaseResource[ClientT]):
        )
        return ToolDefinition(**data)

+    def authorize(self, tool_name: str, user_id: str) -> AuthResponse:
+        """
+        Get the authorization status for a tool.
+        """
+        data = self._client._execute_request(  # type: ignore[attr-defined]
+            "POST",
+            f"{self._base_path}/authorize",
+            json={"tool_name": tool_name, "user_id": user_id},
+        )
+        return AuthResponse(**data)

-class ArcadeClientMixin(Generic[ClientT]):
-    """Mixin for Arcade clients."""

-    def __init__(self, base_url: str = BASE_URL, *args: Any, **kwargs: Any):
-        super().__init__(base_url, *args, **kwargs)  # type: ignore[call-arg]
-        self.auth: AuthResource = AuthResource(self)
-        self.tool: ToolResource = ToolResource(self)
+class AsyncAuthResource(BaseResource[AsyncArcadeClient]):
+    """Asynchronous Authentication resource."""

-    def _handle_http_error(
+    _base_path = f"/{API_VERSION}/auth"
+
+    async def authorize(
        self,
-        e: httpx.HTTPStatusError,
-        error_map: dict[int, type[APIStatusError]],
-    ) -> None:
-        status_code = e.response.status_code
-        error_class = error_map.get(status_code, InternalServerError)
-        raise error_class(str(e), response=e.response)
+        provider: AuthProvider,
+        scopes: list[str],
+        user_id: str,
+        authority: str | None = None,
+    ) -> AuthResponse:
+        """
+        Initiate an asynchronous authorization request.
+        """
+        auth_provider = provider.value

-    def _chat_url(self, base_url: str) -> str:
-        # TODO (sam): make chat a Resource like others but maintain
-        # the ability to call chat directly like the openai clients
-        chat_url = str(base_url)
-        if not base_url.endswith(API_VERSION):
-            chat_url = f"{base_url}/{API_VERSION}"
-        return chat_url
+        body = {
+            "auth_requirement": {
+                "provider": auth_provider,
+                auth_provider: AuthRequest(scope=scopes, authority=authority).model_dump(
+                    exclude_none=True
+                ),
+            },
+            "user_id": user_id,
+        }
+
+        data = await self._client._execute_request(  # type: ignore[attr-defined]
+            "POST",
+            f"{self._base_path}/authorize",
+            json=body,
+        )
+        return AuthResponse(**data)
+
+    async def status(self, auth_id: str) -> AuthResponse:
+        """Poll for the status of an authorization asynchronously"""
+        data = await self._client._execute_request(  # type: ignore[attr-defined]
+            "GET", f"{self._base_path}/status", params={"authorizationID": auth_id}
+        )
+        return AuthResponse(**data)


-class Arcade(ArcadeClientMixin[SyncArcadeClient], SyncArcadeClient):
-    """Synchronous Arcade client.
+class AsyncToolResource(BaseResource[AsyncArcadeClient]):
+    """Asynchronous Tool resource."""

-    Example:
-        from arcade.client import Arcade
+    _base_path = f"/{API_VERSION}/tools"

-        client = Arcade(api_key="your-api-key")
-        client.auth.authorize(...)
-        client.tool.run(...)
-    """
+    async def run(
+        self,
+        tool_name: str,
+        user_id: str,
+        tool_version: str | None = None,
+        inputs: dict[str, Any] | None = None,
+    ) -> ExecuteToolResponse:
+        """
+        Send an asynchronous request to execute a tool and return the response.
+        """
+        request_data = {
+            "tool_name": tool_name,
+            "user_id": user_id,
+            "tool_version": tool_version,
+            "inputs": inputs,
+        }
+        data = await self._client._execute_request(  # type: ignore[attr-defined]
+            "POST", f"{self._base_path}/execute", json=request_data
+        )
+        return ExecuteToolResponse(**data)
+
+    async def get(self, director_id: str, tool_id: str) -> ToolDefinition:
+        """
+        Get the specification for a tool asynchronously.
+        """
+        data = await self._client._execute_request(  # type: ignore[attr-defined]
+            "GET",
+            f"{self._base_path}/definition",
+            params={"director_id": director_id, "tool_id": tool_id},
+        )
+        return ToolDefinition(**data)
+
+    async def authorize(self, tool_name: str, user_id: str) -> AuthResponse:
+        """
+        Get the authorization status for a tool.
+        """
+        data = await self._client._execute_request(  # type: ignore[attr-defined]
+            "POST",
+            f"{self._base_path}/authorize",
+            json={"tool_name": tool_name, "user_id": user_id},
+        )
+        return AuthResponse(**data)
+
+
+class Arcade(SyncArcadeClient):
+    """Synchronous Arcade client."""

    def __init__(self, *args: Any, **kwargs: Any):
        super().__init__(*args, **kwargs)
-
+        self.auth: AuthResource = AuthResource(self)
+        self.tool: ToolResource = ToolResource(self)
        chat_url = self._chat_url(self._base_url)
        self._openai_client = OpenAI(base_url=chat_url, api_key=self._api_key)

@ -181,24 +256,16 @@ class Arcade(ArcadeClientMixin[SyncArcadeClient], SyncArcadeClient):
            response = self._request(method, url, **kwargs)
            return response.json()
        except httpx.HTTPStatusError as e:
-            self._handle_http_error(
-                e,
-                {
-                    400: BadRequestError,
-                    401: UnauthorizedError,
-                    403: PermissionDeniedError,
-                    404: NotFoundError,
-                    500: InternalServerError,
-                },
-            )
+            self._handle_http_error(e)


-class AsyncArcade(ArcadeClientMixin[AsyncArcadeClient], AsyncArcadeClient):
+class AsyncArcade(AsyncArcadeClient):
    """Asynchronous Arcade client."""

    def __init__(self, *args: Any, **kwargs: Any):
        super().__init__(*args, **kwargs)
-
+        self.auth: AsyncAuthResource = AsyncAuthResource(self)
+        self.tool: AsyncToolResource = AsyncToolResource(self)
        chat_url = self._chat_url(self._base_url)
        self._openai_client = AsyncOpenAI(base_url=chat_url, api_key=self._api_key)

@ -214,13 +281,4 @@ class AsyncArcade(ArcadeClientMixin[AsyncArcadeClient], AsyncArcadeClient):
            response = await self._request(method, url, **kwargs)
            return response.json()
        except httpx.HTTPStatusError as e:
-            self._handle_http_error(
-                e,
-                {
-                    400: BadRequestError,
-                    401: UnauthorizedError,
-                    403: PermissionDeniedError,
-                    404: NotFoundError,
-                    500: InternalServerError,
-                },
-            )
+            self._handle_http_error(e)
--- a/arcade/arcade/client/schema.py
+++ b/arcade/arcade/client/schema.py
@ -2,7 +2,7 @@ from enum import Enum

 from pydantic import AnyUrl, BaseModel, Field

-from arcade.core.schema import ToolCallOutput, ToolContext
+from arcade.core.schema import ToolAuthorizationContext, ToolCallOutput


 class AuthProvider(str, Enum):
@ -48,6 +48,9 @@ class AuthResponse(BaseModel):
    auth_id: str = Field(alias="authorizationID")
    """The ID of the authorization request"""

+    scopes: list[str]
+    """The scope(s) requested in the authorization request"""
+
    # TODO: Use AnyUrl?
    auth_url: str | None = Field(None, alias="authorizationURL")
    """The URL for the authorization"""
@ -55,7 +58,7 @@ class AuthResponse(BaseModel):
    status: AuthStatus
    """Only completed implies presence of a token"""

-    context: ToolContext | None = None
+    context: ToolAuthorizationContext | None = None


 class ExecuteToolResponse(BaseModel):
--- a/arcade/tests/client/test_client.py
+++ b/arcade/tests/client/test_client.py
@ -0,0 +1,217 @@
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+from httpx import HTTPStatusError, Response
+
+from arcade.client import Arcade, AsyncArcade, AuthProvider
+from arcade.client.errors import (
+    BadRequestError,
+    InternalServerError,
+    NotFoundError,
+    PermissionDeniedError,
+    UnauthorizedError,
+)
+from arcade.client.schema import AuthResponse, ExecuteToolResponse
+from arcade.core.schema import ToolDefinition
+
+AUTH_RESPONSE_DATA = {
+    "auth_id": "auth_123",
+    "auth_url": "https://example.com/auth",
+    "status": "pending",
+    "authorizationID": "auth_123",
+    "scopes": ["https://www.googleapis.com/auth/gmail.readonly"],
+}
+
+TOOL_RESPONSE_DATA = {
+    "tool_name": "GetEmails",
+    "tool_version": "0.1.0",
+    "output": {"result": "Hello, World!"},
+    "error": None,
+    "invocation_id": "inv_123",
+    "duration": 1.5,
+    "finished_at": "2023-04-01T12:00:00Z",
+    "success": True,
+}
+
+TOOL_DEFINITION_DATA = {
+    "name": "GetEmails",
+    "description": "Retrieve emails from a user's inbox",
+    "input_schema": {"type": "object", "properties": {"n_emails": {"type": "integer"}}},
+    "output_schema": {"type": "array", "items": {"type": "string"}},
+    "version": "0.1.0",
+    "inputs": {"parameters": []},  # Update this line
+    "output": {},
+    "requirements": {"auth_requirements": []},  # Update this line
+}
+
+TOOL_AUTHORIZE_RESPONSE_DATA = {
+    "authorizationID": "auth_456",
+    "authorizationURL": "https://example.com/auth",
+    "scopes": ["scope1", "scope2"],
+    "status": "pending",
+}
+
+
+@pytest.fixture
+def mock_response():
+    """Mock Response object for testing."""
+    response = Mock(spec=Response)
+    response.json.return_value = {}
+    return response
+
+
+@pytest.fixture
+def mock_async_response():
+    """Mock AsyncResponse object for testing."""
+    response = AsyncMock(spec=Response)
+    response.json.return_value = {}
+    return response
+
+
+@pytest.mark.parametrize(
+    "error_code, expected_error",
+    [
+        (400, BadRequestError),
+        (401, UnauthorizedError),
+        (403, PermissionDeniedError),
+        (404, NotFoundError),
+        (500, InternalServerError),
+    ],
+)
+def test_handle_http_error(error_code, expected_error, mock_response):
+    """Test _handle_http_error method for different error codes."""
+    mock_response.status_code = error_code
+    mock_response.json.return_value = {"error": "Test error message"}
+
+    # Create a mock HTTPStatusError
+    mock_http_error = Mock(spec=HTTPStatusError)
+    mock_http_error.response = mock_response
+
+    client = Arcade()  # Create an instance of Arcade
+    with pytest.raises(expected_error):
+        client._handle_http_error(mock_http_error)  # Call the method on the instance
+
+
+def test_arcade_auth_authorize(mock_response, monkeypatch):
+    """Test Arcade.auth.authorize method."""
+    monkeypatch.setattr(Arcade, "_execute_request", lambda *args, **kwargs: AUTH_RESPONSE_DATA)
+    client = Arcade()
+    auth_response = client.auth.authorize(
+        provider=AuthProvider.google,
+        scopes=["https://www.googleapis.com/auth/gmail.readonly"],
+        user_id="sam@arcade-ai.com",
+    )
+    assert auth_response == AuthResponse(**AUTH_RESPONSE_DATA)
+
+
+def test_arcade_auth_poll_authorization(mock_response, monkeypatch):
+    """Test Arcade.auth.poll_authorization method."""
+    monkeypatch.setattr(Arcade, "_execute_request", lambda *args, **kwargs: AUTH_RESPONSE_DATA)
+    client = Arcade()
+    auth_response = client.auth.status("auth_123")
+    assert auth_response == AuthResponse(**AUTH_RESPONSE_DATA)
+
+
+def test_arcade_tool_run(mock_response, monkeypatch):
+    """Test Arcade.tool.run method."""
+    monkeypatch.setattr(Arcade, "_execute_request", lambda *args, **kwargs: TOOL_RESPONSE_DATA)
+    client = Arcade()
+    tool_response = client.tool.run(
+        tool_name="GetEmails",
+        user_id="sam@arcade-ai.com",
+        tool_version="0.1.0",
+        inputs={"n_emails": 5},
+    )
+    assert tool_response == ExecuteToolResponse(**TOOL_RESPONSE_DATA)
+
+
+def test_arcade_tool_get(mock_response, monkeypatch):
+    """Test Arcade.tool.get method."""
+    monkeypatch.setattr(Arcade, "_execute_request", lambda *args, **kwargs: TOOL_DEFINITION_DATA)
+    client = Arcade()
+    tool_definition = client.tool.get(director_id="default", tool_id="GetEmails")
+    assert tool_definition == ToolDefinition(**TOOL_DEFINITION_DATA)
+
+
+def test_arcade_tool_authorize(mock_response, monkeypatch):
+    """Test Arcade.tool.authorize method."""
+    monkeypatch.setattr(
+        Arcade, "_execute_request", lambda *args, **kwargs: TOOL_AUTHORIZE_RESPONSE_DATA
+    )
+    client = Arcade()
+    auth_response = client.tool.authorize(tool_name="GetEmails", user_id="sam@arcade-ai.com")
+    assert auth_response == AuthResponse(**TOOL_AUTHORIZE_RESPONSE_DATA)
+
+
+@pytest.mark.asyncio
+async def test_async_arcade_auth_authorize(mock_async_response, monkeypatch):
+    """Test AsyncArcade.auth.authorize method."""
+
+    async def mock_execute_request(*args, **kwargs):
+        return AUTH_RESPONSE_DATA
+
+    monkeypatch.setattr(AsyncArcade, "_execute_request", mock_execute_request)
+    client = AsyncArcade()
+    auth_response = await client.auth.authorize(
+        provider=AuthProvider.google,
+        scopes=["https://www.googleapis.com/auth/gmail.readonly"],
+        user_id="sam@arcade-ai.com",
+    )
+    assert auth_response == AuthResponse(**AUTH_RESPONSE_DATA)
+
+
+@pytest.mark.asyncio
+async def test_async_arcade_auth_poll_authorization(mock_async_response, monkeypatch):
+    """Test AsyncArcade.auth.poll_authorization method."""
+
+    async def mock_execute_request(*args, **kwargs):
+        return AUTH_RESPONSE_DATA
+
+    monkeypatch.setattr(AsyncArcade, "_execute_request", mock_execute_request)
+    client = AsyncArcade()
+    auth_response = await client.auth.status("auth_123")
+    assert auth_response == AuthResponse(**AUTH_RESPONSE_DATA)
+
+
+@pytest.mark.asyncio
+async def test_async_arcade_tool_run(mock_async_response, monkeypatch):
+    """Test AsyncArcade.tool.run method."""
+
+    async def mock_execute_request(*args, **kwargs):
+        return TOOL_RESPONSE_DATA
+
+    monkeypatch.setattr(AsyncArcade, "_execute_request", mock_execute_request)
+    client = AsyncArcade()
+    tool_response = await client.tool.run(
+        tool_name="GetEmails",
+        user_id="sam@arcade-ai.com",
+        tool_version="0.1.0",
+        inputs={"n_emails": 5},
+    )
+    assert tool_response == ExecuteToolResponse(**TOOL_RESPONSE_DATA)
+
+
+@pytest.mark.asyncio
+async def test_async_arcade_tool_get(mock_async_response, monkeypatch):
+    """Test AsyncArcade.tool.get method."""
+
+    async def mock_execute_request(*args, **kwargs):
+        return TOOL_DEFINITION_DATA
+
+    monkeypatch.setattr(AsyncArcade, "_execute_request", mock_execute_request)
+    client = AsyncArcade()
+    tool_definition = await client.tool.get(director_id="default", tool_id="GetEmails")
+    assert tool_definition == ToolDefinition(**TOOL_DEFINITION_DATA)
+
+
+@pytest.mark.asyncio
+async def test_async_arcade_tool_authorize(mock_async_response, monkeypatch):
+    """Test AsyncArcade.tool.authorize method."""
+
+    async def mock_execute_request(*args, **kwargs):
+        return TOOL_AUTHORIZE_RESPONSE_DATA
+
+    monkeypatch.setattr(AsyncArcade, "_execute_request", mock_execute_request)
+    client = AsyncArcade()
+    auth_response = await client.tool.authorize(tool_name="GetEmails", user_id="sam@arcade-ai.com")
+    assert auth_response == AuthResponse(**TOOL_AUTHORIZE_RESPONSE_DATA)
--- a/examples/langchain/gmail.py
+++ b/examples/langchain/gmail.py
@ -41,13 +41,13 @@ challenge = client.auth.authorize(
 if challenge.status != "completed":
    print(f"Please visit this URL to authorize: {challenge.auth_url}")
    input("Press Enter after you've completed the authorization...")
-    challenge = client.auth.poll_authorization(challenge.auth_id)
+    challenge = client.auth.poll_authorization(challenge)
    if challenge.status != "completed":
        print("Authorization not completed. Please try again.")
        exit(1)


-creds = Credentials(challenge.context.authorization.token)
+creds = Credentials(challenge.context.token)
 api_resource = build_resource_service(credentials=creds)
 toolkit = GmailToolkit(api_resource=api_resource)

--- a/toolkits/slack/arcade_slack/tools/chat.py
+++ b/toolkits/slack/arcade_slack/tools/chat.py
@ -16,8 +16,6 @@ from arcade.sdk.auth import SlackUser
            "im:write",
            "users.profile:read",
            "users:read",
-            "channels:read",
-            "groups:read",
        ],
    )
 )
@ -76,9 +74,6 @@ def format_users(userListResponse: dict) -> str:
    requires_auth=SlackUser(
        scope=[
            "chat:write",
-            "im:write",
-            "users.profile:read",
-            "users:read",
            "channels:read",
            "groups:read",
        ],