diff --git a/arcade/arcade/core/catalog.py b/arcade/arcade/core/catalog.py index 9d50d516..f664e4e7 100644 --- a/arcade/arcade/core/catalog.py +++ b/arcade/arcade/core/catalog.py @@ -3,6 +3,7 @@ import inspect from collections.abc import Iterator from dataclasses import dataclass from datetime import datetime +from enum import Enum from importlib import import_module from types import ModuleType from typing import ( @@ -239,6 +240,9 @@ def create_input_definition(func: Callable) -> ToolInputs: if is_string_literal(tool_field_info.field_type): is_enum = True enum_values = [str(e) for e in get_args(tool_field_info.field_type)] + elif issubclass(tool_field_info.field_type, Enum): + is_enum = True + enum_values = [e.value for e in tool_field_info.field_type] # If the field has a default value, it is not required # If the field is optional, it is not required @@ -479,6 +483,8 @@ def get_wire_type( origin = _type.__origin__ if origin in [list, dict]: return "json" + elif issubclass(_type, Enum): + return "string" elif issubclass(_type, BaseModel): return "json" raise ToolDefinitionError(f"Unsupported parameter type: {_type}") diff --git a/arcade/tests/tool/test_create_tool_definition.py b/arcade/tests/tool/test_create_tool_definition.py index 033d0f81..a6f4a648 100644 --- a/arcade/tests/tool/test_create_tool_definition.py +++ b/arcade/tests/tool/test_create_tool_definition.py @@ -1,3 +1,4 @@ +from enum import Enum from typing import Annotated, Literal, Optional import pytest @@ -88,13 +89,23 @@ def func_with_every_param( pass +class TestEnum(Enum): + FOO_BAR = "foo bar" + BAZ = "baz" + + +@tool(desc="A function that takes an enum") +def func_with_enum_param(param1: Annotated[TestEnum, "an enum"]): + pass + + @tool(desc="A function that takes a dictionary") def func_with_dict_param(param1: Annotated[dict, "a cool dictionary"]): pass -@tool(desc="A function that takes a string enum") -def func_with_string_enum_param(param1: Annotated[Literal["value1", "value2"], "a few choices"]): +@tool(desc="A function that takes a string literal") +def func_with_string_literal_param(param1: Annotated[Literal["value1", "value2"], "a few choices"]): pass @@ -329,6 +340,23 @@ def func_with_complex_return() -> list[dict[str, str]]: }, id="func_with_every_param", ), + pytest.param( + func_with_enum_param, + { + "inputs": ToolInputs( + parameters=[ + InputParameter( + name="param1", + description="an enum", + inferrable=True, + required=True, + value_schema=ValueSchema(val_type="string", enum=["foo bar", "baz"]), + ) + ] + ), + }, + id="func_with_enum_param", + ), pytest.param( func_with_dict_param, { @@ -347,7 +375,7 @@ def func_with_complex_return() -> list[dict[str, str]]: id="func_with_dict_param", ), pytest.param( - func_with_string_enum_param, + func_with_string_literal_param, { "inputs": ToolInputs( parameters=[ diff --git a/examples/fastapi/arcade_example_fastapi/main.py b/examples/fastapi/arcade_example_fastapi/main.py index ee871565..c4fb3ee0 100644 --- a/examples/fastapi/arcade_example_fastapi/main.py +++ b/examples/fastapi/arcade_example_fastapi/main.py @@ -21,6 +21,7 @@ actor = FastAPIActor(app) # actor.register_tool(arithmetic.divide) # actor.register_tool(arithmetic.sqrt) actor.register_tool(gmail.get_emails) +actor.register_tool(gmail.search_emails_by_header) actor.register_tool(gmail.write_draft) actor.register_tool(repo.count_stargazers) actor.register_tool(repo.search_issues) @@ -45,6 +46,7 @@ async def postChat(request: ChatRequest, tool_choice: str = "execute"): max_tokens=500, tools=[ "GetEmails", + "SearchEmailsByHeader", "WriteDraft", "CountStargazers", "SetStarred", diff --git a/toolkits/gmail/arcade_gmail/tools/gmail.py b/toolkits/gmail/arcade_gmail/tools/gmail.py index d29f81d0..5e476a75 100644 --- a/toolkits/gmail/arcade_gmail/tools/gmail.py +++ b/toolkits/gmail/arcade_gmail/tools/gmail.py @@ -1,11 +1,13 @@ import base64 +import datetime +from enum import Enum import json -import re -from base64 import urlsafe_b64decode from email.mime.text import MIMEText -from typing import Annotated, Any, Dict, Optional +from typing import Annotated, Optional +from arcade.core.errors import ToolExecutionError +from googleapiclient.errors import HttpError -from bs4 import BeautifulSoup +from arcade_gmail.tools.utils import parse_email from google.oauth2.credentials import Credentials from googleapiclient.discovery import build @@ -48,6 +50,122 @@ def get_draft_url(draft_id): return f"https://mail.google.com/mail/u/0/#drafts/{draft_id}" +class DateRange(Enum): + TODAY = "today" + YESTERDAY = "yesterday" + LAST_7_DAYS = "last_7_days" + LAST_30_DAYS = "last_30_days" + THIS_MONTH = "this_month" + LAST_MONTH = "last_month" + THIS_YEAR = "this_year" + + def to_date_query(self): + today = datetime.datetime.now() + result = "after:" + comparison_date = today + + if self == DateRange.YESTERDAY: + comparison_date = today - datetime.timedelta(days=1) + elif self == DateRange.LAST_7_DAYS: + comparison_date = today - datetime.timedelta(days=7) + elif self == DateRange.LAST_30_DAYS: + comparison_date = today - datetime.timedelta(days=30) + elif self == DateRange.THIS_MONTH: + comparison_date = today.replace(day=1) + elif self == DateRange.LAST_MONTH: + comparison_date = ( + today.replace(day=1) - datetime.timedelta(days=1) + ).replace(day=1) + elif self == DateRange.THIS_YEAR: + comparison_date = today.replace(month=1, day=1) + elif self == DateRange.LAST_MONTH: + comparison_date = ( + today.replace(month=1, day=1) - datetime.timedelta(days=1) + ).replace(month=1, day=1) + + return result + comparison_date.strftime("%Y/%m/%d") + + +@tool( + requires_auth=Google( + scope=["https://www.googleapis.com/auth/gmail.readonly"], + ) +) +async def search_emails_by_header( + context: ToolContext, + sender: Annotated[ + Optional[str], "The name or email address of the sender of the email" + ] = None, + recipient: Annotated[ + Optional[str], "The name or email address of the recipient" + ] = None, + subject: Annotated[ + Optional[str], "Words to find in the subject of the email" + ] = None, + body: Annotated[Optional[str], "Words to find in the body of the email"] = None, + date_range: Annotated[Optional[DateRange], "The date range of the email"] = None, + limit: Annotated[Optional[int], "The maximum number of emails to return"] = 25, +) -> Annotated[str, "A list of email details in JSON format"]: + """Search for emails by header. + One of the following MUST be provided: sender, recipient, subject, body.""" + + if not any([sender, recipient, subject, body]): + raise ValueError( + "At least one of sender, recipient, subject, or body must be provided." + ) + + # Set up the Gmail API client + service = build("gmail", "v1", credentials=Credentials(context.authorization.token)) + + # 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) + + try: + # Perform the search + response = ( + service.users() + .messages() + .list(userId="me", q=query_string, maxResults=limit or 100) + .execute() + ) + messages = response.get("messages", []) + + 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}") + + return json.dumps({"emails": emails}) + + except HttpError as e: + raise ToolExecutionError( + "Error searching emails", + developer_message=f"Gmail API Error: {e}", + ) + + @tool( requires_auth=Google( scope=["https://www.googleapis.com/auth/gmail.readonly"], @@ -94,96 +212,3 @@ async def get_emails( except Exception as e: print(f"Error reading emails: {e}") return "Error reading emails" - - -def parse_email(email_data: Dict[str, Any]) -> Optional[Dict[str, str]]: - """ - Parse email data and extract relevant information. - - Args: - email_data (Dict[str, Any]): Raw email data from Gmail API. - - Returns: - Optional[Dict[str, str]]: Parsed email details or None if parsing fails. - """ - try: - payload = email_data["payload"] - headers = {d["name"].lower(): d["value"] for d in payload["headers"]} - - body_data = get_email_body(payload) - - return { - "from": headers.get("from", ""), - "date": headers.get("date", ""), - "subject": headers.get("subject", "No subject"), - "body": clean_email_body(body_data) if body_data else "", - } - except Exception as e: - print(f"Error parsing email {email_data.get('id', 'unknown')}: {e}") - return None - - -def get_email_body(payload: Dict[str, Any]) -> Optional[str]: - """ - Extract email body from payload. - - Args: - payload (Dict[str, Any]): Email payload data. - - Returns: - Optional[str]: Decoded email body or None if not found. - """ - if "body" in payload and payload["body"].get("data"): - return urlsafe_b64decode(payload["body"]["data"]).decode() - - for part in payload.get("parts", []): - if part.get("mimeType") == "text/plain" and "data" in part["body"]: - return urlsafe_b64decode(part["body"]["data"]).decode() - - return None - - -def clean_email_body(body: str) -> str: - """ - Remove HTML tags and clean up email body text while preserving most content. - - Args: - body (str): The raw email body text. - - Returns: - str: Cleaned email body text. - """ - try: - # Remove HTML tags using BeautifulSoup - soup = BeautifulSoup(body, "html.parser") - text = soup.get_text(separator=" ") - - # Clean up the text - text = clean_text(text) - - return text.strip() - except Exception as e: - print(f"Error cleaning email body: {e}") - return body - - -def clean_text(text: str) -> str: - """ - Clean up the text while preserving most content. - - Args: - text (str): The input text. - - Returns: - str: Cleaned text. - """ - # Replace multiple newlines with a single newline - text = re.sub(r"\n+", "\n", text) - - # Replace multiple spaces with a single space - text = re.sub(r"\s+", " ", text) - - # Remove leading/trailing whitespace from each line - text = "\n".join(line.strip() for line in text.split("\n")) - - return text diff --git a/toolkits/gmail/arcade_gmail/tools/utils.py b/toolkits/gmail/arcade_gmail/tools/utils.py new file mode 100644 index 00000000..9768b3a9 --- /dev/null +++ b/toolkits/gmail/arcade_gmail/tools/utils.py @@ -0,0 +1,98 @@ +from base64 import urlsafe_b64decode +import re +from typing import Any, Dict, Optional + +from bs4 import BeautifulSoup + + +def parse_email(email_data: Dict[str, Any]) -> Optional[Dict[str, str]]: + """ + Parse email data and extract relevant information. + + Args: + email_data (Dict[str, Any]): Raw email data from Gmail API. + + Returns: + Optional[Dict[str, str]]: Parsed email details or None if parsing fails. + """ + try: + payload = email_data["payload"] + headers = {d["name"].lower(): d["value"] for d in payload["headers"]} + + body_data = _get_email_body(payload) + + return { + "from": headers.get("from", ""), + "date": headers.get("date", ""), + "subject": headers.get("subject", "No subject"), + "body": _clean_email_body(body_data) if body_data else "", + } + except Exception as e: + print(f"Error parsing email {email_data.get('id', 'unknown')}: {e}") + return None + + +def _get_email_body(payload: Dict[str, Any]) -> Optional[str]: + """ + Extract email body from payload. + + Args: + payload (Dict[str, Any]): Email payload data. + + Returns: + Optional[str]: Decoded email body or None if not found. + """ + if "body" in payload and payload["body"].get("data"): + return urlsafe_b64decode(payload["body"]["data"]).decode() + + for part in payload.get("parts", []): + if part.get("mimeType") == "text/plain" and "data" in part["body"]: + return urlsafe_b64decode(part["body"]["data"]).decode() + + return None + + +def _clean_email_body(body: str) -> str: + """ + Remove HTML tags and clean up email body text while preserving most content. + + Args: + body (str): The raw email body text. + + Returns: + str: Cleaned email body text. + """ + try: + # Remove HTML tags using BeautifulSoup + soup = BeautifulSoup(body, "html.parser") + text = soup.get_text(separator=" ") + + # Clean up the text + text = _clean_text(text) + + return text.strip() + except Exception as e: + print(f"Error cleaning email body: {e}") + return body + + +def _clean_text(text: str) -> str: + """ + Clean up the text while preserving most content. + + Args: + text (str): The input text. + + Returns: + str: Cleaned text. + """ + # Replace multiple newlines with a single newline + text = re.sub(r"\n+", "\n", text) + + # Replace multiple spaces with a single space + text = re.sub(r"\s+", " ", text) + + # Remove leading/trailing whitespace from each line + text = "\n".join(line.strip() for line in text.split("\n")) + + return text