datarobot-genai 0.2.4__py3-none-any.whl → 0.2.6__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 +52 -0
- datarobot_genai/drmcp/core/dr_mcp_server.py +11 -8
- datarobot_genai/drmcp/core/tool_config.py +95 -0
- datarobot_genai/drmcp/tools/clients/__init__.py +14 -0
- datarobot_genai/drmcp/tools/clients/atlassian.py +188 -0
- datarobot_genai/drmcp/tools/clients/confluence.py +14 -0
- datarobot_genai/drmcp/tools/clients/jira.py +102 -0
- datarobot_genai/drmcp/tools/clients/s3.py +28 -0
- datarobot_genai/drmcp/tools/confluence/__init__.py +14 -0
- datarobot_genai/drmcp/tools/confluence/tools.py +13 -0
- datarobot_genai/drmcp/tools/jira/__init__.py +14 -0
- datarobot_genai/drmcp/tools/jira/tools.py +58 -0
- datarobot_genai/drmcp/tools/predictive/predict.py +1 -1
- datarobot_genai/drmcp/tools/predictive/predict_realtime.py +1 -1
- datarobot_genai/nat/agent.py +4 -0
- datarobot_genai/nat/datarobot_auth_provider.py +53 -0
- datarobot_genai/nat/datarobot_mcp_client.py +181 -0
- {datarobot_genai-0.2.4.dist-info → datarobot_genai-0.2.6.dist-info}/METADATA +2 -1
- {datarobot_genai-0.2.4.dist-info → datarobot_genai-0.2.6.dist-info}/RECORD +23 -11
- {datarobot_genai-0.2.4.dist-info → datarobot_genai-0.2.6.dist-info}/entry_points.txt +2 -0
- {datarobot_genai-0.2.4.dist-info → datarobot_genai-0.2.6.dist-info}/WHEEL +0 -0
- {datarobot_genai-0.2.4.dist-info → datarobot_genai-0.2.6.dist-info}/licenses/AUTHORS +0 -0
- {datarobot_genai-0.2.4.dist-info → datarobot_genai-0.2.6.dist-info}/licenses/LICENSE +0 -0
|
@@ -197,6 +197,54 @@ class MCPServerConfig(BaseSettings):
|
|
|
197
197
|
description="Enable/disable predictive tools",
|
|
198
198
|
)
|
|
199
199
|
|
|
200
|
+
# Jira tools
|
|
201
|
+
enable_jira_tools: bool = Field(
|
|
202
|
+
default=False,
|
|
203
|
+
validation_alias=AliasChoices(
|
|
204
|
+
RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_JIRA_TOOLS",
|
|
205
|
+
"ENABLE_JIRA_TOOLS",
|
|
206
|
+
),
|
|
207
|
+
description="Enable/disable Jira tools",
|
|
208
|
+
)
|
|
209
|
+
is_jira_oauth_provider_configured: bool = Field(
|
|
210
|
+
default=False,
|
|
211
|
+
validation_alias=AliasChoices(
|
|
212
|
+
RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
|
|
213
|
+
"IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
|
|
214
|
+
),
|
|
215
|
+
description="Whether Jira OAuth provider is configured for Jira integration",
|
|
216
|
+
)
|
|
217
|
+
|
|
218
|
+
@property
|
|
219
|
+
def is_jira_oauth_configured(self) -> bool:
|
|
220
|
+
return self.is_jira_oauth_provider_configured or bool(
|
|
221
|
+
os.getenv("JIRA_CLIENT_ID") and os.getenv("JIRA_CLIENT_SECRET")
|
|
222
|
+
)
|
|
223
|
+
|
|
224
|
+
# Confluence tools
|
|
225
|
+
enable_confluence_tools: bool = Field(
|
|
226
|
+
default=False,
|
|
227
|
+
validation_alias=AliasChoices(
|
|
228
|
+
RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_CONFLUENCE_TOOLS",
|
|
229
|
+
"ENABLE_CONFLUENCE_TOOLS",
|
|
230
|
+
),
|
|
231
|
+
description="Enable/disable Confluence tools",
|
|
232
|
+
)
|
|
233
|
+
is_confluence_oauth_provider_configured: bool = Field(
|
|
234
|
+
default=False,
|
|
235
|
+
validation_alias=AliasChoices(
|
|
236
|
+
RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
|
|
237
|
+
"IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
|
|
238
|
+
),
|
|
239
|
+
description="Whether Confluence OAuth provider is configured for Confluence integration",
|
|
240
|
+
)
|
|
241
|
+
|
|
242
|
+
@property
|
|
243
|
+
def is_confluence_oauth_configured(self) -> bool:
|
|
244
|
+
return self.is_confluence_oauth_provider_configured or bool(
|
|
245
|
+
os.getenv("CONFLUENCE_CLIENT_ID") and os.getenv("CONFLUENCE_CLIENT_SECRET")
|
|
246
|
+
)
|
|
247
|
+
|
|
200
248
|
@field_validator(
|
|
201
249
|
"otel_attributes",
|
|
202
250
|
mode="before",
|
|
@@ -220,6 +268,10 @@ class MCPServerConfig(BaseSettings):
|
|
|
220
268
|
"tool_registration_duplicate_behavior",
|
|
221
269
|
"mcp_server_register_dynamic_prompts_on_startup",
|
|
222
270
|
"enable_predictive_tools",
|
|
271
|
+
"enable_jira_tools",
|
|
272
|
+
"is_jira_oauth_provider_configured",
|
|
273
|
+
"enable_confluence_tools",
|
|
274
|
+
"is_confluence_oauth_provider_configured",
|
|
223
275
|
mode="before",
|
|
224
276
|
)
|
|
225
277
|
@classmethod
|
|
@@ -40,6 +40,8 @@ from .routes_utils import prefix_mount_path
|
|
|
40
40
|
from .server_life_cycle import BaseServerLifecycle
|
|
41
41
|
from .telemetry import OtelASGIMiddleware
|
|
42
42
|
from .telemetry import initialize_telemetry
|
|
43
|
+
from .tool_config import TOOL_CONFIGS
|
|
44
|
+
from .tool_config import is_tool_enabled
|
|
43
45
|
|
|
44
46
|
|
|
45
47
|
def _import_modules_from_dir(
|
|
@@ -142,11 +144,12 @@ class DataRobotMCPServer:
|
|
|
142
144
|
|
|
143
145
|
# Load static tools modules
|
|
144
146
|
base_dir = os.path.dirname(os.path.dirname(__file__))
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
|
|
147
|
+
for tool_type, tool_config in TOOL_CONFIGS.items():
|
|
148
|
+
if is_tool_enabled(tool_type, self._config):
|
|
149
|
+
_import_modules_from_dir(
|
|
150
|
+
os.path.join(base_dir, "tools", tool_config["directory"]),
|
|
151
|
+
tool_config["package_prefix"],
|
|
152
|
+
)
|
|
150
153
|
|
|
151
154
|
# Load memory management tools if available
|
|
152
155
|
if self._memory_manager:
|
|
@@ -213,6 +216,9 @@ class DataRobotMCPServer:
|
|
|
213
216
|
self._logger.info("Registering dynamic prompts from prompt management...")
|
|
214
217
|
asyncio.run(register_prompts_from_datarobot_prompt_management())
|
|
215
218
|
|
|
219
|
+
# Execute pre-server start actions
|
|
220
|
+
asyncio.run(self._lifecycle.pre_server_start(self._mcp))
|
|
221
|
+
|
|
216
222
|
# List registered tools, prompts, and resources before starting server
|
|
217
223
|
tools = asyncio.run(self._mcp._list_tools_mcp())
|
|
218
224
|
prompts = asyncio.run(self._mcp._list_prompts_mcp())
|
|
@@ -232,9 +238,6 @@ class DataRobotMCPServer:
|
|
|
232
238
|
for resource in resources:
|
|
233
239
|
self._logger.info(f" > {resource.name}")
|
|
234
240
|
|
|
235
|
-
# Execute pre-server start actions
|
|
236
|
-
asyncio.run(self._lifecycle.pre_server_start(self._mcp))
|
|
237
|
-
|
|
238
241
|
# Create event loop for async operations
|
|
239
242
|
loop = asyncio.new_event_loop()
|
|
240
243
|
asyncio.set_event_loop(loop)
|
|
@@ -0,0 +1,95 @@
|
|
|
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
|
+
"""Tool configuration and enablement logic."""
|
|
16
|
+
|
|
17
|
+
from collections.abc import Callable
|
|
18
|
+
from enum import Enum
|
|
19
|
+
from typing import TYPE_CHECKING
|
|
20
|
+
from typing import TypedDict
|
|
21
|
+
|
|
22
|
+
if TYPE_CHECKING:
|
|
23
|
+
from .config import MCPServerConfig
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class ToolType(str, Enum):
|
|
27
|
+
"""Enumeration of available tool types."""
|
|
28
|
+
|
|
29
|
+
PREDICTIVE = "predictive"
|
|
30
|
+
JIRA = "jira"
|
|
31
|
+
CONFLUENCE = "confluence"
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
class ToolConfig(TypedDict):
|
|
35
|
+
"""Configuration for a tool type."""
|
|
36
|
+
|
|
37
|
+
name: str
|
|
38
|
+
oauth_check: Callable[["MCPServerConfig"], bool] | None
|
|
39
|
+
directory: str
|
|
40
|
+
package_prefix: str
|
|
41
|
+
config_field_name: str # Name of the config field (e.g., "enable_predictive_tools")
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
# Tool configuration registry
|
|
45
|
+
TOOL_CONFIGS: dict[ToolType, ToolConfig] = {
|
|
46
|
+
ToolType.PREDICTIVE: ToolConfig(
|
|
47
|
+
name="predictive",
|
|
48
|
+
oauth_check=None,
|
|
49
|
+
directory="predictive",
|
|
50
|
+
package_prefix="datarobot_genai.drmcp.tools.predictive",
|
|
51
|
+
config_field_name="enable_predictive_tools",
|
|
52
|
+
),
|
|
53
|
+
ToolType.JIRA: ToolConfig(
|
|
54
|
+
name="jira",
|
|
55
|
+
oauth_check=lambda config: config.is_jira_oauth_configured,
|
|
56
|
+
directory="jira",
|
|
57
|
+
package_prefix="datarobot_genai.drmcp.tools.jira",
|
|
58
|
+
config_field_name="enable_jira_tools",
|
|
59
|
+
),
|
|
60
|
+
ToolType.CONFLUENCE: ToolConfig(
|
|
61
|
+
name="confluence",
|
|
62
|
+
oauth_check=lambda config: config.is_confluence_oauth_configured,
|
|
63
|
+
directory="confluence",
|
|
64
|
+
package_prefix="datarobot_genai.drmcp.tools.confluence",
|
|
65
|
+
config_field_name="enable_confluence_tools",
|
|
66
|
+
),
|
|
67
|
+
}
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def get_tool_enable_config_name(tool_type: ToolType) -> str:
|
|
71
|
+
"""Get the configuration field name for enabling a tool."""
|
|
72
|
+
return TOOL_CONFIGS[tool_type]["config_field_name"]
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def is_tool_enabled(tool_type: ToolType, config: "MCPServerConfig") -> bool:
|
|
76
|
+
"""
|
|
77
|
+
Check if a tool is enabled based on configuration.
|
|
78
|
+
|
|
79
|
+
Args:
|
|
80
|
+
tool_type: The type of tool to check
|
|
81
|
+
config: The server configuration
|
|
82
|
+
|
|
83
|
+
Returns
|
|
84
|
+
-------
|
|
85
|
+
True if the tool is enabled, False otherwise
|
|
86
|
+
"""
|
|
87
|
+
tool_config = TOOL_CONFIGS[tool_type]
|
|
88
|
+
enable_config_name = tool_config["config_field_name"]
|
|
89
|
+
is_enabled = getattr(config, enable_config_name)
|
|
90
|
+
|
|
91
|
+
# If tool is enabled, check OAuth requirements if needed
|
|
92
|
+
if is_enabled and tool_config["oauth_check"] is not None:
|
|
93
|
+
return tool_config["oauth_check"](config)
|
|
94
|
+
|
|
95
|
+
return is_enabled
|
|
@@ -0,0 +1,14 @@
|
|
|
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
|
+
|
|
@@ -0,0 +1,188 @@
|
|
|
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
|
+
"""Atlassian API client utilities for OAuth and cloud ID management."""
|
|
16
|
+
|
|
17
|
+
import logging
|
|
18
|
+
from typing import Any
|
|
19
|
+
from typing import Literal
|
|
20
|
+
|
|
21
|
+
import httpx
|
|
22
|
+
from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
|
|
23
|
+
from fastmcp.exceptions import ToolError
|
|
24
|
+
|
|
25
|
+
from datarobot_genai.drmcp.core.auth import get_access_token
|
|
26
|
+
|
|
27
|
+
logger = logging.getLogger(__name__)
|
|
28
|
+
|
|
29
|
+
# Atlassian Cloud API base URL
|
|
30
|
+
ATLASSIAN_API_BASE = "https://api.atlassian.com"
|
|
31
|
+
|
|
32
|
+
# API endpoint paths
|
|
33
|
+
OAUTH_ACCESSIBLE_RESOURCES_PATH = "/oauth/token/accessible-resources"
|
|
34
|
+
|
|
35
|
+
# Supported Atlassian service types
|
|
36
|
+
AtlassianServiceType = Literal["jira", "confluence"]
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
async def get_atlassian_access_token() -> str | ToolError:
|
|
40
|
+
"""
|
|
41
|
+
Get Atlassian OAuth access token with error handling.
|
|
42
|
+
|
|
43
|
+
Returns
|
|
44
|
+
-------
|
|
45
|
+
Access token string on success, ToolError on failure
|
|
46
|
+
|
|
47
|
+
Example:
|
|
48
|
+
```python
|
|
49
|
+
token = await get_atlassian_access_token()
|
|
50
|
+
if isinstance(token, ToolError):
|
|
51
|
+
# Handle error
|
|
52
|
+
return token
|
|
53
|
+
# Use token
|
|
54
|
+
```
|
|
55
|
+
"""
|
|
56
|
+
try:
|
|
57
|
+
access_token = await get_access_token("atlassian")
|
|
58
|
+
if not access_token:
|
|
59
|
+
logger.warning("Empty access token received")
|
|
60
|
+
return ToolError("Received empty access token. Please complete the OAuth flow.")
|
|
61
|
+
return access_token
|
|
62
|
+
except OAuthServiceClientErr as e:
|
|
63
|
+
logger.error(f"OAuth client error: {e}", exc_info=True)
|
|
64
|
+
return ToolError(
|
|
65
|
+
"Could not obtain access token for Atlassian. Make sure the OAuth "
|
|
66
|
+
"permission was granted for the application to act on your behalf."
|
|
67
|
+
)
|
|
68
|
+
except Exception as e:
|
|
69
|
+
logger.error(f"Unexpected error obtaining access token: {e}", exc_info=True)
|
|
70
|
+
return ToolError("An unexpected error occurred while obtaining access token for Atlassian.")
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def _find_resource_by_service(
|
|
74
|
+
resources: list[dict[str, Any]], service_type: str
|
|
75
|
+
) -> dict[str, Any] | None:
|
|
76
|
+
"""
|
|
77
|
+
Find a resource that matches the specified service type.
|
|
78
|
+
|
|
79
|
+
Args:
|
|
80
|
+
resources: List of accessible resources from Atlassian API
|
|
81
|
+
service_type: Service type to filter by (e.g., "jira", "confluence")
|
|
82
|
+
|
|
83
|
+
Returns
|
|
84
|
+
-------
|
|
85
|
+
Resource dictionary if found, None otherwise
|
|
86
|
+
"""
|
|
87
|
+
service_lower = service_type.lower()
|
|
88
|
+
for resource in resources:
|
|
89
|
+
if not resource.get("id"):
|
|
90
|
+
continue
|
|
91
|
+
scopes = resource.get("scopes", [])
|
|
92
|
+
if any(service_lower in scope.lower() for scope in scopes):
|
|
93
|
+
return resource
|
|
94
|
+
return None
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def _find_first_resource_with_id(resources: list[dict[str, Any]]) -> dict[str, Any] | None:
|
|
98
|
+
"""
|
|
99
|
+
Find the first resource that has an ID.
|
|
100
|
+
|
|
101
|
+
Args:
|
|
102
|
+
resources: List of accessible resources from Atlassian API
|
|
103
|
+
|
|
104
|
+
Returns
|
|
105
|
+
-------
|
|
106
|
+
Resource dictionary if found, None otherwise
|
|
107
|
+
"""
|
|
108
|
+
for resource in resources:
|
|
109
|
+
if resource.get("id"):
|
|
110
|
+
return resource
|
|
111
|
+
return None
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
async def get_atlassian_cloud_id(
|
|
115
|
+
client: httpx.AsyncClient,
|
|
116
|
+
service_type: AtlassianServiceType | None = None,
|
|
117
|
+
) -> str:
|
|
118
|
+
"""
|
|
119
|
+
Get the cloud ID for the authenticated Atlassian instance.
|
|
120
|
+
|
|
121
|
+
According to Atlassian OAuth 2.0 documentation, API calls should use:
|
|
122
|
+
https://api.atlassian.com/ex/{service}/{cloudId}/rest/api/3/...
|
|
123
|
+
|
|
124
|
+
Args:
|
|
125
|
+
client: HTTP client with authentication headers configured
|
|
126
|
+
service_type: Optional service type to filter by (e.g., "jira", "confluence").
|
|
127
|
+
If provided, will prioritize resources matching this service type.
|
|
128
|
+
|
|
129
|
+
Returns
|
|
130
|
+
-------
|
|
131
|
+
Cloud ID string for the Atlassian instance
|
|
132
|
+
|
|
133
|
+
Raises
|
|
134
|
+
------
|
|
135
|
+
ValueError: If cloud ID cannot be retrieved due to:
|
|
136
|
+
- No accessible resources found
|
|
137
|
+
- No cloud ID found in resources
|
|
138
|
+
- Authentication failure (401)
|
|
139
|
+
- HTTP request failure
|
|
140
|
+
|
|
141
|
+
Example:
|
|
142
|
+
```python
|
|
143
|
+
client = httpx.AsyncClient(headers={"Authorization": f"Bearer {token}"})
|
|
144
|
+
cloud_id = await get_atlassian_cloud_id(client, service_type="jira")
|
|
145
|
+
```
|
|
146
|
+
"""
|
|
147
|
+
url = f"{ATLASSIAN_API_BASE}{OAUTH_ACCESSIBLE_RESOURCES_PATH}"
|
|
148
|
+
|
|
149
|
+
try:
|
|
150
|
+
response = await client.get(url)
|
|
151
|
+
response.raise_for_status()
|
|
152
|
+
resources = response.json()
|
|
153
|
+
|
|
154
|
+
if not resources:
|
|
155
|
+
raise ValueError(
|
|
156
|
+
"No accessible resources found. Ensure OAuth token has required scopes."
|
|
157
|
+
)
|
|
158
|
+
|
|
159
|
+
# If service_type is specified, try to find matching resource
|
|
160
|
+
if service_type:
|
|
161
|
+
resource = _find_resource_by_service(resources, service_type)
|
|
162
|
+
if resource:
|
|
163
|
+
cloud_id = resource["id"]
|
|
164
|
+
logger.debug(f"Using {service_type} cloud ID: {cloud_id}")
|
|
165
|
+
return cloud_id
|
|
166
|
+
logger.warning(
|
|
167
|
+
f"No {service_type} resource found, falling back to first available resource"
|
|
168
|
+
)
|
|
169
|
+
|
|
170
|
+
# Fallback: use the first resource with an ID
|
|
171
|
+
resource = _find_first_resource_with_id(resources)
|
|
172
|
+
if resource:
|
|
173
|
+
cloud_id = resource["id"]
|
|
174
|
+
logger.debug(f"Using cloud ID (fallback): {cloud_id}")
|
|
175
|
+
return cloud_id
|
|
176
|
+
|
|
177
|
+
raise ValueError("No cloud ID found in accessible resources")
|
|
178
|
+
except httpx.HTTPStatusError as e:
|
|
179
|
+
if e.response.status_code == 401:
|
|
180
|
+
raise ValueError(
|
|
181
|
+
"Authentication failed. Token may be expired. "
|
|
182
|
+
"Complete OAuth flow again: GET /oauth/atlassian/authorize"
|
|
183
|
+
) from e
|
|
184
|
+
logger.error(f"HTTP error getting cloud ID: {e.response.status_code}")
|
|
185
|
+
raise ValueError(f"Failed to get Atlassian cloud ID: HTTP {e.response.status_code}") from e
|
|
186
|
+
except httpx.RequestError as e:
|
|
187
|
+
logger.error(f"Request error getting cloud ID: {e}")
|
|
188
|
+
raise ValueError("Failed to get Atlassian cloud ID: Network error") from e
|
|
@@ -0,0 +1,14 @@
|
|
|
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
|
+
|
|
@@ -0,0 +1,102 @@
|
|
|
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
|
+
import logging
|
|
16
|
+
from typing import Any
|
|
17
|
+
|
|
18
|
+
import httpx
|
|
19
|
+
|
|
20
|
+
from .atlassian import ATLASSIAN_API_BASE
|
|
21
|
+
from .atlassian import get_atlassian_cloud_id
|
|
22
|
+
|
|
23
|
+
logger = logging.getLogger(__name__)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class JiraClient:
|
|
27
|
+
"""Client for interacting with Jira API using OAuth access token."""
|
|
28
|
+
|
|
29
|
+
def __init__(self, access_token: str):
|
|
30
|
+
"""
|
|
31
|
+
Initialize Jira client with access token.
|
|
32
|
+
|
|
33
|
+
Args:
|
|
34
|
+
access_token: OAuth access token for Atlassian API
|
|
35
|
+
"""
|
|
36
|
+
self.access_token = access_token
|
|
37
|
+
self._client = httpx.AsyncClient(
|
|
38
|
+
headers={
|
|
39
|
+
"Authorization": f"Bearer {access_token}",
|
|
40
|
+
"Accept": "application/json",
|
|
41
|
+
"Content-Type": "application/json",
|
|
42
|
+
},
|
|
43
|
+
timeout=30.0,
|
|
44
|
+
)
|
|
45
|
+
self._cloud_id: str | None = None
|
|
46
|
+
|
|
47
|
+
async def _get_cloud_id(self) -> str:
|
|
48
|
+
"""
|
|
49
|
+
Get the cloud ID for the authenticated Atlassian Jira instance.
|
|
50
|
+
|
|
51
|
+
According to Atlassian OAuth 2.0 documentation, API calls should use:
|
|
52
|
+
https://api.atlassian.com/ex/jira/{cloudId}/rest/api/3/...
|
|
53
|
+
|
|
54
|
+
Returns
|
|
55
|
+
-------
|
|
56
|
+
Cloud ID string
|
|
57
|
+
|
|
58
|
+
Raises
|
|
59
|
+
------
|
|
60
|
+
ValueError: If cloud ID cannot be retrieved
|
|
61
|
+
"""
|
|
62
|
+
if self._cloud_id:
|
|
63
|
+
return self._cloud_id
|
|
64
|
+
|
|
65
|
+
self._cloud_id = await get_atlassian_cloud_id(self._client, service_type="jira")
|
|
66
|
+
return self._cloud_id
|
|
67
|
+
|
|
68
|
+
async def get_jira_issue(self, issue_key: str) -> dict[str, Any]:
|
|
69
|
+
"""
|
|
70
|
+
Get a Jira issue by its key.
|
|
71
|
+
|
|
72
|
+
Args:
|
|
73
|
+
issue_key: The key (ID) of the Jira issue, e.g., 'PROJ-123'
|
|
74
|
+
|
|
75
|
+
Returns
|
|
76
|
+
-------
|
|
77
|
+
Dictionary containing the issue data
|
|
78
|
+
|
|
79
|
+
Raises
|
|
80
|
+
------
|
|
81
|
+
httpx.HTTPStatusError: If the API request fails
|
|
82
|
+
"""
|
|
83
|
+
cloud_id = await self._get_cloud_id()
|
|
84
|
+
url = f"{ATLASSIAN_API_BASE}/ex/jira/{cloud_id}/rest/api/3/issue/{issue_key}"
|
|
85
|
+
|
|
86
|
+
response = await self._client.get(url)
|
|
87
|
+
response.raise_for_status()
|
|
88
|
+
return response.json()
|
|
89
|
+
|
|
90
|
+
async def close(self) -> None:
|
|
91
|
+
"""Close the HTTP client."""
|
|
92
|
+
await self._client.aclose()
|
|
93
|
+
|
|
94
|
+
async def __aenter__(self) -> "JiraClient":
|
|
95
|
+
"""Async context manager entry."""
|
|
96
|
+
return self
|
|
97
|
+
|
|
98
|
+
async def __aexit__(
|
|
99
|
+
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
|
|
100
|
+
) -> None:
|
|
101
|
+
"""Async context manager exit."""
|
|
102
|
+
await self.close()
|
|
@@ -0,0 +1,28 @@
|
|
|
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
|
+
import logging
|
|
16
|
+
|
|
17
|
+
from datarobot_genai.drmcp.core.credentials import get_credentials
|
|
18
|
+
|
|
19
|
+
logger = logging.getLogger(__name__)
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def get_s3_bucket_info() -> dict[str, str]:
|
|
23
|
+
"""Get S3 bucket configuration."""
|
|
24
|
+
credentials = get_credentials()
|
|
25
|
+
return {
|
|
26
|
+
"bucket": credentials.aws_predictions_s3_bucket,
|
|
27
|
+
"prefix": credentials.aws_predictions_s3_prefix,
|
|
28
|
+
}
|
|
@@ -0,0 +1,14 @@
|
|
|
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
|
+
|
|
@@ -0,0 +1,13 @@
|
|
|
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.
|
|
@@ -0,0 +1,14 @@
|
|
|
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
|
+
|
|
@@ -0,0 +1,58 @@
|
|
|
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
|
+
import logging
|
|
16
|
+
from typing import Annotated
|
|
17
|
+
|
|
18
|
+
from fastmcp.exceptions import ToolError
|
|
19
|
+
from fastmcp.tools.tool import ToolResult
|
|
20
|
+
|
|
21
|
+
from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
|
|
22
|
+
from datarobot_genai.drmcp.tools.clients.atlassian import get_atlassian_access_token
|
|
23
|
+
from datarobot_genai.drmcp.tools.clients.jira import JiraClient
|
|
24
|
+
|
|
25
|
+
logger = logging.getLogger(__name__)
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
@dr_mcp_tool(tags={"jira", "read", "get", "issue"})
|
|
29
|
+
async def jira_get_issue(
|
|
30
|
+
*, issue_key: Annotated[str, "The key (ID) of the Jira issue to retrieve, e.g., 'PROJ-123'."]
|
|
31
|
+
) -> ToolResult:
|
|
32
|
+
"""Retrieve all fields and details for a single Jira issue by its key."""
|
|
33
|
+
if not issue_key:
|
|
34
|
+
raise ToolError("Argument validation error: 'issue_key' cannot be empty.")
|
|
35
|
+
|
|
36
|
+
access_token = await get_atlassian_access_token()
|
|
37
|
+
if isinstance(access_token, ToolError):
|
|
38
|
+
return access_token
|
|
39
|
+
|
|
40
|
+
try:
|
|
41
|
+
client = JiraClient(
|
|
42
|
+
access_token,
|
|
43
|
+
)
|
|
44
|
+
issue = await client.get_jira_issue(issue_key)
|
|
45
|
+
except Exception as e:
|
|
46
|
+
logger.error(f"Unexpected error getting Jira issue: {e}")
|
|
47
|
+
return ToolError(
|
|
48
|
+
f"An unexpected error occurred while getting Jira issue '{issue_key}': {str(e)}"
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
return ToolResult(
|
|
52
|
+
content=f"Successfully retrieved details for issue '{issue_key}'.",
|
|
53
|
+
# TODO: Add more fields to the structured content, note fields here are just examples
|
|
54
|
+
structured_content={
|
|
55
|
+
"key": issue.get("key", issue_key),
|
|
56
|
+
"status": issue.get("fields", {}).get("status", {}).get("name", "Unknown"),
|
|
57
|
+
},
|
|
58
|
+
)
|
|
@@ -22,10 +22,10 @@ from fastmcp.resources import HttpResource
|
|
|
22
22
|
from fastmcp.resources import ResourceManager
|
|
23
23
|
|
|
24
24
|
from datarobot_genai.drmcp.core.clients import get_credentials
|
|
25
|
-
from datarobot_genai.drmcp.core.clients import get_s3_bucket_info
|
|
26
25
|
from datarobot_genai.drmcp.core.clients import get_sdk_client
|
|
27
26
|
from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
|
|
28
27
|
from datarobot_genai.drmcp.core.utils import generate_presigned_url
|
|
28
|
+
from datarobot_genai.drmcp.tools.clients.s3 import get_s3_bucket_info
|
|
29
29
|
|
|
30
30
|
logger = logging.getLogger(__name__)
|
|
31
31
|
|
|
@@ -23,11 +23,11 @@ from datarobot_predict import TimeSeriesType
|
|
|
23
23
|
from datarobot_predict.deployment import predict as dr_predict
|
|
24
24
|
from pydantic import BaseModel
|
|
25
25
|
|
|
26
|
-
from datarobot_genai.drmcp.core.clients import get_s3_bucket_info
|
|
27
26
|
from datarobot_genai.drmcp.core.clients import get_sdk_client
|
|
28
27
|
from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
|
|
29
28
|
from datarobot_genai.drmcp.core.utils import PredictionResponse
|
|
30
29
|
from datarobot_genai.drmcp.core.utils import predictions_result_response
|
|
30
|
+
from datarobot_genai.drmcp.tools.clients.s3 import get_s3_bucket_info
|
|
31
31
|
|
|
32
32
|
logger = logging.getLogger(__name__)
|
|
33
33
|
|
datarobot_genai/nat/agent.py
CHANGED
|
@@ -123,6 +123,8 @@ class NatAgent(BaseAgent[None]):
|
|
|
123
123
|
model: str | None = None,
|
|
124
124
|
verbose: bool | str | None = True,
|
|
125
125
|
timeout: int | None = 90,
|
|
126
|
+
authorization_context: dict[str, Any] | None = None,
|
|
127
|
+
forwarded_headers: dict[str, str] | None = None,
|
|
126
128
|
**kwargs: Any,
|
|
127
129
|
) -> None:
|
|
128
130
|
super().__init__(
|
|
@@ -131,6 +133,8 @@ class NatAgent(BaseAgent[None]):
|
|
|
131
133
|
model=model,
|
|
132
134
|
verbose=verbose,
|
|
133
135
|
timeout=timeout,
|
|
136
|
+
authorization_context=authorization_context,
|
|
137
|
+
forwarded_headers=forwarded_headers,
|
|
134
138
|
**kwargs,
|
|
135
139
|
)
|
|
136
140
|
self.workflow_path = workflow_path
|
|
@@ -0,0 +1,53 @@
|
|
|
1
|
+
# Copyright 2025 DataRobot, Inc. and its affiliates.
|
|
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
|
+
from datarobot.core.config import DataRobotAppFrameworkBaseSettings
|
|
15
|
+
from nat.authentication.api_key.api_key_auth_provider import APIKeyAuthProvider
|
|
16
|
+
from nat.authentication.api_key.api_key_auth_provider_config import APIKeyAuthProviderConfig
|
|
17
|
+
from nat.builder.builder import Builder
|
|
18
|
+
from nat.cli.register_workflow import register_auth_provider
|
|
19
|
+
from pydantic import Field
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class Config(DataRobotAppFrameworkBaseSettings):
|
|
23
|
+
"""
|
|
24
|
+
Finds variables in the priority order of: env
|
|
25
|
+
variables (including Runtime Parameters), .env, file_secrets, then
|
|
26
|
+
Pulumi output variables.
|
|
27
|
+
"""
|
|
28
|
+
|
|
29
|
+
datarobot_api_token: str | None = None
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
config = Config()
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class DataRobotAPIKeyAuthProviderConfig(APIKeyAuthProviderConfig, name="datarobot_api_key"): # type: ignore[call-arg]
|
|
36
|
+
raw_key: str = Field(
|
|
37
|
+
description=(
|
|
38
|
+
"Raw API token or credential to be injected into the request parameter. "
|
|
39
|
+
"Used for 'bearer','x-api-key','custom', and other schemes. "
|
|
40
|
+
),
|
|
41
|
+
default=config.datarobot_api_token,
|
|
42
|
+
)
|
|
43
|
+
default_user_id: str | None = Field(default="default-user", description="Default user ID")
|
|
44
|
+
allow_default_user_id_for_tool_calls: bool = Field(
|
|
45
|
+
default=True, description="Allow default user ID for tool calls"
|
|
46
|
+
)
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
@register_auth_provider(config_type=DataRobotAPIKeyAuthProviderConfig)
|
|
50
|
+
async def datarobot_api_key_client(
|
|
51
|
+
config: DataRobotAPIKeyAuthProviderConfig, builder: Builder
|
|
52
|
+
) -> APIKeyAuthProviderConfig:
|
|
53
|
+
yield APIKeyAuthProvider(config=config)
|
|
@@ -0,0 +1,181 @@
|
|
|
1
|
+
# Copyright 2025 DataRobot, Inc. and its affiliates.
|
|
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
|
+
import logging
|
|
16
|
+
from typing import Literal
|
|
17
|
+
|
|
18
|
+
from nat.builder.builder import Builder
|
|
19
|
+
from nat.cli.register_workflow import register_function_group
|
|
20
|
+
from nat.data_models.component_ref import AuthenticationRef
|
|
21
|
+
from nat.plugins.mcp.client_base import MCPSSEClient
|
|
22
|
+
from nat.plugins.mcp.client_base import MCPStdioClient
|
|
23
|
+
from nat.plugins.mcp.client_base import MCPStreamableHTTPClient
|
|
24
|
+
from nat.plugins.mcp.client_config import MCPServerConfig
|
|
25
|
+
from nat.plugins.mcp.client_impl import MCPClientConfig
|
|
26
|
+
from nat.plugins.mcp.client_impl import MCPFunctionGroup
|
|
27
|
+
from nat.plugins.mcp.client_impl import mcp_apply_tool_alias_and_description
|
|
28
|
+
from nat.plugins.mcp.client_impl import mcp_session_tool_function
|
|
29
|
+
from pydantic import Field
|
|
30
|
+
from pydantic import HttpUrl
|
|
31
|
+
|
|
32
|
+
from datarobot_genai.core.mcp.common import MCPConfig
|
|
33
|
+
|
|
34
|
+
logger = logging.getLogger(__name__)
|
|
35
|
+
|
|
36
|
+
config = MCPConfig().server_config
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
class DataRobotMCPServerConfig(MCPServerConfig):
|
|
40
|
+
transport: Literal["streamable-http", "sse", "stdio"] = Field(
|
|
41
|
+
default=config["transport"] if config else "stdio",
|
|
42
|
+
description="Transport type to connect to the MCP server (sse or streamable-http)",
|
|
43
|
+
)
|
|
44
|
+
url: HttpUrl | None = Field(
|
|
45
|
+
default=config["url"] if config else None,
|
|
46
|
+
description="URL of the MCP server (for sse or streamable-http transport)",
|
|
47
|
+
)
|
|
48
|
+
# Authentication configuration
|
|
49
|
+
auth_provider: str | AuthenticationRef | None = Field(
|
|
50
|
+
default="datarobot_api_key" if config else None,
|
|
51
|
+
description="Reference to authentication provider",
|
|
52
|
+
)
|
|
53
|
+
command: str | None = Field(
|
|
54
|
+
default=None if config else "docker",
|
|
55
|
+
description="Command to run for stdio transport (e.g. 'python' or 'docker')",
|
|
56
|
+
)
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
class DataRobotMCPClientConfig(MCPClientConfig, name="datarobot_mcp_client"): # type: ignore[call-arg]
|
|
60
|
+
server: DataRobotMCPServerConfig = Field(
|
|
61
|
+
default=DataRobotMCPServerConfig(), description="DataRobot MCP Server configuration"
|
|
62
|
+
)
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
@register_function_group(config_type=DataRobotMCPClientConfig)
|
|
66
|
+
async def datarobot_mcp_client_function_group(
|
|
67
|
+
config: DataRobotMCPClientConfig, _builder: Builder
|
|
68
|
+
) -> MCPFunctionGroup:
|
|
69
|
+
"""
|
|
70
|
+
Connect to an MCP server and expose tools as a function group.
|
|
71
|
+
|
|
72
|
+
Args:
|
|
73
|
+
config: The configuration for the MCP client
|
|
74
|
+
_builder: The builder
|
|
75
|
+
Returns:
|
|
76
|
+
The function group
|
|
77
|
+
"""
|
|
78
|
+
# Resolve auth provider if specified
|
|
79
|
+
auth_provider = None
|
|
80
|
+
if config.server.auth_provider:
|
|
81
|
+
auth_provider = await _builder.get_auth_provider(config.server.auth_provider)
|
|
82
|
+
|
|
83
|
+
# Build the appropriate client
|
|
84
|
+
if config.server.transport == "stdio":
|
|
85
|
+
if not config.server.command:
|
|
86
|
+
raise ValueError("command is required for stdio transport")
|
|
87
|
+
client = MCPStdioClient(
|
|
88
|
+
config.server.command,
|
|
89
|
+
config.server.args,
|
|
90
|
+
config.server.env,
|
|
91
|
+
tool_call_timeout=config.tool_call_timeout,
|
|
92
|
+
auth_flow_timeout=config.auth_flow_timeout,
|
|
93
|
+
reconnect_enabled=config.reconnect_enabled,
|
|
94
|
+
reconnect_max_attempts=config.reconnect_max_attempts,
|
|
95
|
+
reconnect_initial_backoff=config.reconnect_initial_backoff,
|
|
96
|
+
reconnect_max_backoff=config.reconnect_max_backoff,
|
|
97
|
+
)
|
|
98
|
+
elif config.server.transport == "sse":
|
|
99
|
+
client = MCPSSEClient(
|
|
100
|
+
str(config.server.url),
|
|
101
|
+
tool_call_timeout=config.tool_call_timeout,
|
|
102
|
+
auth_flow_timeout=config.auth_flow_timeout,
|
|
103
|
+
reconnect_enabled=config.reconnect_enabled,
|
|
104
|
+
reconnect_max_attempts=config.reconnect_max_attempts,
|
|
105
|
+
reconnect_initial_backoff=config.reconnect_initial_backoff,
|
|
106
|
+
reconnect_max_backoff=config.reconnect_max_backoff,
|
|
107
|
+
)
|
|
108
|
+
elif config.server.transport == "streamable-http":
|
|
109
|
+
# Use default_user_id for the base client
|
|
110
|
+
base_user_id = auth_provider.config.default_user_id if auth_provider else None
|
|
111
|
+
client = MCPStreamableHTTPClient(
|
|
112
|
+
str(config.server.url),
|
|
113
|
+
auth_provider=auth_provider,
|
|
114
|
+
user_id=base_user_id,
|
|
115
|
+
tool_call_timeout=config.tool_call_timeout,
|
|
116
|
+
auth_flow_timeout=config.auth_flow_timeout,
|
|
117
|
+
reconnect_enabled=config.reconnect_enabled,
|
|
118
|
+
reconnect_max_attempts=config.reconnect_max_attempts,
|
|
119
|
+
reconnect_initial_backoff=config.reconnect_initial_backoff,
|
|
120
|
+
reconnect_max_backoff=config.reconnect_max_backoff,
|
|
121
|
+
)
|
|
122
|
+
else:
|
|
123
|
+
raise ValueError(f"Unsupported transport: {config.server.transport}")
|
|
124
|
+
|
|
125
|
+
logger.info("Configured to use MCP server at %s", client.server_name)
|
|
126
|
+
|
|
127
|
+
# Create the MCP function group
|
|
128
|
+
group = MCPFunctionGroup(config=config)
|
|
129
|
+
|
|
130
|
+
# Store shared components for session client creation
|
|
131
|
+
group._shared_auth_provider = auth_provider
|
|
132
|
+
group._client_config = config
|
|
133
|
+
|
|
134
|
+
async with client:
|
|
135
|
+
# Expose the live MCP client on the function group instance so other components
|
|
136
|
+
# (e.g., HTTP endpoints) can reuse the already-established session instead of creating a
|
|
137
|
+
# new client per request.
|
|
138
|
+
group.mcp_client = client
|
|
139
|
+
group.mcp_client_server_name = client.server_name
|
|
140
|
+
group.mcp_client_transport = client.transport
|
|
141
|
+
|
|
142
|
+
all_tools = await client.get_tools()
|
|
143
|
+
tool_overrides = mcp_apply_tool_alias_and_description(all_tools, config.tool_overrides)
|
|
144
|
+
|
|
145
|
+
# Add each tool as a function to the group
|
|
146
|
+
for tool_name, tool in all_tools.items():
|
|
147
|
+
# Get override if it exists
|
|
148
|
+
override = tool_overrides.get(tool_name)
|
|
149
|
+
|
|
150
|
+
# Use override values or defaults
|
|
151
|
+
function_name = override.alias if override and override.alias else tool_name
|
|
152
|
+
description = (
|
|
153
|
+
override.description if override and override.description else tool.description
|
|
154
|
+
)
|
|
155
|
+
|
|
156
|
+
# Create the tool function according to configuration
|
|
157
|
+
tool_fn = mcp_session_tool_function(tool, group)
|
|
158
|
+
|
|
159
|
+
# Normalize optional typing for linter/type-checker compatibility
|
|
160
|
+
single_fn = tool_fn.single_fn
|
|
161
|
+
if single_fn is None:
|
|
162
|
+
# Should not happen because FunctionInfo always sets a single_fn
|
|
163
|
+
logger.warning("Skipping tool %s because single_fn is None", function_name)
|
|
164
|
+
continue
|
|
165
|
+
|
|
166
|
+
input_schema = tool_fn.input_schema
|
|
167
|
+
# Convert NoneType sentinel to None for FunctionGroup.add_function signature
|
|
168
|
+
if input_schema is type(None): # noqa: E721
|
|
169
|
+
input_schema = None
|
|
170
|
+
|
|
171
|
+
# Add to group
|
|
172
|
+
logger.info("Adding tool %s to group", function_name)
|
|
173
|
+
group.add_function(
|
|
174
|
+
name=function_name,
|
|
175
|
+
description=description,
|
|
176
|
+
fn=single_fn,
|
|
177
|
+
input_schema=input_schema,
|
|
178
|
+
converters=tool_fn.converters,
|
|
179
|
+
)
|
|
180
|
+
|
|
181
|
+
yield group
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: datarobot-genai
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.6
|
|
4
4
|
Summary: Generic helpers for GenAI
|
|
5
5
|
Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
|
|
6
6
|
Author: DataRobot, Inc.
|
|
@@ -63,6 +63,7 @@ Requires-Dist: anyio==4.11.0; extra == 'nat'
|
|
|
63
63
|
Requires-Dist: crewai>=1.1.0; (python_version >= '3.11') and extra == 'nat'
|
|
64
64
|
Requires-Dist: llama-index-llms-litellm<0.7.0,>=0.4.1; extra == 'nat'
|
|
65
65
|
Requires-Dist: nvidia-nat-langchain==1.3.0; (python_version >= '3.11') and extra == 'nat'
|
|
66
|
+
Requires-Dist: nvidia-nat-mcp==1.3.0; (python_version >= '3.11') and extra == 'nat'
|
|
66
67
|
Requires-Dist: nvidia-nat-opentelemetry==1.3.0; (python_version >= '3.11') and extra == 'nat'
|
|
67
68
|
Requires-Dist: nvidia-nat==1.3.0; (python_version >= '3.11') and extra == 'nat'
|
|
68
69
|
Requires-Dist: opentelemetry-instrumentation-crewai<1.0.0,>=0.40.5; extra == 'nat'
|
|
@@ -27,11 +27,11 @@ 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=0Y0OjbLq-buP66U9tiXMyqrX9wWIn557vKvkBvjM_cM,11797
|
|
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
|
|
34
|
-
datarobot_genai/drmcp/core/dr_mcp_server.py,sha256=
|
|
34
|
+
datarobot_genai/drmcp/core/dr_mcp_server.py,sha256=yKaIe7Qq23Zgny7Q1dc48iEcpfM8z2Ne6iouGYL58AE,14392
|
|
35
35
|
datarobot_genai/drmcp/core/dr_mcp_server_logo.py,sha256=hib-nfR1SNTW6CnpFsFCkL9H_OMwa4YYyinV7VNOuLk,4708
|
|
36
36
|
datarobot_genai/drmcp/core/exceptions.py,sha256=eqsGI-lxybgvWL5w4BFhbm3XzH1eU5tetwjnhJxelpc,905
|
|
37
37
|
datarobot_genai/drmcp/core/logging.py,sha256=Y_hig4eBWiXGaVV7B_3wBcaYVRNH4ydptbEQhrP9-mY,3414
|
|
@@ -41,6 +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=5OkC3e-vekZtdqg-DbwcyadGSrDxmZqSDey2YyGVn1M,2978
|
|
44
45
|
datarobot_genai/drmcp/core/tool_filter.py,sha256=tLOcG50QBvS48cOVHM6OqoODYiiS6KeM_F-2diaHkW0,2858
|
|
45
46
|
datarobot_genai/drmcp/core/utils.py,sha256=dSjrayWVcnC5GxQcvOIOSHaoEymPIVtG_s2ZBMlmSOw,4336
|
|
46
47
|
datarobot_genai/drmcp/core/dynamic_prompts/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
|
|
@@ -73,13 +74,22 @@ datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py,sha256=Va3_5c2ToZyfIsE
|
|
|
73
74
|
datarobot_genai/drmcp/test_utils/tool_base_ete.py,sha256=-mKHBkGkyOKQCVS2LHFhSnRofIqJBbeAPRkwizBDtTg,6104
|
|
74
75
|
datarobot_genai/drmcp/test_utils/utils.py,sha256=esGKFv8aO31-Qg3owayeWp32BYe1CdYOEutjjdbweCw,3048
|
|
75
76
|
datarobot_genai/drmcp/tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
|
|
77
|
+
datarobot_genai/drmcp/tools/clients/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
|
|
78
|
+
datarobot_genai/drmcp/tools/clients/atlassian.py,sha256=__M_uz7FrcbKCYRzeMn24DCEYD6OmFx_LuywHCxgXsA,6472
|
|
79
|
+
datarobot_genai/drmcp/tools/clients/confluence.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
|
|
80
|
+
datarobot_genai/drmcp/tools/clients/jira.py,sha256=JjvssdMAWgZ3HWZkQg0a3HjpE7yz7jfRtzO4LOp47Uw,3080
|
|
81
|
+
datarobot_genai/drmcp/tools/clients/s3.py,sha256=GmwzvurFdNfvxOooA8g5S4osRysHYU0S9ypg_177Glg,953
|
|
82
|
+
datarobot_genai/drmcp/tools/confluence/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
|
|
83
|
+
datarobot_genai/drmcp/tools/confluence/tools.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
|
|
84
|
+
datarobot_genai/drmcp/tools/jira/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
|
|
85
|
+
datarobot_genai/drmcp/tools/jira/tools.py,sha256=EBf8T_VwKXNticorRN8ut9ktEMS0XFeb0Uj_6TKwMio,2191
|
|
76
86
|
datarobot_genai/drmcp/tools/predictive/__init__.py,sha256=WuOHlNNEpEmcF7gVnhckruJRKU2qtmJLE3E7zoCGLDo,1030
|
|
77
87
|
datarobot_genai/drmcp/tools/predictive/data.py,sha256=k4EJxJrl8DYVGVfJ0DM4YTfnZlC_K3OUHZ0eRUzfluI,3165
|
|
78
88
|
datarobot_genai/drmcp/tools/predictive/deployment.py,sha256=lm02Ayuo11L1hP41fgi3QpR1Eyty-Wc16rM0c8SgliM,3277
|
|
79
89
|
datarobot_genai/drmcp/tools/predictive/deployment_info.py,sha256=BGEF_dmbxOBJR0n1Tt9TO2-iNTQSBTr-oQUyaxLZ0ZI,15297
|
|
80
90
|
datarobot_genai/drmcp/tools/predictive/model.py,sha256=Yih5-KedJ-1yupPLXCJsCXOdyWWi9pRvgapXDlgXWJA,4891
|
|
81
|
-
datarobot_genai/drmcp/tools/predictive/predict.py,sha256=
|
|
82
|
-
datarobot_genai/drmcp/tools/predictive/predict_realtime.py,sha256=
|
|
91
|
+
datarobot_genai/drmcp/tools/predictive/predict.py,sha256=Qoob2_t2crfWtyPzkXMRz2ITZumnczU6Dq4C7q9RBMI,9370
|
|
92
|
+
datarobot_genai/drmcp/tools/predictive/predict_realtime.py,sha256=urq6rPyZFsAP-bPyclSNzrkvb6FTamdlFau8q0IWWJ0,13472
|
|
83
93
|
datarobot_genai/drmcp/tools/predictive/project.py,sha256=KaMDAvJY4s12j_4ybA7-KcCS1yMOj-KPIKNBgCSE2iM,2536
|
|
84
94
|
datarobot_genai/drmcp/tools/predictive/training.py,sha256=kxeDVLqUh9ajDk8wK7CZRRydDK8UNuTVZCB3huUihF8,23660
|
|
85
95
|
datarobot_genai/langgraph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
@@ -90,12 +100,14 @@ datarobot_genai/llama_index/agent.py,sha256=V6ZsD9GcBDJS-RJo1tJtIHhyW69_78gM6_fO
|
|
|
90
100
|
datarobot_genai/llama_index/base.py,sha256=ovcQQtC-djD_hcLrWdn93jg23AmD6NBEj7xtw4a6K6c,14481
|
|
91
101
|
datarobot_genai/llama_index/mcp.py,sha256=leXqF1C4zhuYEKFwNEfZHY4dsUuGZk3W7KArY-zxVL8,2645
|
|
92
102
|
datarobot_genai/nat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
|
93
|
-
datarobot_genai/nat/agent.py,sha256=
|
|
103
|
+
datarobot_genai/nat/agent.py,sha256=jDeIS9f-8vGbeLy5gQkSjeuHINx5Fh_4BvXYERsgIIk,10516
|
|
104
|
+
datarobot_genai/nat/datarobot_auth_provider.py,sha256=5SBs0xBEZAlBvK-_zOR2cfE_rnn2CoEXeb-Du-rqotc,2117
|
|
94
105
|
datarobot_genai/nat/datarobot_llm_clients.py,sha256=STzAZ4OF8U-Y_cUTywxmKBGVotwsnbGP6vTojnu6q0g,9921
|
|
95
106
|
datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
|
|
96
|
-
datarobot_genai
|
|
97
|
-
datarobot_genai-0.2.
|
|
98
|
-
datarobot_genai-0.2.
|
|
99
|
-
datarobot_genai-0.2.
|
|
100
|
-
datarobot_genai-0.2.
|
|
101
|
-
datarobot_genai-0.2.
|
|
107
|
+
datarobot_genai/nat/datarobot_mcp_client.py,sha256=XK3-Tp4hdQmTdI-Zwrl3Xd81qH5RhFTw_UaLTqwnYn4,7531
|
|
108
|
+
datarobot_genai-0.2.6.dist-info/METADATA,sha256=2F0MD9IDjdYCmuhCO76dJHljxSk0cJP0rX17TJQoPaQ,6172
|
|
109
|
+
datarobot_genai-0.2.6.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
|
|
110
|
+
datarobot_genai-0.2.6.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
|
|
111
|
+
datarobot_genai-0.2.6.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
|
|
112
|
+
datarobot_genai-0.2.6.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
|
|
113
|
+
datarobot_genai-0.2.6.dist-info/RECORD,,
|
|
@@ -1,3 +1,5 @@
|
|
|
1
1
|
[nat.plugins]
|
|
2
|
+
datarobot_auth_provider = datarobot_genai.nat.datarobot_auth_provider
|
|
2
3
|
datarobot_llm_clients = datarobot_genai.nat.datarobot_llm_clients
|
|
3
4
|
datarobot_llm_providers = datarobot_genai.nat.datarobot_llm_providers
|
|
5
|
+
datarobot_mcp_client = datarobot_genai.nat.datarobot_mcp_client
|
|
File without changes
|
|
File without changes
|
|
File without changes
|