diff --git a/toolkits/google/arcade_google/tools/constants.py b/toolkits/google/arcade_google/tools/constants.py index 2ef97c93..98390598 100644 --- a/toolkits/google/arcade_google/tools/constants.py +++ b/toolkits/google/arcade_google/tools/constants.py @@ -16,3 +16,6 @@ except ValueError as e: f"'{os.getenv('ARCADE_GMAIL_DEFAULT_REPLY_TO')}'. Expected one of " f"{list(GmailReplyToWhom.__members__.keys())}" ) from e + + +DEFAULT_SEARCH_CONTACTS_LIMIT = 30 diff --git a/toolkits/google/arcade_google/tools/contacts.py b/toolkits/google/arcade_google/tools/contacts.py index 9ff5ab98..ee037124 100644 --- a/toolkits/google/arcade_google/tools/contacts.py +++ b/toolkits/google/arcade_google/tools/contacts.py @@ -6,6 +6,9 @@ from arcade.sdk.auth import Google from google.oauth2.credentials import Credentials from googleapiclient.discovery import build +from arcade_google.tools.constants import DEFAULT_SEARCH_CONTACTS_LIMIT +from arcade_google.tools.utils import build_people_service, search_contacts + async def _warmup_cache(service) -> None: # type: ignore[no-untyped-def] """ @@ -18,47 +21,46 @@ async def _warmup_cache(service) -> None: # type: ignore[no-untyped-def] @tool(requires_auth=Google(scopes=["https://www.googleapis.com/auth/contacts.readonly"])) -async def search_contacts( +async def search_contacts_by_email( context: ToolContext, - query: Annotated[ - str, - "The search query for filtering contacts.", - ], + email: Annotated[str, "The email address to search for"], limit: Annotated[ Optional[int], - "The maximum number of contacts to return (default 10, max 30)", - ] = 10, + "The maximum number of contacts to return (30 is the max allowed by Google API)", + ] = DEFAULT_SEARCH_CONTACTS_LIMIT, ) -> Annotated[dict, "A dictionary containing the list of matching contacts"]: """ - Search the user's contacts in Google Contacts. - - Up to 30 contacts with a name or email address containing the query will be returned. - If the query matches more than 30 contacts, only the first 30 will be returned. + Search the user's contacts in Google Contacts by email address. """ - # Build the People API service - service = build( - "people", - "v1", - credentials=Credentials( - context.authorization.token - if context.authorization and context.authorization.token - else "" - ), + service = build_people_service( + context.authorization.token if context.authorization and context.authorization.token else "" ) - # Warm-up the cache before performing search. # TODO: Ideally we should warmup only if this user (or google domain?) hasn't warmed up recently await _warmup_cache(service) - # Search primary contacts using searchContacts - primary_response = ( - service.people() - .searchContacts(query=query, pageSize=limit, readMask="names,emailAddresses") - .execute() - ) - primary_results = primary_response.get("results", []) + return {"contacts": search_contacts(service, email, limit)} - return {"contacts": primary_results} + +@tool(requires_auth=Google(scopes=["https://www.googleapis.com/auth/contacts.readonly"])) +async def search_contacts_by_name( + context: ToolContext, + name: Annotated[str, "The full name to search for"], + limit: Annotated[ + Optional[int], + "The maximum number of contacts to return (30 is the max allowed by Google API)", + ] = DEFAULT_SEARCH_CONTACTS_LIMIT, +) -> Annotated[dict, "A dictionary containing the list of matching contacts"]: + """ + Search the user's contacts in Google Contacts by name. + """ + service = build_people_service( + context.authorization.token if context.authorization and context.authorization.token else "" + ) + # Warm-up the cache before performing search. + # TODO: Ideally we should warmup only if this user (or google domain?) hasn't warmed up recently + await _warmup_cache(service) + return {"contacts": search_contacts(service, name, limit)} @tool(requires_auth=Google(scopes=["https://www.googleapis.com/auth/contacts"])) diff --git a/toolkits/google/arcade_google/tools/utils.py b/toolkits/google/arcade_google/tools/utils.py index 330590ef..6f0a2c27 100644 --- a/toolkits/google/arcade_google/tools/utils.py +++ b/toolkits/google/arcade_google/tools/utils.py @@ -5,7 +5,7 @@ from datetime import datetime, timedelta from email.message import EmailMessage from email.mime.text import MIMEText from enum import Enum -from typing import Any, Optional, Union +from typing import Any, Optional, Union, cast from zoneinfo import ZoneInfo from arcade.sdk import ToolContext @@ -13,6 +13,7 @@ from bs4 import BeautifulSoup from google.oauth2.credentials import Credentials from googleapiclient.discovery import Resource, build +from arcade_google.tools.constants import DEFAULT_SEARCH_CONTACTS_LIMIT from arcade_google.tools.exceptions import GmailToolError, GoogleServiceError from arcade_google.tools.models import Day, GmailAction, GmailReplyToWhom, TimeSlot @@ -598,3 +599,39 @@ def build_docs_service(auth_token: Optional[str]) -> Resource: # type: ignore[n """ auth_token = auth_token or "" return build("docs", "v1", credentials=Credentials(auth_token)) + + +# Contacts utils +def build_people_service(auth_token: Optional[str]) -> Resource: # type: ignore[no-any-unimported] + """ + Build a People service object. + """ + auth_token = auth_token or "" + return build("people", "v1", credentials=Credentials(auth_token)) + + +def search_contacts(service: Any, query: str, limit: Optional[int]) -> list[dict[str, Any]]: + """ + Search the user's contacts in Google Contacts. + """ + response = ( + service.people() + .searchContacts( + query=query, + pageSize=limit or DEFAULT_SEARCH_CONTACTS_LIMIT, + readMask=",".join([ + "names", + "nicknames", + "emailAddresses", + "phoneNumbers", + "addresses", + "organizations", + "biographies", + "urls", + "userDefined", + ]), + ) + .execute() + ) + + return cast(list[dict[str, Any]], response.get("results", [])) diff --git a/toolkits/google/evals/eval_google_contacts.py b/toolkits/google/evals/eval_google_contacts.py index 653fbbe4..354f45aa 100644 --- a/toolkits/google/evals/eval_google_contacts.py +++ b/toolkits/google/evals/eval_google_contacts.py @@ -8,7 +8,11 @@ from arcade.sdk.eval import ( from arcade.sdk.eval.critic import BinaryCritic import arcade_google -from arcade_google.tools.contacts import create_contact, search_contacts +from arcade_google.tools.contacts import ( + create_contact, + search_contacts_by_email, + search_contacts_by_name, +) # Evaluation rubric rubric = EvalRubric( @@ -31,12 +35,27 @@ def contacts_eval_suite() -> EvalSuite: ) suite.add_case( - name="Find a contact by name", + name="Search contacts by name", user_message="Find my contact Bob", expected_tool_calls=[ ExpectedToolCall( - func=search_contacts, - args={"query": "Bob"}, + func=search_contacts_by_name, + args={ + "name": "Bob", + }, + ) + ], + ) + + suite.add_case( + name="Search contacts by email", + user_message="Find my contact alice@example.com", + expected_tool_calls=[ + ExpectedToolCall( + func=search_contacts_by_email, + args={ + "email": "alice@example.com", + }, ) ], ) @@ -46,9 +65,9 @@ def contacts_eval_suite() -> EvalSuite: user_message="Find 5 contacts whose names include 'Alice'", expected_tool_calls=[ ExpectedToolCall( - func=search_contacts, + func=search_contacts_by_name, args={ - "query": "Alice", + "name": "Alice", "limit": 5, }, ) diff --git a/toolkits/google/tests/test_contacts.py b/toolkits/google/tests/test_contacts.py index c6bb57c6..6f62433b 100644 --- a/toolkits/google/tests/test_contacts.py +++ b/toolkits/google/tests/test_contacts.py @@ -3,7 +3,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from arcade.sdk import ToolContext -from arcade_google.tools.contacts import create_contact, search_contacts +from arcade_google.tools.contacts import create_contact @pytest.fixture @@ -14,67 +14,6 @@ def mock_context(): return context -@pytest.mark.asyncio -async def test_search_contacts_success(mock_context): - search_response_data = { - "results": [ - { - "resourceName": "people/1", - "names": [{"displayName": "John Doe"}], - "emailAddresses": [{"value": "john@example.com"}], - }, - { - "resourceName": "people/2", - "names": [{"displayName": "Jane Doe"}], - "emailAddresses": [{"value": "jane@example.com"}], - }, - ] - } - search_call = MagicMock() - search_call.execute.return_value = search_response_data - - people_mock = MagicMock() - people_mock.searchContacts.return_value = search_call - - service_mock = MagicMock() - service_mock.people.return_value = people_mock - - with ( - patch("arcade_google.tools.contacts.build", return_value=service_mock) as mock_build, - patch( - "arcade_google.tools.contacts._warmup_cache", new=AsyncMock(return_value=None) - ) as mock_warmup, - ): - result = await search_contacts(mock_context, query="Doe", limit=2) - assert "contacts" in result - assert result["contacts"] == search_response_data["results"] - - assert mock_warmup.call_count == 1 - assert people_mock.searchContacts.call_count == 1 - - # Check that the People API service was built with the expected parameters. - mock_build.assert_called_once() - - -@pytest.mark.asyncio -async def test_search_contacts_error(mock_context): - error_call = MagicMock() - error_call.execute.side_effect = Exception("Search error") - - people_mock = MagicMock() - people_mock.searchContacts.return_value = error_call - - service_mock = MagicMock() - service_mock.people.return_value = people_mock - - with ( - patch("arcade_google.tools.contacts.build", return_value=service_mock), - patch("arcade_google.tools.contacts._warmup_cache", new=AsyncMock(return_value=None)), - pytest.raises(Exception, match="Error in execution of SearchContacts"), - ): - await search_contacts(mock_context, query="Doe") - - @pytest.mark.asyncio async def test_create_contact_success(mock_context): # Test create_contact with all parameters (given, family names and email)