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.
@@ -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
- if self._config.enable_predictive_tools:
146
- _import_modules_from_dir(
147
- os.path.join(base_dir, "tools", "predictive"),
148
- "datarobot_genai.drmcp.tools.predictive",
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
 
@@ -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.4
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=D7bSi40Yc5J71_JxmpfppG83snbIJW9iz1J7qbiJrRs,9855
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=peAq0TL4ZL0P6XJjwCEWNfi0OVPPSQMSRQyKFH6yjtY,14228
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=7h73VidPJ4nlxtdFWRQ_S1epDvYuXCypKT8DXzDl4-g,9366
82
- datarobot_genai/drmcp/tools/predictive/predict_realtime.py,sha256=t7f28y_ealZoA6itrCzlJbTc3KqEU0H-a41ah-7U0XI,13468
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=siBLDWAff2-JwZ8Q3iNpM_e4_IoSwG9IvY0hyEjNenw,10292
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-0.2.4.dist-info/METADATA,sha256=V_woRsjfbnIZPMWWB6LPPfcuXDcLb50elyugsB31S0w,6088
97
- datarobot_genai-0.2.4.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
98
- datarobot_genai-0.2.4.dist-info/entry_points.txt,sha256=CZhmZcSyt_RBltgLN_b9xasJD6J5SaDc_z7K0wuOY9Y,150
99
- datarobot_genai-0.2.4.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
100
- datarobot_genai-0.2.4.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
101
- datarobot_genai-0.2.4.dist-info/RECORD,,
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