datarobot-genai 0.2.19__py3-none-any.whl → 0.2.21__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- datarobot_genai/drmcp/core/config.py +24 -0
- datarobot_genai/drmcp/core/tool_config.py +8 -0
- datarobot_genai/drmcp/tools/clients/confluence.py +99 -0
- datarobot_genai/drmcp/tools/clients/gdrive.py +248 -0
- datarobot_genai/drmcp/tools/confluence/tools.py +66 -0
- datarobot_genai/drmcp/tools/gdrive/__init__.py +0 -0
- datarobot_genai/drmcp/tools/gdrive/tools.py +88 -0
- {datarobot_genai-0.2.19.dist-info → datarobot_genai-0.2.21.dist-info}/METADATA +1 -1
- {datarobot_genai-0.2.19.dist-info → datarobot_genai-0.2.21.dist-info}/RECORD +13 -10
- {datarobot_genai-0.2.19.dist-info → datarobot_genai-0.2.21.dist-info}/WHEEL +0 -0
- {datarobot_genai-0.2.19.dist-info → datarobot_genai-0.2.21.dist-info}/entry_points.txt +0 -0
- {datarobot_genai-0.2.19.dist-info → datarobot_genai-0.2.21.dist-info}/licenses/AUTHORS +0 -0
- {datarobot_genai-0.2.19.dist-info → datarobot_genai-0.2.21.dist-info}/licenses/LICENSE +0 -0
|
@@ -245,6 +245,30 @@ class MCPServerConfig(BaseSettings):
|
|
|
245
245
|
os.getenv("CONFLUENCE_CLIENT_ID") and os.getenv("CONFLUENCE_CLIENT_SECRET")
|
|
246
246
|
)
|
|
247
247
|
|
|
248
|
+
# Gdrive tools
|
|
249
|
+
enable_gdrive_tools: bool = Field(
|
|
250
|
+
default=False,
|
|
251
|
+
validation_alias=AliasChoices(
|
|
252
|
+
RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_GDRIVE_TOOLS",
|
|
253
|
+
"ENABLE_GDRIVE_TOOLS",
|
|
254
|
+
),
|
|
255
|
+
description="Enable/disable GDrive tools",
|
|
256
|
+
)
|
|
257
|
+
is_gdrive_oauth_provider_configured: bool = Field(
|
|
258
|
+
default=False,
|
|
259
|
+
validation_alias=AliasChoices(
|
|
260
|
+
RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
|
|
261
|
+
"IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
|
|
262
|
+
),
|
|
263
|
+
description="Whether GDrive OAuth provider is configured for GDrive integration",
|
|
264
|
+
)
|
|
265
|
+
|
|
266
|
+
@property
|
|
267
|
+
def is_gdrive_oauth_configured(self) -> bool:
|
|
268
|
+
return self.is_gdrive_oauth_provider_configured or bool(
|
|
269
|
+
os.getenv("GDRIVE_CLIENT_ID") and os.getenv("GDRIVE_CLIENT_SECRET")
|
|
270
|
+
)
|
|
271
|
+
|
|
248
272
|
@field_validator(
|
|
249
273
|
"otel_attributes",
|
|
250
274
|
mode="before",
|
|
@@ -29,6 +29,7 @@ class ToolType(str, Enum):
|
|
|
29
29
|
PREDICTIVE = "predictive"
|
|
30
30
|
JIRA = "jira"
|
|
31
31
|
CONFLUENCE = "confluence"
|
|
32
|
+
GDRIVE = "gdrive"
|
|
32
33
|
|
|
33
34
|
|
|
34
35
|
class ToolConfig(TypedDict):
|
|
@@ -64,6 +65,13 @@ TOOL_CONFIGS: dict[ToolType, ToolConfig] = {
|
|
|
64
65
|
package_prefix="datarobot_genai.drmcp.tools.confluence",
|
|
65
66
|
config_field_name="enable_confluence_tools",
|
|
66
67
|
),
|
|
68
|
+
ToolType.GDRIVE: ToolConfig(
|
|
69
|
+
name="gdrive",
|
|
70
|
+
oauth_check=lambda config: config.is_gdrive_oauth_configured,
|
|
71
|
+
directory="gdrive",
|
|
72
|
+
package_prefix="datarobot_genai.drmcp.tools.gdrive",
|
|
73
|
+
config_field_name="enable_gdrive_tools",
|
|
74
|
+
),
|
|
67
75
|
}
|
|
68
76
|
|
|
69
77
|
|
|
@@ -30,6 +30,9 @@ from .atlassian import get_atlassian_cloud_id
|
|
|
30
30
|
|
|
31
31
|
logger = logging.getLogger(__name__)
|
|
32
32
|
|
|
33
|
+
# Search expand fields for CQL search - content.space gives us space.key directly
|
|
34
|
+
SEARCH_EXPAND_FIELDS = "content.space"
|
|
35
|
+
|
|
33
36
|
|
|
34
37
|
class ConfluenceError(Exception):
|
|
35
38
|
"""Exception for Confluence API errors."""
|
|
@@ -75,6 +78,32 @@ class ConfluenceComment(BaseModel):
|
|
|
75
78
|
}
|
|
76
79
|
|
|
77
80
|
|
|
81
|
+
class ContentSearchResult(BaseModel):
|
|
82
|
+
"""Pydantic model for Confluence search result item."""
|
|
83
|
+
|
|
84
|
+
id: str
|
|
85
|
+
title: str
|
|
86
|
+
type: str
|
|
87
|
+
space_key: str = ""
|
|
88
|
+
space_name: str = ""
|
|
89
|
+
excerpt: str = ""
|
|
90
|
+
last_modified: str | None = None
|
|
91
|
+
url: str = ""
|
|
92
|
+
|
|
93
|
+
def as_flat_dict(self) -> dict[str, Any]:
|
|
94
|
+
"""Return a flat dictionary representation of the search result."""
|
|
95
|
+
return {
|
|
96
|
+
"id": self.id,
|
|
97
|
+
"title": self.title,
|
|
98
|
+
"type": self.type,
|
|
99
|
+
"spaceKey": self.space_key,
|
|
100
|
+
"spaceName": self.space_name,
|
|
101
|
+
"excerpt": self.excerpt,
|
|
102
|
+
"lastModified": self.last_modified,
|
|
103
|
+
"url": self.url,
|
|
104
|
+
}
|
|
105
|
+
|
|
106
|
+
|
|
78
107
|
class ConfluenceClient:
|
|
79
108
|
"""
|
|
80
109
|
Client for interacting with Confluence API using OAuth access token.
|
|
@@ -382,6 +411,76 @@ class ConfluenceClient:
|
|
|
382
411
|
|
|
383
412
|
return self._parse_comment_response(response.json(), page_id)
|
|
384
413
|
|
|
414
|
+
async def search_confluence_content(
|
|
415
|
+
self, cql_query: str, max_results: int
|
|
416
|
+
) -> list[ContentSearchResult]:
|
|
417
|
+
"""
|
|
418
|
+
Search Confluence content using CQL (Confluence Query Language).
|
|
419
|
+
|
|
420
|
+
Args:
|
|
421
|
+
cql_query: CQL Query
|
|
422
|
+
max_results: Maximum number of results to return
|
|
423
|
+
|
|
424
|
+
Returns
|
|
425
|
+
-------
|
|
426
|
+
List of Confluence content search results
|
|
427
|
+
|
|
428
|
+
Raises
|
|
429
|
+
------
|
|
430
|
+
ConfluenceError: If the API request fails (400, 403, 429)
|
|
431
|
+
"""
|
|
432
|
+
cloud_id = await self._get_cloud_id()
|
|
433
|
+
url = f"{ATLASSIAN_API_BASE}/ex/confluence/{cloud_id}/wiki/rest/api/search"
|
|
434
|
+
|
|
435
|
+
response = await self._client.get(
|
|
436
|
+
url,
|
|
437
|
+
params={
|
|
438
|
+
"cql": cql_query,
|
|
439
|
+
"limit": max_results,
|
|
440
|
+
"expand": SEARCH_EXPAND_FIELDS,
|
|
441
|
+
},
|
|
442
|
+
)
|
|
443
|
+
|
|
444
|
+
if response.status_code == HTTPStatus.BAD_REQUEST:
|
|
445
|
+
error_msg = self._extract_error_message(response)
|
|
446
|
+
raise ConfluenceError(f"Invalid CQL query: {error_msg}", status_code=400)
|
|
447
|
+
|
|
448
|
+
if response.status_code == HTTPStatus.FORBIDDEN:
|
|
449
|
+
raise ConfluenceError(
|
|
450
|
+
"Permission denied: you don't have access to search this content",
|
|
451
|
+
status_code=403,
|
|
452
|
+
)
|
|
453
|
+
|
|
454
|
+
if response.status_code == HTTPStatus.TOO_MANY_REQUESTS:
|
|
455
|
+
raise ConfluenceError("Rate limit exceeded. Please try again later.", status_code=429)
|
|
456
|
+
|
|
457
|
+
response.raise_for_status()
|
|
458
|
+
raw_results = response.json().get("results", [])
|
|
459
|
+
results = [ContentSearchResult(**self._parse_search_item(item)) for item in raw_results]
|
|
460
|
+
return results
|
|
461
|
+
|
|
462
|
+
def _parse_search_item(self, item: dict) -> dict:
|
|
463
|
+
"""Parse raw search API response item into model-compatible dict."""
|
|
464
|
+
content = item.get("content", item)
|
|
465
|
+
links = content.get("_links", {})
|
|
466
|
+
base_url = links.get("base", "")
|
|
467
|
+
webui = links.get("webui", "")
|
|
468
|
+
url = f"{base_url}{webui}" if base_url and webui else webui
|
|
469
|
+
|
|
470
|
+
# Get space from content.space (requires expand=content.space)
|
|
471
|
+
content_space = content.get("space", {})
|
|
472
|
+
|
|
473
|
+
return {
|
|
474
|
+
"id": str(content.get("id", "")),
|
|
475
|
+
"title": content.get("title", ""),
|
|
476
|
+
"type": content.get("type", "page"),
|
|
477
|
+
"space_key": content_space.get("key", ""),
|
|
478
|
+
"space_name": content_space.get("name", ""),
|
|
479
|
+
"excerpt": item.get("excerpt", ""),
|
|
480
|
+
"last_modified": item.get("lastModified"),
|
|
481
|
+
"url": url,
|
|
482
|
+
}
|
|
483
|
+
|
|
385
484
|
async def __aenter__(self) -> "ConfluenceClient":
|
|
386
485
|
"""Async context manager entry."""
|
|
387
486
|
return self
|
|
@@ -0,0 +1,248 @@
|
|
|
1
|
+
# Copyright 2025 DataRobot, Inc.
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
|
|
15
|
+
"""Google Drive API Client and utilities for OAuth."""
|
|
16
|
+
|
|
17
|
+
import logging
|
|
18
|
+
from typing import Annotated
|
|
19
|
+
from typing import Any
|
|
20
|
+
|
|
21
|
+
import httpx
|
|
22
|
+
from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
|
|
23
|
+
from fastmcp.exceptions import ToolError
|
|
24
|
+
from pydantic import BaseModel
|
|
25
|
+
from pydantic import ConfigDict
|
|
26
|
+
from pydantic import Field
|
|
27
|
+
|
|
28
|
+
from datarobot_genai.drmcp.core.auth import get_access_token
|
|
29
|
+
|
|
30
|
+
logger = logging.getLogger(__name__)
|
|
31
|
+
|
|
32
|
+
DEFAULT_FIELDS = "nextPageToken,files(id,name,size,mimeType,webViewLink,createdTime,modifiedTime)"
|
|
33
|
+
DEFAULT_ORDER = "modifiedTime desc"
|
|
34
|
+
MAX_PAGE_SIZE = 100
|
|
35
|
+
LIMIT = 500
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
async def get_gdrive_access_token() -> str | ToolError:
|
|
39
|
+
"""
|
|
40
|
+
Get Google Drive OAuth access token with error handling.
|
|
41
|
+
|
|
42
|
+
Returns
|
|
43
|
+
-------
|
|
44
|
+
Access token string on success, ToolError on failure
|
|
45
|
+
|
|
46
|
+
Example:
|
|
47
|
+
```python
|
|
48
|
+
token = await get_gdrive_access_token()
|
|
49
|
+
if isinstance(token, ToolError):
|
|
50
|
+
# Handle error
|
|
51
|
+
return token
|
|
52
|
+
# Use token
|
|
53
|
+
```
|
|
54
|
+
"""
|
|
55
|
+
try:
|
|
56
|
+
access_token = await get_access_token("google")
|
|
57
|
+
if not access_token:
|
|
58
|
+
logger.warning("Empty access token received")
|
|
59
|
+
return ToolError("Received empty access token. Please complete the OAuth flow.")
|
|
60
|
+
return access_token
|
|
61
|
+
except OAuthServiceClientErr as e:
|
|
62
|
+
logger.error(f"OAuth client error: {e}", exc_info=True)
|
|
63
|
+
return ToolError(
|
|
64
|
+
"Could not obtain access token for Google. Make sure the OAuth "
|
|
65
|
+
"permission was granted for the application to act on your behalf."
|
|
66
|
+
)
|
|
67
|
+
except Exception as e:
|
|
68
|
+
logger.error(f"Unexpected error obtaining access token: {e}", exc_info=True)
|
|
69
|
+
return ToolError("An unexpected error occurred while obtaining access token for Google.")
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
class GoogleDriveError(Exception):
|
|
73
|
+
"""Exception for Google Drive API errors."""
|
|
74
|
+
|
|
75
|
+
def __init__(self, message: str) -> None:
|
|
76
|
+
super().__init__(message)
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
PrimitiveData = str | int | float | bool | None
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
class GoogleDriveFile(BaseModel):
|
|
83
|
+
"""Represents a file from Google Drive."""
|
|
84
|
+
|
|
85
|
+
id: str
|
|
86
|
+
name: str
|
|
87
|
+
mime_type: Annotated[str, Field(alias="mimeType")]
|
|
88
|
+
size: int | None = None
|
|
89
|
+
web_view_link: Annotated[str | None, Field(alias="webViewLink")] = None
|
|
90
|
+
created_time: Annotated[str | None, Field(alias="createdTime")] = None
|
|
91
|
+
modified_time: Annotated[str | None, Field(alias="modifiedTime")] = None
|
|
92
|
+
|
|
93
|
+
model_config = ConfigDict(populate_by_name=True)
|
|
94
|
+
|
|
95
|
+
@classmethod
|
|
96
|
+
def from_api_response(cls, data: dict[str, Any]) -> "GoogleDriveFile":
|
|
97
|
+
"""Create a GoogleDriveFile from API response data."""
|
|
98
|
+
return cls(
|
|
99
|
+
id=data.get("id", "Unknown"),
|
|
100
|
+
name=data.get("name", "Unknown"),
|
|
101
|
+
mime_type=data.get("mimeType", "Unknown"),
|
|
102
|
+
size=int(data["size"]) if data.get("size") else None,
|
|
103
|
+
web_view_link=data.get("webViewLink"),
|
|
104
|
+
created_time=data.get("createdTime"),
|
|
105
|
+
modified_time=data.get("modifiedTime"),
|
|
106
|
+
)
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
class PaginatedResult(BaseModel):
|
|
110
|
+
"""Result of a paginated API call."""
|
|
111
|
+
|
|
112
|
+
files: list[GoogleDriveFile]
|
|
113
|
+
next_page_token: str | None = None
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
class GoogleDriveClient:
|
|
117
|
+
"""Client for interacting with Google Drive API."""
|
|
118
|
+
|
|
119
|
+
def __init__(self, access_token: str) -> None:
|
|
120
|
+
self._client = httpx.AsyncClient(
|
|
121
|
+
base_url="https://www.googleapis.com/drive/v3/files",
|
|
122
|
+
headers={"Authorization": f"Bearer {access_token}"},
|
|
123
|
+
timeout=30.0,
|
|
124
|
+
)
|
|
125
|
+
|
|
126
|
+
async def list_files(
|
|
127
|
+
self,
|
|
128
|
+
page_size: int,
|
|
129
|
+
limit: int,
|
|
130
|
+
page_token: str | None = None,
|
|
131
|
+
query: str | None = None,
|
|
132
|
+
) -> PaginatedResult:
|
|
133
|
+
"""
|
|
134
|
+
List files from Google Drive.
|
|
135
|
+
|
|
136
|
+
It's public API for GoogleDriveClient.
|
|
137
|
+
|
|
138
|
+
Args:
|
|
139
|
+
page_size: Number of files to return per 1 gdrive api request.
|
|
140
|
+
limit: Maximum number of files to return.
|
|
141
|
+
page_token: Optional token (specific for gdrive api) allowing to query next page.
|
|
142
|
+
query: Optional query to filter results.
|
|
143
|
+
If not provided it'll list all authorized user files.
|
|
144
|
+
If the query doesn't contain operators (contains, =, etc.), it will be treated as
|
|
145
|
+
a name search: "name contains '{query}'".
|
|
146
|
+
|
|
147
|
+
Returns
|
|
148
|
+
-------
|
|
149
|
+
List of Google Drive files.
|
|
150
|
+
"""
|
|
151
|
+
if page_size <= 0:
|
|
152
|
+
raise GoogleDriveError("Error: page size must be positive.")
|
|
153
|
+
if limit <= 0:
|
|
154
|
+
raise GoogleDriveError("Error: limit must be positive.")
|
|
155
|
+
if limit < page_size:
|
|
156
|
+
raise GoogleDriveError("Error: limit must be bigger than or equal to page size.")
|
|
157
|
+
if limit % page_size != 0:
|
|
158
|
+
raise GoogleDriveError("Error: limit must be multiplication of page size.")
|
|
159
|
+
|
|
160
|
+
page_size = min(page_size, MAX_PAGE_SIZE)
|
|
161
|
+
limit = min(limit, LIMIT)
|
|
162
|
+
fetched = 0
|
|
163
|
+
|
|
164
|
+
formatted_query = self._get_formatted_query(query)
|
|
165
|
+
|
|
166
|
+
files: list[GoogleDriveFile] = []
|
|
167
|
+
|
|
168
|
+
while fetched < limit:
|
|
169
|
+
data = await self._list_files(
|
|
170
|
+
page_size=page_size,
|
|
171
|
+
page_token=page_token,
|
|
172
|
+
query=formatted_query,
|
|
173
|
+
)
|
|
174
|
+
files.extend(data.files)
|
|
175
|
+
fetched += len(data.files)
|
|
176
|
+
page_token = data.next_page_token
|
|
177
|
+
|
|
178
|
+
if not page_token:
|
|
179
|
+
break
|
|
180
|
+
|
|
181
|
+
return PaginatedResult(files=files, next_page_token=page_token)
|
|
182
|
+
|
|
183
|
+
async def _list_files(
|
|
184
|
+
self,
|
|
185
|
+
page_size: int,
|
|
186
|
+
page_token: str | None = None,
|
|
187
|
+
query: str | None = None,
|
|
188
|
+
) -> PaginatedResult:
|
|
189
|
+
"""Fetch a page of files from Google Drive."""
|
|
190
|
+
params: dict[str, PrimitiveData] = {
|
|
191
|
+
"pageSize": page_size,
|
|
192
|
+
"fields": DEFAULT_FIELDS,
|
|
193
|
+
"orderBy": DEFAULT_ORDER,
|
|
194
|
+
}
|
|
195
|
+
if page_token:
|
|
196
|
+
params["pageToken"] = page_token
|
|
197
|
+
if query:
|
|
198
|
+
params["q"] = query
|
|
199
|
+
|
|
200
|
+
response = await self._client.get(url="/", params=params)
|
|
201
|
+
response.raise_for_status()
|
|
202
|
+
data = response.json()
|
|
203
|
+
|
|
204
|
+
files = [
|
|
205
|
+
GoogleDriveFile.from_api_response(file_data) for file_data in data.get("files", [])
|
|
206
|
+
]
|
|
207
|
+
next_page_token = data.get("nextPageToken")
|
|
208
|
+
return PaginatedResult(files=files, next_page_token=next_page_token)
|
|
209
|
+
|
|
210
|
+
@staticmethod
|
|
211
|
+
def _get_formatted_query(query: str | None) -> str | None:
|
|
212
|
+
"""Get formatted Google Drive API query.
|
|
213
|
+
|
|
214
|
+
Args:
|
|
215
|
+
query: Optional search query string (e.g., "name contains 'report'"").
|
|
216
|
+
If the query doesn't contain operators (contains, =, etc.), it will be treated as
|
|
217
|
+
a name search: "name contains '{query}'".
|
|
218
|
+
|
|
219
|
+
Returns
|
|
220
|
+
-------
|
|
221
|
+
Correctly formatted query (if provided)
|
|
222
|
+
"""
|
|
223
|
+
if not query:
|
|
224
|
+
return None
|
|
225
|
+
|
|
226
|
+
# If query doesn't look like a formatted query (no operators), format it as a name search
|
|
227
|
+
# Check if query already has Google Drive API operators
|
|
228
|
+
has_operator = any(
|
|
229
|
+
op in query for op in [" contains ", "=", "!=", " in ", " and ", " or ", " not "]
|
|
230
|
+
)
|
|
231
|
+
formatted_query = query
|
|
232
|
+
if not has_operator and query.strip():
|
|
233
|
+
# Simple text search - format as name contains query
|
|
234
|
+
# Escape backslashes first, then single quotes for Google Drive API
|
|
235
|
+
escaped_query = query.replace("\\", "\\\\").replace("'", "\\'")
|
|
236
|
+
formatted_query = f"name contains '{escaped_query}'"
|
|
237
|
+
logger.debug(f"Auto-formatted query '{query}' to '{formatted_query}'")
|
|
238
|
+
return formatted_query
|
|
239
|
+
|
|
240
|
+
async def __aenter__(self) -> "GoogleDriveClient":
|
|
241
|
+
"""Async context manager entry."""
|
|
242
|
+
return self
|
|
243
|
+
|
|
244
|
+
async def __aexit__(
|
|
245
|
+
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
|
|
246
|
+
) -> None:
|
|
247
|
+
"""Async context manager exit."""
|
|
248
|
+
await self._client.aclose()
|
|
@@ -186,3 +186,69 @@ async def confluence_add_comment(
|
|
|
186
186
|
"page_id": page_id,
|
|
187
187
|
},
|
|
188
188
|
)
|
|
189
|
+
|
|
190
|
+
|
|
191
|
+
@dr_mcp_tool(tags={"confluence", "search", "content"})
|
|
192
|
+
async def confluence_search(
|
|
193
|
+
*,
|
|
194
|
+
cql_query: Annotated[
|
|
195
|
+
str,
|
|
196
|
+
"The CQL (Confluence Query Language) string used to filter content, "
|
|
197
|
+
"e.g., 'type=page and space=DOC'.",
|
|
198
|
+
],
|
|
199
|
+
max_results: Annotated[int, "Maximum number of content items to return. Default is 10."] = 10,
|
|
200
|
+
include_body: Annotated[
|
|
201
|
+
bool,
|
|
202
|
+
"If True, fetch full page body content for each result (slower, "
|
|
203
|
+
"makes additional API calls). Default is False, which returns only excerpts.",
|
|
204
|
+
] = False,
|
|
205
|
+
) -> ToolResult:
|
|
206
|
+
"""
|
|
207
|
+
Search Confluence pages and content efficiently using a CQL query string.
|
|
208
|
+
This pushes the search logic to the Confluence API (Push-Down).
|
|
209
|
+
|
|
210
|
+
Refer to Confluence documentation for advanced searching using CQL:
|
|
211
|
+
https://developer.atlassian.com/cloud/confluence/advanced-searching-using-cql/
|
|
212
|
+
"""
|
|
213
|
+
if not cql_query:
|
|
214
|
+
raise ToolError("Argument validation error: 'cql_query' cannot be empty.")
|
|
215
|
+
|
|
216
|
+
if max_results < 1 or max_results > 100:
|
|
217
|
+
raise ToolError("Argument validation error: 'max_results' must be between 1 and 100.")
|
|
218
|
+
|
|
219
|
+
access_token = await get_atlassian_access_token()
|
|
220
|
+
if isinstance(access_token, ToolError):
|
|
221
|
+
raise access_token
|
|
222
|
+
|
|
223
|
+
try:
|
|
224
|
+
async with ConfluenceClient(access_token) as client:
|
|
225
|
+
results = await client.search_confluence_content(
|
|
226
|
+
cql_query=cql_query, max_results=max_results
|
|
227
|
+
)
|
|
228
|
+
|
|
229
|
+
# If include_body is True, fetch full content for each page
|
|
230
|
+
if include_body and results:
|
|
231
|
+
data = []
|
|
232
|
+
for result in results:
|
|
233
|
+
flat = result.as_flat_dict()
|
|
234
|
+
try:
|
|
235
|
+
page = await client.get_page_by_id(result.id)
|
|
236
|
+
flat["body"] = page.body
|
|
237
|
+
except ConfluenceError:
|
|
238
|
+
flat["body"] = None # Keep excerpt if page fetch fails
|
|
239
|
+
data.append(flat)
|
|
240
|
+
else:
|
|
241
|
+
data = [result.as_flat_dict() for result in results]
|
|
242
|
+
|
|
243
|
+
except ConfluenceError as e:
|
|
244
|
+
logger.error(f"Confluence error searching content: {e}")
|
|
245
|
+
raise ToolError(str(e))
|
|
246
|
+
except Exception as e:
|
|
247
|
+
logger.error(f"Unexpected error searching Confluence content: {e}")
|
|
248
|
+
raise ToolError(f"An unexpected error occurred while searching Confluence: {str(e)}")
|
|
249
|
+
|
|
250
|
+
n = len(results)
|
|
251
|
+
return ToolResult(
|
|
252
|
+
content=f"Successfully executed CQL query and retrieved {n} result(s).",
|
|
253
|
+
structured_content={"data": data, "count": n},
|
|
254
|
+
)
|
|
File without changes
|
|
@@ -0,0 +1,88 @@
|
|
|
1
|
+
# Copyright 2025 DataRobot, Inc.
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
|
|
15
|
+
"""Google Drive MCP tools for interacting with Google Drive API."""
|
|
16
|
+
|
|
17
|
+
import logging
|
|
18
|
+
from typing import Annotated
|
|
19
|
+
|
|
20
|
+
from fastmcp.exceptions import ToolError
|
|
21
|
+
from fastmcp.tools.tool import ToolResult
|
|
22
|
+
|
|
23
|
+
from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
|
|
24
|
+
from datarobot_genai.drmcp.tools.clients.gdrive import LIMIT
|
|
25
|
+
from datarobot_genai.drmcp.tools.clients.gdrive import MAX_PAGE_SIZE
|
|
26
|
+
from datarobot_genai.drmcp.tools.clients.gdrive import GoogleDriveClient
|
|
27
|
+
from datarobot_genai.drmcp.tools.clients.gdrive import GoogleDriveError
|
|
28
|
+
from datarobot_genai.drmcp.tools.clients.gdrive import get_gdrive_access_token
|
|
29
|
+
|
|
30
|
+
logger = logging.getLogger(__name__)
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
@dr_mcp_tool(tags={"google", "gdrive", "list", "files"})
|
|
34
|
+
async def google_drive_list_files(
|
|
35
|
+
*,
|
|
36
|
+
page_size: Annotated[
|
|
37
|
+
int, f"Maximum number of files to return per page (max {MAX_PAGE_SIZE})."
|
|
38
|
+
] = 10,
|
|
39
|
+
limit: Annotated[int, f"Total maximum number of files to return (max {LIMIT})."] = 50,
|
|
40
|
+
page_token: Annotated[
|
|
41
|
+
str | None, "The token for the next page of results, retrieved from a previous call."
|
|
42
|
+
] = None,
|
|
43
|
+
query: Annotated[
|
|
44
|
+
str | None, "Optional filter to narrow results (e.g., 'trashed = false')."
|
|
45
|
+
] = None,
|
|
46
|
+
) -> ToolResult | ToolError:
|
|
47
|
+
"""
|
|
48
|
+
List files in the user's Google Drive with pagination and filtering support.
|
|
49
|
+
Use this tool to discover file names and IDs for use with other tools.
|
|
50
|
+
|
|
51
|
+
Limit must be bigger than or equal to page size and it must be multiplication of page size.
|
|
52
|
+
Ex.
|
|
53
|
+
page size = 10 limit = 50
|
|
54
|
+
page size = 3 limit = 3
|
|
55
|
+
page size = 12 limit = 36
|
|
56
|
+
"""
|
|
57
|
+
access_token = await get_gdrive_access_token()
|
|
58
|
+
if isinstance(access_token, ToolError):
|
|
59
|
+
raise access_token
|
|
60
|
+
|
|
61
|
+
try:
|
|
62
|
+
async with GoogleDriveClient(access_token) as client:
|
|
63
|
+
data = await client.list_files(
|
|
64
|
+
page_size=page_size, page_token=page_token, query=query, limit=limit
|
|
65
|
+
)
|
|
66
|
+
except GoogleDriveError as e:
|
|
67
|
+
logger.error(f"Google Drive error listing files: {e}")
|
|
68
|
+
raise ToolError(str(e))
|
|
69
|
+
except Exception as e:
|
|
70
|
+
logger.error(f"Unexpected error listing Google Drive files: {e}")
|
|
71
|
+
raise ToolError(f"An unexpected error occurred while listing Google Drive files: {str(e)}")
|
|
72
|
+
|
|
73
|
+
number_of_files = len(data.files)
|
|
74
|
+
next_page_info = (
|
|
75
|
+
f"Next page token needed to fetch more data: {data.next_page_token}"
|
|
76
|
+
if data.next_page_token
|
|
77
|
+
else "There're no more pages."
|
|
78
|
+
)
|
|
79
|
+
return ToolResult(
|
|
80
|
+
content=f"Successfully listed {number_of_files} files. {next_page_info}",
|
|
81
|
+
structured_content={
|
|
82
|
+
"files": [
|
|
83
|
+
file.model_dump(by_alias=True, include={"id", "name"}) for file in data.files
|
|
84
|
+
],
|
|
85
|
+
"count": number_of_files,
|
|
86
|
+
"nextPageToken": data.next_page_token,
|
|
87
|
+
},
|
|
88
|
+
)
|
|
@@ -27,7 +27,7 @@ datarobot_genai/drmcp/server.py,sha256=KE4kjS5f9bfdYftG14HBHrfvxDfCD4pwCXePfvl1O
|
|
|
27
27
|
datarobot_genai/drmcp/core/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
|
|
28
28
|
datarobot_genai/drmcp/core/auth.py,sha256=E-5wrGbBFEBlD5377g6Exddrc7HsazamwX8tWr2RLXY,5815
|
|
29
29
|
datarobot_genai/drmcp/core/clients.py,sha256=y-yG8617LbmiZ_L7FWfMrk4WjIekyr76u_Q80aLqGpI,5524
|
|
30
|
-
datarobot_genai/drmcp/core/config.py,sha256=
|
|
30
|
+
datarobot_genai/drmcp/core/config.py,sha256=69QDsVVSvjzv1uIHOjtQGzdg7_Ic4sA3vLA6ZbZJ1Ok,12674
|
|
31
31
|
datarobot_genai/drmcp/core/config_utils.py,sha256=U-aieWw7MyP03cGDFIp97JH99ZUfr3vD9uuTzBzxn7w,6428
|
|
32
32
|
datarobot_genai/drmcp/core/constants.py,sha256=lUwoW_PTrbaBGqRJifKqCn3EoFacoEgdO-CpoFVrUoU,739
|
|
33
33
|
datarobot_genai/drmcp/core/credentials.py,sha256=PYEUDNMVw1BoMzZKLkPVTypNkVevEPtmk3scKnE-zYg,6706
|
|
@@ -41,7 +41,7 @@ datarobot_genai/drmcp/core/routes.py,sha256=dqE2M0UzAyyN9vQjlyTjYW4rpju3LT039po5
|
|
|
41
41
|
datarobot_genai/drmcp/core/routes_utils.py,sha256=vSseXWlplMSnRgoJgtP_rHxWSAVYcx_tpTv4lyTpQoc,944
|
|
42
42
|
datarobot_genai/drmcp/core/server_life_cycle.py,sha256=WKGJWGxalvqxupzJ2y67Kklc_9PgpZT0uyjlv_sr5wc,3419
|
|
43
43
|
datarobot_genai/drmcp/core/telemetry.py,sha256=NEkSTC1w6uQgtukLHI-sWvR4EMgInysgATcvfQ5CplM,15378
|
|
44
|
-
datarobot_genai/drmcp/core/tool_config.py,sha256=
|
|
44
|
+
datarobot_genai/drmcp/core/tool_config.py,sha256=5JCWO70ZH-K-34yS7vYJG2nl4i9UO_q_W9NCoWSXXno,3271
|
|
45
45
|
datarobot_genai/drmcp/core/tool_filter.py,sha256=tLOcG50QBvS48cOVHM6OqoODYiiS6KeM_F-2diaHkW0,2858
|
|
46
46
|
datarobot_genai/drmcp/core/utils.py,sha256=EvfpqKZ3tECMoxpIQ_tA_3rOgy6KJEYKC0lWZo_Daag,4517
|
|
47
47
|
datarobot_genai/drmcp/core/dynamic_prompts/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
|
|
@@ -78,11 +78,14 @@ datarobot_genai/drmcp/test_utils/utils.py,sha256=esGKFv8aO31-Qg3owayeWp32BYe1CdY
|
|
|
78
78
|
datarobot_genai/drmcp/tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
|
|
79
79
|
datarobot_genai/drmcp/tools/clients/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
|
|
80
80
|
datarobot_genai/drmcp/tools/clients/atlassian.py,sha256=__M_uz7FrcbKCYRzeMn24DCEYD6OmFx_LuywHCxgXsA,6472
|
|
81
|
-
datarobot_genai/drmcp/tools/clients/confluence.py,sha256=
|
|
81
|
+
datarobot_genai/drmcp/tools/clients/confluence.py,sha256=YS5XsKd-jK5Yg0rgwOcC76v9e8fDJgUZIW5B9kcq5B0,17101
|
|
82
|
+
datarobot_genai/drmcp/tools/clients/gdrive.py,sha256=QmNTmJdSqYO5Y5Vnp3roNZiNNJeocBVjF9UcSzcjgRY,8635
|
|
82
83
|
datarobot_genai/drmcp/tools/clients/jira.py,sha256=Rm91JAyrNIqxu66-9rU1YqoRXVnWbEy-Ahvy6f6HlVg,9823
|
|
83
84
|
datarobot_genai/drmcp/tools/clients/s3.py,sha256=GmwzvurFdNfvxOooA8g5S4osRysHYU0S9ypg_177Glg,953
|
|
84
85
|
datarobot_genai/drmcp/tools/confluence/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
|
|
85
|
-
datarobot_genai/drmcp/tools/confluence/tools.py,sha256=
|
|
86
|
+
datarobot_genai/drmcp/tools/confluence/tools.py,sha256=ySwABe8osAzky3BO3lRaF6UHnXQgaurkmvM0iHFfL30,9849
|
|
87
|
+
datarobot_genai/drmcp/tools/gdrive/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
88
|
+
datarobot_genai/drmcp/tools/gdrive/tools.py,sha256=wmCUSaCWqepdlOIApA8tZ-grPYSV7wZKoer6uRy26Qg,3459
|
|
86
89
|
datarobot_genai/drmcp/tools/jira/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
|
|
87
90
|
datarobot_genai/drmcp/tools/jira/tools.py,sha256=dfkqTU2HH-7n44hX80ODFacKq0p0LOchFcZtIIKFNMM,9687
|
|
88
91
|
datarobot_genai/drmcp/tools/predictive/__init__.py,sha256=WuOHlNNEpEmcF7gVnhckruJRKU2qtmJLE3E7zoCGLDo,1030
|
|
@@ -108,9 +111,9 @@ datarobot_genai/nat/datarobot_llm_clients.py,sha256=Yu208Ed_p_4P3HdpuM7fYnKcXtim
|
|
|
108
111
|
datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
|
|
109
112
|
datarobot_genai/nat/datarobot_mcp_client.py,sha256=35FzilxNp4VqwBYI0NsOc91-xZm1C-AzWqrOdDy962A,9612
|
|
110
113
|
datarobot_genai/nat/helpers.py,sha256=Q7E3ADZdtFfS8E6OQPyw2wgA6laQ58N3bhLj5CBWwJs,3265
|
|
111
|
-
datarobot_genai-0.2.
|
|
112
|
-
datarobot_genai-0.2.
|
|
113
|
-
datarobot_genai-0.2.
|
|
114
|
-
datarobot_genai-0.2.
|
|
115
|
-
datarobot_genai-0.2.
|
|
116
|
-
datarobot_genai-0.2.
|
|
114
|
+
datarobot_genai-0.2.21.dist-info/METADATA,sha256=wdbafq2Z77RRyxy7nG-kYCUU5pfCbsDbBTH7QdENnRU,6301
|
|
115
|
+
datarobot_genai-0.2.21.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
|
|
116
|
+
datarobot_genai-0.2.21.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
|
|
117
|
+
datarobot_genai-0.2.21.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
|
|
118
|
+
datarobot_genai-0.2.21.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
|
|
119
|
+
datarobot_genai-0.2.21.dist-info/RECORD,,
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|