feat: Add optional creds argument to GoogleCalendarToolSpec (#19826)

Author: MaitreyaBuddha

llama-index-integrations/tools/llama-index-tools-google/llama_index/tools/google/calendar/base.py

@@ -35,6 +35,17 @@ class GoogleCalendarToolSpec(BaseToolSpec):
 
     spec_functions = ["load_data", "create_event", "get_date"]
 
+    def __init__(self, creds: Optional[Any] = None):
+        """
+        Initialize the GoogleCalendarToolSpec.
+
+        Args:
+            creds (Optional[Any]): Pre-configured credentials to use for authentication.
+                                 If provided, these will be used instead of the OAuth flow.
+
+        """
+        self.creds = creds
+
     def load_data(
         self,
         number_of_results: Optional[int] = 100,
@@ -119,6 +130,9 @@ def _get_credentials(self) -> Any:
             Credentials, the obtained credential.
 
         """
+        if self.creds is not None:
+            return self.creds
+
         from google_auth_oauthlib.flow import InstalledAppFlow
 
         from google.auth.transport.requests import Request

llama-index-integrations/tools/llama-index-tools-google/pyproject.toml

@@ -26,7 +26,7 @@ dev = [
 
 [project]
 name = "llama-index-tools-google"
-version = "0.6.1"
+version = "0.6.2"
 description = "llama-index tools google integration"
 authors = [{name = "Your Name", email = "[email protected]"}]
 requires-python = ">=3.9,<4.0"

llama-index-integrations/tools/llama-index-tools-google/tests/test_tools_google.py

@@ -1,3 +1,4 @@
+from unittest.mock import Mock, patch
 from llama_index.core.tools.tool_spec.base import BaseToolSpec
 from llama_index.tools.google import (
     GmailToolSpec,
@@ -15,3 +16,45 @@ def test_class():
 
     names_of_base_classes = [b.__name__ for b in GoogleSearchToolSpec.__mro__]
     assert BaseToolSpec.__name__ in names_of_base_classes
+
+
+def test_google_calendar_tool_spec_init_without_creds():
+    """Test GoogleCalendarToolSpec initialization without credentials."""
+    tool = GoogleCalendarToolSpec()
+    assert tool.creds is None
+
+
+def test_google_calendar_tool_spec_init_with_creds():
+    """Test GoogleCalendarToolSpec initialization with credentials."""
+    mock_creds = Mock()
+    tool = GoogleCalendarToolSpec(creds=mock_creds)
+    assert tool.creds is mock_creds
+
+
+def test_google_calendar_tool_spec_get_credentials_with_provided_creds():
+    """Test _get_credentials returns provided credentials when available."""
+    mock_creds = Mock()
+    tool = GoogleCalendarToolSpec(creds=mock_creds)
+
+    credentials = tool._get_credentials()
+    assert credentials is mock_creds
+
+
+@patch("os.path.exists")
+@patch("google.oauth2.credentials.Credentials.from_authorized_user_file")
+def test_google_calendar_tool_spec_get_credentials_oauth_flow(
+    mock_from_file, mock_exists
+):
+    """Test _get_credentials falls back to OAuth flow when no creds provided."""
+    mock_exists.return_value = True
+    mock_creds = Mock()
+    mock_creds.valid = True
+    mock_from_file.return_value = mock_creds
+
+    tool = GoogleCalendarToolSpec()  # No creds provided
+
+    credentials = tool._get_credentials()
+    assert credentials is mock_creds
+    mock_from_file.assert_called_once_with(
+        "token.json", ["https://www.googleapis.com/auth/calendar"]
+    )