GitHub toolkit (#16)

the changes needed in the SDK to handle tool auth,
 and multiple tool auth providers.
This commit is contained in:
Nate Barbettini 2024-08-21 19:22:46 -07:00 committed by GitHub
parent 6a2f37edea
commit 3154298572
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 220 additions and 27 deletions

View file

@ -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))

View file

@ -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"))

View file

@ -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),

View file

@ -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):

View file

@ -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"

View file

@ -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,

View file

@ -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",
)

View file

@ -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"

View file

@ -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"]

View file

@ -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
}
]

View file

@ -0,0 +1 @@
__all__ = ["public_repo", "user"]

View file

@ -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}"
)

View file

@ -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}"
)

View file

@ -0,0 +1,17 @@
[tool.poetry]
name = "arcade_github"
version = "0.1.0"
description = "AI tools for interacting with GitHub"
authors = ["Nate Barbettini <nate@arcade-ai.com>"]
[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"

View file

@ -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"],
)
)