Toolkit lint cleanup (#72)

Included toolkits as part of the linting process.

Cleaned up any tools that needed to be updated because of this.

This portion of the PR description was added via arcade chat!
This commit is contained in:
Eric Gustin 2024-10-01 10:41:38 -07:00 committed by GitHub
parent aabceae135
commit 83cf070c82
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 94 additions and 73 deletions

View file

@ -56,7 +56,6 @@ ignore = [
[lint.per-file-ignores]
"**/tests/*" = ["S101"]
"toolkits/*" = ["A002", "TRY300", "C901", "C416", "S113", "RUF013", "SIM103"] # TODO: Remove everything here
[format]
preview = true

View file

@ -14,6 +14,8 @@ from arcade.sdk import tool
from arcade.sdk.auth import Google
from arcade_google.tools.utils import (
DateRange,
build_query_string,
fetch_messages,
get_draft_url,
get_email_in_trash_url,
get_sent_email_url,
@ -79,7 +81,7 @@ async def send_email(
)
)
async def send_draft_email(
context: ToolContext, id: Annotated[str, "The ID of the draft to send"]
context: ToolContext, email_id: Annotated[str, "The ID of the draft to send"]
) -> Annotated[str, "A confirmation message with the sent email ID and URL"]:
"""
Send a draft email using the Gmail API.
@ -90,7 +92,7 @@ async def send_draft_email(
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Send the draft email
sent_message = service.users().drafts().send(userId="me", body={"id": id}).execute()
sent_message = service.users().drafts().send(userId="me", body={"id": email_id}).execute()
# Construct the URL to the sent email
return f"Draft email with ID {sent_message['id']} sent: {get_sent_email_url(sent_message['id'])}"
@ -162,7 +164,7 @@ async def write_draft_email(
)
async def update_draft_email(
context: ToolContext,
id: Annotated[str, "The ID of the draft email to update."],
draft_email_id: Annotated[str, "The ID of the draft email to update."],
subject: Annotated[str, "The subject of the draft email"],
body: Annotated[str, "The body of the draft email"],
recipient: Annotated[str, "The recipient of the draft email"],
@ -189,10 +191,10 @@ async def update_draft_email(
raw_message = base64.urlsafe_b64encode(message.as_bytes()).decode()
# Update the draft
draft = {"id": id, "message": {"raw": raw_message}}
draft = {"id": draft_email_id, "message": {"raw": raw_message}}
updated_draft_message = (
service.users().drafts().update(userId="me", id=id, body=draft).execute()
service.users().drafts().update(userId="me", id=draft_email_id, body=draft).execute()
)
return f"Draft email with ID {updated_draft_message['id']} updated: {get_draft_url(updated_draft_message['id'])}"
except HttpError as e:
@ -214,7 +216,7 @@ async def update_draft_email(
)
async def delete_draft_email(
context: ToolContext,
id: Annotated[str, "The ID of the draft email to delete"],
draft_email_id: Annotated[str, "The ID of the draft email to delete"],
) -> Annotated[str, "A confirmation message indicating successful deletion"]:
"""
Delete a draft email using the Gmail API.
@ -225,8 +227,7 @@ async def delete_draft_email(
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Delete the draft
service.users().drafts().delete(userId="me", id=id).execute()
return f"Draft email with ID {id} deleted successfully."
service.users().drafts().delete(userId="me", id=draft_email_id).execute()
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{delete_draft_email.__name__}' tool.",
@ -237,6 +238,8 @@ async def delete_draft_email(
f"Unexpected Error encountered during execution of '{delete_draft_email.__name__}' tool.",
str(e),
)
else:
return f"Draft email with ID {draft_email_id} deleted successfully."
# Email Management Tools
@ -246,7 +249,7 @@ async def delete_draft_email(
)
)
async def trash_email(
context: ToolContext, id: Annotated[str, "The ID of the email to trash"]
context: ToolContext, email_id: Annotated[str, "The ID of the email to trash"]
) -> Annotated[str, "A confirmation message with the trashed email ID and URL"]:
"""
Move an email to the trash folder using the Gmail API.
@ -257,9 +260,9 @@ async def trash_email(
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Trash the email
service.users().messages().trash(userId="me", id=id).execute()
service.users().messages().trash(userId="me", id=email_id).execute()
return f"Email with ID {id} trashed successfully: {get_email_in_trash_url(id)}"
return f"Email with ID {email_id} trashed successfully: {get_email_in_trash_url(email_id)}"
except HttpError as e:
raise ToolExecutionError(
f"HttpError during execution of '{trash_email.__name__}' tool.", str(e)
@ -339,53 +342,21 @@ async def list_emails_by_header(
Search for emails by header using the Gmail API.
At least one of the following parametersMUST be provided: sender, recipient, subject, body.
"""
if not any([sender, recipient, subject, body]):
raise ToolInputError(
"At least one of sender, recipient, subject, or body must be provided."
)
# Build the query string
query = []
if sender:
query.append(f"from:{sender}")
if recipient:
query.append(f"to:{recipient}")
if subject:
query.append(f"subject:{subject}")
if body:
query.append(body)
if date_range:
query.append(date_range.to_date_query())
query_string = " ".join(query)
query = build_query_string(sender, recipient, subject, body, date_range)
try:
# Set up the Gmail API client
service = build("gmail", "v1", credentials=Credentials(context.authorization.token))
# Perform the search
response = (
service.users()
.messages()
.list(userId="me", q=query_string, maxResults=limit or 100)
.execute()
)
messages = response.get("messages", [])
messages = fetch_messages(service, query, limit)
if not messages:
return json.dumps({"emails": []})
emails = []
for msg in messages:
try:
email_data = service.users().messages().get(userId="me", id=msg["id"]).execute()
email_details = parse_email(email_data)
if email_details:
emails.append(email_details)
except HttpError as e:
print(f"Error reading email {msg['id']}: {e}")
emails = process_messages(service, messages)
return json.dumps({"emails": emails})
except HttpError as e:
raise ToolExecutionError(
@ -399,6 +370,18 @@ async def list_emails_by_header(
)
def process_messages(service, messages):
emails = []
for msg in messages:
try:
email_data = service.users().messages().get(userId="me", id=msg["id"]).execute()
email_details = parse_email(email_data)
emails += email_details if email_details else []
except HttpError as e:
print(f"Error reading email {msg['id']}: {e}")
return emails
@tool(
requires_auth=Google(
scopes=["https://www.googleapis.com/auth/gmail.readonly"],

View file

@ -174,3 +174,34 @@ def _clean_text(text: str) -> str:
text = "\n".join(line.strip() for line in text.split("\n"))
return text
def build_query_string(sender, recipient, subject, body, date_range):
"""
Helper function to build a query string for Gmail list_emails_by_header tool.
"""
query = []
if sender:
query.append(f"from:{sender}")
if recipient:
query.append(f"to:{recipient}")
if subject:
query.append(f"subject:{subject}")
if body:
query.append(body)
if date_range:
query.append(date_range.to_date_query())
return " ".join(query)
def fetch_messages(service, query_string, limit):
"""
Helper function to fetch messages from Gmail API for the list_emails_by_header tool.
"""
response = (
service.users()
.messages()
.list(userId="me", q=query_string, maxResults=limit or 100)
.execute()
)
return response.get("messages", [])

View file

@ -99,7 +99,7 @@ async def test_update_draft_email(mock_build, mock_context):
# Test happy path
result = await update_draft_email(
context=mock_context,
id="draft123",
draft_email_id="draft123",
subject="Updated Subject",
body="Updated Body",
recipient="updated@example.com",
@ -117,7 +117,7 @@ async def test_update_draft_email(mock_build, mock_context):
with pytest.raises(ToolExecutionError):
await update_draft_email(
context=mock_context,
id="nonexistent_draft",
draft_email_id="nonexistent_draft",
subject="Updated Subject",
body="Updated Body",
recipient="updated@example.com",
@ -131,7 +131,7 @@ async def test_send_draft_email(mock_build, mock_context):
mock_build.return_value = mock_service
# Test happy path
result = await send_draft_email(context=mock_context, id="draft456")
result = await send_draft_email(context=mock_context, email_id="draft456")
assert "Draft email with ID" in result
assert "sent" in result
@ -143,7 +143,7 @@ async def test_send_draft_email(mock_build, mock_context):
)
with pytest.raises(ToolExecutionError):
await send_draft_email(context=mock_context, id="nonexistent_draft")
await send_draft_email(context=mock_context, email_id="nonexistent_draft")
@pytest.mark.asyncio
@ -153,7 +153,7 @@ async def test_delete_draft_email(mock_build, mock_context):
mock_build.return_value = mock_service
# Test happy path
result = await delete_draft_email(context=mock_context, id="draft789")
result = await delete_draft_email(context=mock_context, draft_email_id="draft789")
assert "Draft email with ID" in result
assert "deleted successfully" in result
@ -165,7 +165,7 @@ async def test_delete_draft_email(mock_build, mock_context):
)
with pytest.raises(ToolExecutionError):
await delete_draft_email(context=mock_context, id="nonexistent_draft")
await delete_draft_email(context=mock_context, draft_email_id="nonexistent_draft")
@pytest.mark.asyncio
@ -404,7 +404,7 @@ async def test_trash_email(mock_build, mock_context):
# Test happy path
email_id = "123456"
result = await trash_email(context=mock_context, id=email_id)
result = await trash_email(context=mock_context, email_id=email_id)
assert (
f"Email with ID {email_id} trashed successfully: https://mail.google.com/mail/u/0/#trash/{email_id}"
@ -418,4 +418,4 @@ async def test_trash_email(mock_build, mock_context):
)
with pytest.raises(ToolExecutionError):
await trash_email(context=mock_context, id="nonexistent_email")
await trash_email(context=mock_context, email_id="nonexistent_email")

View file

@ -72,4 +72,4 @@ def sum_range(
"""
Sum all numbers from start through end
"""
return sum([i for i in range(start, end + 1)])
return sum(list(range(start, end + 1)))

View file

@ -1,6 +1,6 @@
from typing import Annotated
import requests
import httpx
from arcade.core.errors import ToolExecutionError
from arcade.core.schema import ToolContext
@ -13,7 +13,7 @@ TWEETS_URL = "https://api.x.com/2/tweets"
# Manage Tweets Tools. See developer docs for additional available parameters: https://developer.x.com/en/docs/x-api/tweets/manage-tweets/api-reference
@tool(requires_auth=X(scopes=["tweet.read", "tweet.write", "users.read"]))
def post_tweet(
async def post_tweet(
context: ToolContext,
tweet_text: Annotated[str, "The text content of the tweet you want to post"],
) -> Annotated[str, "Success string and the URL of the tweet"]:
@ -25,7 +25,8 @@ def post_tweet(
}
payload = {"text": tweet_text}
response = requests.post(TWEETS_URL, headers=headers, json=payload)
async with httpx.AsyncClient() as client:
response = await client.post(TWEETS_URL, headers=headers, json=payload, timeout=10)
if response.status_code != 201:
raise ToolExecutionError(
@ -37,7 +38,7 @@ def post_tweet(
@tool(requires_auth=X(scopes=["tweet.read", "tweet.write", "users.read"]))
def delete_tweet_by_id(
async def delete_tweet_by_id(
context: ToolContext,
tweet_id: Annotated[str, "The ID of the tweet you want to delete"],
) -> Annotated[str, "Success string confirming the tweet deletion"]:
@ -46,7 +47,8 @@ def delete_tweet_by_id(
headers = {"Authorization": f"Bearer {context.authorization.token}"}
url = f"{TWEETS_URL}/{tweet_id}"
response = requests.delete(url, headers=headers)
async with httpx.AsyncClient() as client:
response = await client.delete(url, headers=headers, timeout=10)
if response.status_code != 200:
raise ToolExecutionError(
@ -57,7 +59,7 @@ def delete_tweet_by_id(
@tool(requires_auth=X(scopes=["tweet.read", "users.read"]))
def search_recent_tweets_by_username(
async def search_recent_tweets_by_username(
context: ToolContext,
username: Annotated[str, "The username of the X (Twitter) user to look up"],
max_results: Annotated[
@ -78,7 +80,8 @@ def search_recent_tweets_by_username(
"https://api.x.com/2/tweets/search/recent?expansions=author_id&user.fields=id,name,username"
)
response = requests.get(url, headers=headers, params=params)
async with httpx.AsyncClient() as client:
response = await client.get(url, headers=headers, params=params, timeout=10)
if response.status_code != 200:
raise ToolExecutionError(
@ -91,10 +94,14 @@ def search_recent_tweets_by_username(
@tool(requires_auth=X(scopes=["tweet.read", "users.read"]))
def search_recent_tweets_by_keywords(
async def search_recent_tweets_by_keywords(
context: ToolContext,
keywords: Annotated[list[str], "List of keywords that must be present in the tweet"] = None,
phrases: Annotated[list[str], "List of phrases that must be present in the tweet"] = None,
keywords: Annotated[
list[str] | None, "List of keywords that must be present in the tweet"
] = None,
phrases: Annotated[
list[str] | None, "List of phrases that must be present in the tweet"
] = None,
max_results: Annotated[
int, "The maximum number of results to return. Cannot be less than 10"
] = 10,
@ -125,7 +132,8 @@ def search_recent_tweets_by_keywords(
"https://api.x.com/2/tweets/search/recent?expansions=author_id&user.fields=id,name,username"
)
response = requests.get(url, headers=headers, params=params)
async with httpx.AsyncClient() as client:
response = await client.get(url, headers=headers, params=params, timeout=10)
if response.status_code != 200:
raise ToolExecutionError(

View file

@ -1,6 +1,6 @@
from typing import Annotated
import requests
import httpx
from arcade.core.errors import ToolExecutionError
from arcade.core.schema import ToolContext
@ -10,7 +10,7 @@ from arcade.sdk.auth import X
# Users Lookup Tools. See developer docs for additional available query parameters: https://developer.x.com/en/docs/x-api/users/lookup/api-reference
@tool(requires_auth=X(scopes=["users.read", "tweet.read"]))
def lookup_single_user_by_username(
async def lookup_single_user_by_username(
context: ToolContext,
username: Annotated[str, "The username of the X (Twitter) user to look up"],
) -> Annotated[str, "User information including id, name, username, and description"]:
@ -21,7 +21,8 @@ def lookup_single_user_by_username(
}
url = f"https://api.x.com/2/users/by/username/{username}?user.fields=created_at,description,id,location,most_recent_tweet_id,name,pinned_tweet_id,profile_image_url,protected,public_metrics,url,username,verified,verified_type,withheld"
response = requests.get(url, headers=headers)
async with httpx.AsyncClient() as client:
response = await client.get(url, headers=headers, timeout=10)
if response.status_code != 200:
raise ToolExecutionError(

View file

@ -53,6 +53,4 @@ def sanity_check_tweets_data(tweets_data: dict) -> bool:
"""
if not tweets_data.get("data", []):
return False
if not tweets_data.get("includes", {}).get("users", []):
return False
return True
return tweets_data.get("includes", {}).get("users", [])

View file

@ -7,6 +7,7 @@ authors = ["Eric Gustin <eric@arcade-ai.com>"]
[tool.poetry.dependencies]
python = "^3.10"
arcade-ai = "^0.1.0"
httpx = "^0.27.2"
[tool.poetry.dev-dependencies]
pytest = "^8.3.0"