From 3154298572fad9be37162d0127387d39ff8df702 Mon Sep 17 00:00:00 2001 From: Nate Barbettini Date: Wed, 21 Aug 2024 19:22:46 -0700 Subject: [PATCH] GitHub toolkit (#16) the changes needed in the SDK to handle tool auth, and multiple tool auth providers. --- arcade/arcade/actor/core/base.py | 6 ++- arcade/arcade/cli/new.py | 8 +++- arcade/arcade/core/catalog.py | 14 +++++-- arcade/arcade/core/schema.py | 14 +++++++ arcade/arcade/sdk/auth.py | 27 ++++++++++++- .../tests/tool/test_create_tool_definition.py | 28 ++++++++++++- .../fastapi/arcade_example_fastapi/main.py | 13 ++++++- examples/fastapi/poetry.lock | 39 +++++++++++++------ examples/fastapi/pyproject.toml | 1 + schemas/preview/tool_definition.schema.jsonc | 8 +++- toolkits/github/arcade_github/__init__.py | 0 .../github/arcade_github/tools/__init__.py | 1 + .../github/arcade_github/tools/public_repo.py | 31 +++++++++++++++ toolkits/github/arcade_github/tools/user.py | 35 +++++++++++++++++ toolkits/github/pyproject.toml | 17 ++++++++ toolkits/gmail/arcade_gmail/tools/gmail.py | 5 +-- 16 files changed, 220 insertions(+), 27 deletions(-) create mode 100644 toolkits/github/arcade_github/__init__.py create mode 100644 toolkits/github/arcade_github/tools/__init__.py create mode 100644 toolkits/github/arcade_github/tools/public_repo.py create mode 100644 toolkits/github/arcade_github/tools/user.py create mode 100644 toolkits/github/pyproject.toml diff --git a/arcade/arcade/actor/core/base.py b/arcade/arcade/actor/core/base.py index 8b570c50..ce97841a 100644 --- a/arcade/arcade/actor/core/base.py +++ b/arcade/arcade/actor/core/base.py @@ -80,7 +80,11 @@ class BaseActor(Actor): **tool_request.inputs or {}, ) if response.code == 200 and response.data is not None: - output = ToolCallOutput(value=response.data.result) + output = ( + ToolCallOutput(value=response.data.result) + if hasattr(response.data, "result") and response.data.result + else ToolCallOutput(value=f"Tool {tool_name} called successfully") + ) else: output = ToolCallOutput(error=ToolCallError(message=response.msg)) diff --git a/arcade/arcade/cli/new.py b/arcade/arcade/cli/new.py index ef7ee5b0..2f4bb0ba 100644 --- a/arcade/arcade/cli/new.py +++ b/arcade/arcade/cli/new.py @@ -23,7 +23,7 @@ def ask_question(question: str, default: Optional[str] = None) -> str: """ if default: question = f"{question} [{default}]" - answer = typer.prompt(question) + answer = typer.prompt(question, default=default) if not answer and default: return default return str(answer) @@ -79,7 +79,7 @@ build-backend = "poetry.core.masonry.api" def create_new_toolkit(directory: str) -> None: """Generate a new Toolkit package based on user input.""" name = ask_question("Name of the new toolkit?") - toolkit_name = f"arcade_{name}" + toolkit_name = name if name.startswith("arcade_") else f"arcade_{name}" # Check for illegal characters in the toolkit name if not re.match(r"^[\w_]+$", toolkit_name): @@ -103,9 +103,13 @@ def create_new_toolkit(directory: str) -> None: # Create the top level toolkit directory create_directory(top_level_dir) + # Create the toolkit directory create_directory(toolkit_dir) + # Create the __init__.py file in the toolkit directory + create_file(os.path.join(toolkit_dir, "__init__.py"), "") + # Create the tools directory create_directory(os.path.join(toolkit_dir, "tools")) diff --git a/arcade/arcade/core/catalog.py b/arcade/arcade/core/catalog.py index 7a6f3795..48534c92 100644 --- a/arcade/arcade/core/catalog.py +++ b/arcade/arcade/core/catalog.py @@ -23,6 +23,7 @@ from pydantic_core import PydanticUndefined from arcade.core.errors import ToolDefinitionError from arcade.core.schema import ( + GoogleRequirement, InputParameter, OAuth2Requirement, ToolAuthRequirement, @@ -41,7 +42,7 @@ from arcade.core.utils import ( snake_to_pascal_case, ) from arcade.sdk.annotations import Inferrable -from arcade.sdk.auth import OAuth2 +from arcade.sdk.auth import Google, OAuth2, ToolAuthorization WireType = Literal["string", "integer", "float", "boolean", "json"] @@ -181,10 +182,15 @@ class ToolCatalog(BaseModel): raise ToolDefinitionError(f"Tool {tool_name} must have a return type annotation") auth_requirement = getattr(tool, "__tool_requires_auth__", None) - if isinstance(auth_requirement, OAuth2): - auth_requirement = ToolAuthRequirement( - oauth2=OAuth2Requirement(**auth_requirement.model_dump()) + if isinstance(auth_requirement, ToolAuthorization): + new_auth_requirement = ToolAuthRequirement( + provider=auth_requirement.get_provider(), ) + if isinstance(auth_requirement, OAuth2): + new_auth_requirement.oauth2 = OAuth2Requirement(**auth_requirement.model_dump()) + elif isinstance(auth_requirement, Google): + new_auth_requirement.google = GoogleRequirement(**auth_requirement.model_dump()) + auth_requirement = new_auth_requirement return ToolDefinition( name=snake_to_pascal_case(tool_name), diff --git a/arcade/arcade/core/schema.py b/arcade/arcade/core/schema.py index 5c0c3ab7..2efdd356 100644 --- a/arcade/arcade/core/schema.py +++ b/arcade/arcade/core/schema.py @@ -72,10 +72,24 @@ class OAuth2Requirement(BaseModel): """The scope(s) needed for authorization.""" +class GoogleRequirement(BaseModel): + """Indicates that the tool requires Google authorization.""" + + scope: Optional[list[str]] = None + """The scope(s) needed for authorization.""" + + class ToolAuthRequirement(BaseModel): """A requirement for authorization to use a tool.""" + provider: str + """The provider type.""" + oauth2: Optional[OAuth2Requirement] = None + """The OAuth 2.0 requirement, if any.""" + + google: Optional[GoogleRequirement] = None + """The Google requirement, if any.""" class ToolRequirements(BaseModel): diff --git a/arcade/arcade/sdk/auth.py b/arcade/arcade/sdk/auth.py index 51899057..c2765162 100644 --- a/arcade/arcade/sdk/auth.py +++ b/arcade/arcade/sdk/auth.py @@ -1,4 +1,4 @@ -from abc import ABC +from abc import ABC, abstractmethod from typing import Optional from pydantic import AnyUrl, BaseModel @@ -7,14 +7,39 @@ from pydantic import AnyUrl, BaseModel class ToolAuthorization(BaseModel, ABC): """Marks a tool as requiring authorization.""" + @abstractmethod + def get_provider(self) -> str: + """Return the name of the authorization method.""" + pass + pass class OAuth2(ToolAuthorization): """Marks a tool as requiring OAuth 2.0 authorization.""" + def get_provider(self) -> str: + return "oauth2" + authority: AnyUrl """The URL of the OAuth 2.0 authorization server.""" scope: Optional[list[str]] = None """The scope(s) needed for the authorized action.""" + + +class Google(ToolAuthorization): + """Marks a tool as requiring Google authorization.""" + + def get_provider(self) -> str: + return "google" + + scope: Optional[list[str]] = None + """The scope(s) needed for the authorized action.""" + + +class GitHubApp(ToolAuthorization): + """Marks a tool as requiring GitHub App authorization.""" + + def get_provider(self) -> str: + return "github_app" diff --git a/arcade/tests/tool/test_create_tool_definition.py b/arcade/tests/tool/test_create_tool_definition.py index 8d44d650..033d0f81 100644 --- a/arcade/tests/tool/test_create_tool_definition.py +++ b/arcade/tests/tool/test_create_tool_definition.py @@ -4,6 +4,7 @@ import pytest from arcade.core.catalog import ToolCatalog from arcade.core.schema import ( + GoogleRequirement, InputParameter, OAuth2Requirement, ToolAuthRequirement, @@ -15,7 +16,7 @@ from arcade.core.schema import ( ) from arcade.sdk import tool from arcade.sdk.annotations import Inferrable -from arcade.sdk.auth import OAuth2 +from arcade.sdk.auth import Google, OAuth2 ### Tests on @tool decorator @@ -52,6 +53,14 @@ def func_with_auth_requirement(): pass +@tool( + desc="A function that requires authentication", + requires_auth=Google(scope=["https://www.googleapis.com/auth/gmail.readonly"]), +) +def func_with_google_auth_requirement(): + pass + + ### Tests on input params @tool(desc="A function with an input parameter") def func_with_param(context: Annotated[str, "First param"]): @@ -199,15 +208,30 @@ def func_with_complex_return() -> list[dict[str, str]]: { "requirements": ToolRequirements( authorization=ToolAuthRequirement( + provider="oauth2", oauth2=OAuth2Requirement( authority="https://example.com/oauth2/auth", scope=["scope1", "scope2"], - ) + ), ) ) }, id="func_with_auth_requirement", ), + pytest.param( + func_with_google_auth_requirement, + { + "requirements": ToolRequirements( + authorization=ToolAuthRequirement( + provider="google", + google=GoogleRequirement( + scope=["https://www.googleapis.com/auth/gmail.readonly"], + ), + ) + ) + }, + id="func_with_google_auth_requirement", + ), # Tests on input params pytest.param( func_with_value_return, diff --git a/examples/fastapi/arcade_example_fastapi/main.py b/examples/fastapi/arcade_example_fastapi/main.py index 08f4e96c..7121ae64 100644 --- a/examples/fastapi/arcade_example_fastapi/main.py +++ b/examples/fastapi/arcade_example_fastapi/main.py @@ -4,6 +4,7 @@ from pydantic import BaseModel from arcade_arithmetic.tools import arithmetic from arcade_gmail.tools import gmail +from arcade_github.tools import public_repo, user from arcade.actor.fastapi.actor import FastAPIActor @@ -17,6 +18,8 @@ actor.register_tool(arithmetic.multiply) actor.register_tool(arithmetic.divide) actor.register_tool(arithmetic.sqrt) actor.register_tool(gmail.get_emails) +actor.register_tool(public_repo.count_stargazers) +actor.register_tool(user.set_starred) class ChatRequest(BaseModel): @@ -34,7 +37,15 @@ async def chat(request: ChatRequest, tool_choice: str = "execute"): model="gpt-4o-mini", max_tokens=150, # TODO tests for tool choice - tools=["Add", "Multiply", "Divide", "Sqrt", "GetEmails"], + tools=[ + "Add", + "Multiply", + "Divide", + "Sqrt", + "GetEmails", + "CountStargazers", + "SetStarred", + ], tool_choice=tool_choice, user="sam", ) diff --git a/examples/fastapi/poetry.lock b/examples/fastapi/poetry.lock index 1a898e27..91fbca01 100644 --- a/examples/fastapi/poetry.lock +++ b/examples/fastapi/poetry.lock @@ -77,6 +77,23 @@ arcade-ai = {path = "../../arcade", develop = true} type = "directory" url = "../../toolkits/math" +[[package]] +name = "arcade-github" +version = "0.1.0" +description = "AI tools for interacting with GitHub" +optional = false +python-versions = "^3.10" +files = [] +develop = true + +[package.dependencies] +arcade-ai = "^0.1.0" +requests = "^2.32.3" + +[package.source] +type = "directory" +url = "../../toolkits/github" + [[package]] name = "arcade-gmail" version = "0.1.0" @@ -272,18 +289,18 @@ test = ["pytest (>=6)"] [[package]] name = "fastapi" -version = "0.112.0" +version = "0.112.1" description = "FastAPI framework, high performance, easy to learn, fast to code, ready for production" optional = false python-versions = ">=3.8" files = [ - {file = "fastapi-0.112.0-py3-none-any.whl", hash = "sha256:3487ded9778006a45834b8c816ec4a48d522e2631ca9e75ec5a774f1b052f821"}, - {file = "fastapi-0.112.0.tar.gz", hash = "sha256:d262bc56b7d101d1f4e8fc0ad2ac75bb9935fec504d2b7117686cec50710cf05"}, + {file = "fastapi-0.112.1-py3-none-any.whl", hash = "sha256:bcbd45817fc2a1cd5da09af66815b84ec0d3d634eb173d1ab468ae3103e183e4"}, + {file = "fastapi-0.112.1.tar.gz", hash = "sha256:b2537146f8c23389a7faa8b03d0bd38d4986e6983874557d95eed2acc46448ef"}, ] [package.dependencies] pydantic = ">=1.7.4,<1.8 || >1.8,<1.8.1 || >1.8.1,<2.0.0 || >2.0.0,<2.0.1 || >2.0.1,<2.1.0 || >2.1.0,<3.0.0" -starlette = ">=0.37.2,<0.38.0" +starlette = ">=0.37.2,<0.39.0" typing-extensions = ">=4.8.0" [package.extras] @@ -608,13 +625,13 @@ signedtoken = ["cryptography (>=3.0.0)", "pyjwt (>=2.0.0,<3)"] [[package]] name = "openai" -version = "1.40.6" +version = "1.40.8" description = "The official Python library for the openai API" optional = false python-versions = ">=3.7.1" files = [ - {file = "openai-1.40.6-py3-none-any.whl", hash = "sha256:b36372124a779381a420a34dd96f762baa748b6bdfaf83a6b9f2745f72ccc1c5"}, - {file = "openai-1.40.6.tar.gz", hash = "sha256:2239232bcb7f4bd4ce8e02544b5769618582411cf399816d96686d1b6c1e5c8d"}, + {file = "openai-1.40.8-py3-none-any.whl", hash = "sha256:3ed4ddad48e0dde059c9b4d3dc240e47781beca2811e52ba449ddc4a471a2fd4"}, + {file = "openai-1.40.8.tar.gz", hash = "sha256:e225f830b946378e214c5b2cfa8df28ba2aeb7e9d44f738cb2a926fd971f5bc0"}, ] [package.dependencies] @@ -978,13 +995,13 @@ files = [ [[package]] name = "starlette" -version = "0.37.2" +version = "0.38.2" description = "The little ASGI library that shines." optional = false python-versions = ">=3.8" files = [ - {file = "starlette-0.37.2-py3-none-any.whl", hash = "sha256:6fe59f29268538e5d0d182f2791a479a0c64638e6935d1c6989e63fb2699c6ee"}, - {file = "starlette-0.37.2.tar.gz", hash = "sha256:9af890290133b79fc3db55474ade20f6220a364a0402e0b556e7cd5e1e093823"}, + {file = "starlette-0.38.2-py3-none-any.whl", hash = "sha256:4ec6a59df6bbafdab5f567754481657f7ed90dc9d69b0c9ff017907dd54faeff"}, + {file = "starlette-0.38.2.tar.gz", hash = "sha256:c7c0441065252160993a1a37cf2a73bb64d271b17303e0b0c1eb7191cfb12d75"}, ] [package.dependencies] @@ -1098,4 +1115,4 @@ zstd = ["zstandard (>=0.18.0)"] [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "2b4d72bc0026e52378d00b4d6ff6a6d0b95efc34cfd1acf2496fd64827baa378" +content-hash = "ba5d3be15ef9c2adf5c108b4f8ff3388f3a4df546fa84f0bcfeb56c2768171ad" diff --git a/examples/fastapi/pyproject.toml b/examples/fastapi/pyproject.toml index 22c09d10..f7f1cdd8 100644 --- a/examples/fastapi/pyproject.toml +++ b/examples/fastapi/pyproject.toml @@ -10,6 +10,7 @@ fastapi = "^0.112.0" arcade-ai = {path = "../../arcade", develop = true} arcade_arithmetic = {path = "../../toolkits/math", develop = true} arcade_gmail = {path = "../../toolkits/gmail", develop = true} +arcade_github = {path = "../../toolkits/github", develop = true} [build-system] requires = ["poetry-core"] diff --git a/schemas/preview/tool_definition.schema.jsonc b/schemas/preview/tool_definition.schema.jsonc index 8d2bc6d9..bfd4e365 100644 --- a/schemas/preview/tool_definition.schema.jsonc +++ b/schemas/preview/tool_definition.schema.jsonc @@ -53,7 +53,7 @@ "properties": { "parameters": { "type": "array", - "minItems": 1, + "minItems": 0, "items": { "type": "object", "properties": { @@ -121,6 +121,10 @@ { "type": "object", "properties": { + "provider": { + "type": "string", + "enum": ["oauth2", "github_app"] + }, "oauth2": { "type": "object", "properties": { @@ -139,7 +143,7 @@ "additionalProperties": false } }, - "required": ["oauth2"], + "required": [], "additionalProperties": false } ] diff --git a/toolkits/github/arcade_github/__init__.py b/toolkits/github/arcade_github/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/toolkits/github/arcade_github/tools/__init__.py b/toolkits/github/arcade_github/tools/__init__.py new file mode 100644 index 00000000..d6d496f1 --- /dev/null +++ b/toolkits/github/arcade_github/tools/__init__.py @@ -0,0 +1 @@ +__all__ = ["public_repo", "user"] diff --git a/toolkits/github/arcade_github/tools/public_repo.py b/toolkits/github/arcade_github/tools/public_repo.py new file mode 100644 index 00000000..5850e04c --- /dev/null +++ b/toolkits/github/arcade_github/tools/public_repo.py @@ -0,0 +1,31 @@ +from typing import Annotated +from arcade.sdk import tool +import requests + + +# TODO: This does not support private repositories. https://app.clickup.com/t/86b1r3mhe +@tool +def count_stargazers( + owner: Annotated[str, "The owner of the repository"], + name: Annotated[str, "The name of the repository"], +) -> int: + """Count the number of stargazers (stars) for a public GitHub repository. + + For example, to count the number of stars for microsoft/vscode, you would use: + ``` + count_stargazers(owner="microsoft", name="vscode") + ``` + """ + + url = f"https://api.github.com/repos/{owner}/{name}" + response = requests.get(url) + + print(response) + + if response.status_code == 200: + data = response.json() + return data.get("stargazers_count", 0) + else: + raise Exception( + f"Failed to fetch repository data. Status code: {response.status_code}" + ) diff --git a/toolkits/github/arcade_github/tools/user.py b/toolkits/github/arcade_github/tools/user.py new file mode 100644 index 00000000..bea151a4 --- /dev/null +++ b/toolkits/github/arcade_github/tools/user.py @@ -0,0 +1,35 @@ +from typing import Annotated +from arcade.core.schema import ToolContext +from arcade.sdk import tool +from arcade.sdk.auth import GitHubApp +import requests + + +@tool(requires_auth=GitHubApp()) +def set_starred( + context: ToolContext, + owner: Annotated[str, "The owner of the repository"], + name: Annotated[str, "The name of the repository"], + starred: Annotated[bool, "Whether to star the repository or not"], +): + """ + Star or un-star a GitHub repository. + + For example, to star microsoft/vscode, you would use: + ``` + set_starred(owner="microsoft", name="vscode", starred=True) + ``` + """ + + url = f"https://api.github.com/user/starred/{owner}/{name}" + authorization_header = f"Bearer {context.authorization.token}" + response = ( + requests.put(url, headers={"Authorization": authorization_header}) + if starred + else requests.delete(url, headers={"Authorization": authorization_header}) + ) + + if not 200 <= response.status_code < 300: + raise Exception( + f"Failed to star/unstar repository. Status code: {response.status_code}" + ) diff --git a/toolkits/github/pyproject.toml b/toolkits/github/pyproject.toml new file mode 100644 index 00000000..74b7adf8 --- /dev/null +++ b/toolkits/github/pyproject.toml @@ -0,0 +1,17 @@ +[tool.poetry] +name = "arcade_github" +version = "0.1.0" +description = "AI tools for interacting with GitHub" +authors = ["Nate Barbettini "] + +[tool.poetry.dependencies] +python = "^3.10" +arcade-ai = "^0.1.0" +requests = "^2.32.3" + +[tool.poetry.dev-dependencies] +pytest = "^7.4.0" + +[build-system] +requires = ["poetry-core>=1.0.0"] +build-backend = "poetry.core.masonry.api" diff --git a/toolkits/gmail/arcade_gmail/tools/gmail.py b/toolkits/gmail/arcade_gmail/tools/gmail.py index a52697f7..118d57bf 100644 --- a/toolkits/gmail/arcade_gmail/tools/gmail.py +++ b/toolkits/gmail/arcade_gmail/tools/gmail.py @@ -8,12 +8,11 @@ from googleapiclient.discovery import build from arcade.core.schema import ToolContext from arcade.sdk import tool -from arcade.sdk.auth import OAuth2 +from arcade.sdk.auth import Google @tool( - requires_auth=OAuth2( - authority="https://accounts.google.com", + requires_auth=Google( scope=["https://www.googleapis.com/auth/gmail.readonly"], ) )