diff --git a/llama-index-integrations/tools/llama-index-tools-google/Makefile b/llama-index-integrations/tools/llama-index-tools-google/Makefile index b9eab05aa3..653075c04e 100644 --- a/llama-index-integrations/tools/llama-index-tools-google/Makefile +++ b/llama-index-integrations/tools/llama-index-tools-google/Makefile @@ -3,6 +3,12 @@ GIT_ROOT ?= $(shell git rev-parse --show-toplevel) help: ## Show all Makefile targets. @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}' + +pre-commit: sync lint format test + +sync: + uv sync + format: ## Run code autoformatters (black). pre-commit install git ls-files | xargs pre-commit run black --files diff --git a/llama-index-integrations/tools/llama-index-tools-google/README.md b/llama-index-integrations/tools/llama-index-tools-google/README.md index a721848f11..4796fe97c6 100644 --- a/llama-index-integrations/tools/llama-index-tools-google/README.md +++ b/llama-index-integrations/tools/llama-index-tools-google/README.md @@ -30,6 +30,20 @@ google_spec = GoogleSearchToolSpec(key="your-key", engine="your-engine") - store oAuth`credentials.json` in the same directory as the runnable agent. - you will need to manually approve the Oath every time this tool is invoked +**Multi-calendar support:** + +```python +from llama_index.tools.google import GoogleCalendarToolSpec, all_calendars + +# List all available calendars +calendar_ids = all_calendars(creds) + +# Create tool spec with specific allowed calendars +calendar_spec = GoogleCalendarToolSpec( + creds=creds, allowed_calendar_ids=calendar_ids +) +``` + #### [gmail read, create]() - same as calendar diff --git a/llama-index-integrations/tools/llama-index-tools-google/llama_index/tools/google/__init__.py b/llama-index-integrations/tools/llama-index-tools-google/llama_index/tools/google/__init__.py index cb7d77af25..ce881ba307 100644 --- a/llama-index-integrations/tools/llama-index-tools-google/llama_index/tools/google/__init__.py +++ b/llama-index-integrations/tools/llama-index-tools-google/llama_index/tools/google/__init__.py @@ -1,4 +1,4 @@ -from llama_index.tools.google.calendar.base import GoogleCalendarToolSpec +from llama_index.tools.google.calendar.base import GoogleCalendarToolSpec, all_calendars from llama_index.tools.google.gmail.base import GmailToolSpec from llama_index.tools.google.search.base import ( QUERY_URL_TMPL, @@ -7,6 +7,7 @@ __all__ = [ "GoogleCalendarToolSpec", + "all_calendars", "GmailToolSpec", "GoogleSearchToolSpec", "QUERY_URL_TMPL", diff --git a/llama-index-integrations/tools/llama-index-tools-google/llama_index/tools/google/calendar/base.py b/llama-index-integrations/tools/llama-index-tools-google/llama_index/tools/google/calendar/base.py index afa34141fb..5e727ffdea 100644 --- a/llama-index-integrations/tools/llama-index-tools-google/llama_index/tools/google/calendar/base.py +++ b/llama-index-integrations/tools/llama-index-tools-google/llama_index/tools/google/calendar/base.py @@ -35,7 +35,11 @@ class GoogleCalendarToolSpec(BaseToolSpec): spec_functions = ["load_data", "create_event", "get_date"] - def __init__(self, creds: Optional[Any] = None): + def __init__( + self, + creds: Optional[Any] = None, + allowed_calendar_ids: Optional[List[str]] = None, + ): """ Initialize the GoogleCalendarToolSpec. @@ -45,11 +49,13 @@ def __init__(self, creds: Optional[Any] = None): """ self.creds = creds + self.allowed_calendar_ids = allowed_calendar_ids or [] def load_data( self, number_of_results: Optional[int] = 100, start_date: Optional[Union[str, datetime.date]] = None, + calendar_id: Optional[str] = None, ) -> List[Document]: """ Load data from user's calendar. @@ -57,8 +63,16 @@ def load_data( Args: number_of_results (Optional[int]): the number of events to return. Defaults to 100. start_date (Optional[Union[str, datetime.date]]): the start date to return events from in date isoformat. Defaults to today. + calendar_id (Optional[str]): the calendar ID to load events from. Must be provided and in allowed_calendar_ids list. """ + if calendar_id is None: + return {"error": "calendar_id is required"} + + validation_error = self._validate_calendar_id(calendar_id) + if validation_error: + return validation_error + from googleapiclient.discovery import build credentials = self._get_credentials() @@ -75,7 +89,7 @@ def load_data( events_result = ( service.events() .list( - calendarId="primary", + calendarId=calendar_id, timeMin=start_datetime_utc, maxResults=number_of_results, singleEvents=True, @@ -164,6 +178,7 @@ def create_event( start_datetime: Optional[Union[str, datetime.datetime]] = None, end_datetime: Optional[Union[str, datetime.datetime]] = None, attendees: Optional[List[str]] = None, + calendar_id: Optional[str] = None, ) -> str: """ Create an event on the users calendar. @@ -175,8 +190,16 @@ def create_event( start_datetime Optional[Union[str, datetime.datetime]]: The start datetime for the event end_datetime Optional[Union[str, datetime.datetime]]: The end datetime for the event attendees Optional[List[str]]: A list of email address to invite to the event + calendar_id (Optional[str]): The calendar ID to create the event in. Must be provided and in allowed_calendar_ids list. """ + if calendar_id is None: + return "Error: calendar_id is required" + + validation_error = self._validate_calendar_id(calendar_id) + if validation_error: + return validation_error + from googleapiclient.discovery import build credentials = self._get_credentials() @@ -209,14 +232,37 @@ def create_event( }, "attendees": attendees_list, } - event = service.events().insert(calendarId="primary", body=event).execute() + event = service.events().insert(calendarId=calendar_id, body=event).execute() return ( "Your calendar event has been created successfully! You can move on to the" " next step." ) + def _validate_calendar_id(self, calendar_id: str) -> dict: + if calendar_id not in self.allowed_calendar_ids: + import logging + + logger = logging.getLogger(__name__) + logger.error( + f"Invalid calendar ID '{calendar_id}' attempted. Valid IDs: {self.allowed_calendar_ids}" + ) + return { + "error": "Invalid calendar_id", + "allowed_values": list(self.allowed_calendar_ids), + } + return None + def get_date(self): """ A function to return todays date. Call this before any other functions if you are unaware of the date. """ return datetime.date.today() + + +def all_calendars(creds) -> List[str]: + """List all accessible calendar IDs for configuration purposes.""" + from googleapiclient.discovery import build + + service = build("calendar", "v3", credentials=creds) + calendar_list = service.calendarList().list().execute() + return [cal["id"] for cal in calendar_list.get("items", [])] diff --git a/llama-index-integrations/tools/llama-index-tools-google/tests/test_tools_google.py b/llama-index-integrations/tools/llama-index-tools-google/tests/test_tools_google.py index 0662e60bee..02cdf3f9a1 100644 --- a/llama-index-integrations/tools/llama-index-tools-google/tests/test_tools_google.py +++ b/llama-index-integrations/tools/llama-index-tools-google/tests/test_tools_google.py @@ -58,3 +58,29 @@ def test_google_calendar_tool_spec_get_credentials_oauth_flow( mock_from_file.assert_called_once_with( "token.json", ["https://www.googleapis.com/auth/calendar"] ) + + +def test_google_calendar_tool_spec_init_with_allowed_calendar_ids(): + """Test GoogleCalendarToolSpec initialization with allowed calendar IDs.""" + allowed_ids = ["test@example.com", "secondary@example.com"] + tool = GoogleCalendarToolSpec(allowed_calendar_ids=allowed_ids) + assert tool.allowed_calendar_ids == allowed_ids + + +def test_calendar_id_validation_error(): + """Test calendar ID validation returns structured error.""" + allowed_ids = ["test@example.com"] + tool = GoogleCalendarToolSpec(allowed_calendar_ids=allowed_ids) + + result = tool._validate_calendar_id("invalid@example.com") + expected = {"error": "Invalid calendar_id", "allowed_values": ["test@example.com"]} + assert result == expected + + +def test_calendar_id_validation_success(): + """Test calendar ID validation returns None for valid ID.""" + allowed_ids = ["test@example.com", "secondary@example.com"] + tool = GoogleCalendarToolSpec(allowed_calendar_ids=allowed_ids) + + result = tool._validate_calendar_id("test@example.com") + assert result is None