datarobot-genai 0.2.4__tar.gz → 0.2.6__tar.gz

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.
Files changed (113) hide show
  1. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/PKG-INFO +2 -1
  2. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/pyproject.toml +4 -1
  3. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/config.py +52 -0
  4. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dr_mcp_server.py +11 -8
  5. datarobot_genai-0.2.6/src/datarobot_genai/drmcp/core/tool_config.py +95 -0
  6. datarobot_genai-0.2.6/src/datarobot_genai/drmcp/tools/clients/__init__.py +14 -0
  7. datarobot_genai-0.2.6/src/datarobot_genai/drmcp/tools/clients/atlassian.py +188 -0
  8. datarobot_genai-0.2.6/src/datarobot_genai/drmcp/tools/clients/confluence.py +14 -0
  9. datarobot_genai-0.2.6/src/datarobot_genai/drmcp/tools/clients/jira.py +102 -0
  10. datarobot_genai-0.2.6/src/datarobot_genai/drmcp/tools/clients/s3.py +28 -0
  11. datarobot_genai-0.2.6/src/datarobot_genai/drmcp/tools/confluence/__init__.py +14 -0
  12. datarobot_genai-0.2.6/src/datarobot_genai/drmcp/tools/confluence/tools.py +13 -0
  13. datarobot_genai-0.2.6/src/datarobot_genai/drmcp/tools/jira/__init__.py +14 -0
  14. datarobot_genai-0.2.6/src/datarobot_genai/drmcp/tools/jira/tools.py +58 -0
  15. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/tools/predictive/predict.py +1 -1
  16. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/tools/predictive/predict_realtime.py +1 -1
  17. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/nat/agent.py +4 -0
  18. datarobot_genai-0.2.6/src/datarobot_genai/nat/datarobot_auth_provider.py +53 -0
  19. datarobot_genai-0.2.6/src/datarobot_genai/nat/datarobot_mcp_client.py +181 -0
  20. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/.gitignore +0 -0
  21. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/AUTHORS +0 -0
  22. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/LICENSE +0 -0
  23. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/README.md +0 -0
  24. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/__init__.py +0 -0
  25. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/__init__.py +0 -0
  26. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/agents/__init__.py +0 -0
  27. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/agents/base.py +0 -0
  28. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/chat/__init__.py +0 -0
  29. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/chat/auth.py +0 -0
  30. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/chat/client.py +0 -0
  31. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/chat/responses.py +0 -0
  32. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/cli/__init__.py +0 -0
  33. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/cli/agent_environment.py +0 -0
  34. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/cli/agent_kernel.py +0 -0
  35. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/custom_model.py +0 -0
  36. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/mcp/__init__.py +0 -0
  37. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/mcp/common.py +0 -0
  38. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/telemetry_agent.py +0 -0
  39. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/utils/__init__.py +0 -0
  40. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/utils/auth.py +0 -0
  41. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/core/utils/urls.py +0 -0
  42. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/crewai/__init__.py +0 -0
  43. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/crewai/agent.py +0 -0
  44. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/crewai/base.py +0 -0
  45. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/crewai/events.py +0 -0
  46. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/crewai/mcp.py +0 -0
  47. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/__init__.py +0 -0
  48. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/__init__.py +0 -0
  49. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/auth.py +0 -0
  50. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/clients.py +0 -0
  51. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/config_utils.py +0 -0
  52. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/constants.py +0 -0
  53. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/credentials.py +0 -0
  54. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dr_mcp_server_logo.py +0 -0
  55. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_prompts/__init__.py +0 -0
  56. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_prompts/controllers.py +0 -0
  57. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py +0 -0
  58. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_prompts/register.py +0 -0
  59. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_prompts/utils.py +0 -0
  60. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/__init__.py +0 -0
  61. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/__init__.py +0 -0
  62. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/__init__.py +0 -0
  63. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/base.py +0 -0
  64. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/default.py +0 -0
  65. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/drum.py +0 -0
  66. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/config.py +0 -0
  67. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/controllers.py +0 -0
  68. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/metadata.py +0 -0
  69. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/register.py +0 -0
  70. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_agentic_fallback_schema.json +0 -0
  71. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_prediction_fallback_schema.json +0 -0
  72. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/register.py +0 -0
  73. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/dynamic_tools/schema.py +0 -0
  74. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/exceptions.py +0 -0
  75. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/logging.py +0 -0
  76. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/mcp_instance.py +0 -0
  77. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/mcp_server_tools.py +0 -0
  78. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/memory_management/__init__.py +0 -0
  79. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/memory_management/manager.py +0 -0
  80. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/memory_management/memory_tools.py +0 -0
  81. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/routes.py +0 -0
  82. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/routes_utils.py +0 -0
  83. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/server_life_cycle.py +0 -0
  84. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/telemetry.py +0 -0
  85. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/tool_filter.py +0 -0
  86. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/core/utils.py +0 -0
  87. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/server.py +0 -0
  88. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/test_utils/__init__.py +0 -0
  89. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/test_utils/integration_mcp_server.py +0 -0
  90. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/test_utils/mcp_utils_ete.py +0 -0
  91. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/test_utils/mcp_utils_integration.py +0 -0
  92. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py +0 -0
  93. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/test_utils/tool_base_ete.py +0 -0
  94. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/test_utils/utils.py +0 -0
  95. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/tools/__init__.py +0 -0
  96. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/tools/predictive/__init__.py +0 -0
  97. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/tools/predictive/data.py +0 -0
  98. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/tools/predictive/deployment.py +0 -0
  99. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/tools/predictive/deployment_info.py +0 -0
  100. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/tools/predictive/model.py +0 -0
  101. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/tools/predictive/project.py +0 -0
  102. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/drmcp/tools/predictive/training.py +0 -0
  103. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/langgraph/__init__.py +0 -0
  104. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/langgraph/agent.py +0 -0
  105. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/langgraph/mcp.py +0 -0
  106. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/llama_index/__init__.py +0 -0
  107. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/llama_index/agent.py +0 -0
  108. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/llama_index/base.py +0 -0
  109. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/llama_index/mcp.py +0 -0
  110. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/nat/__init__.py +0 -0
  111. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/nat/datarobot_llm_clients.py +0 -0
  112. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/nat/datarobot_llm_providers.py +0 -0
  113. {datarobot_genai-0.2.4 → datarobot_genai-0.2.6}/src/datarobot_genai/py.typed +0 -0
@@ -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'
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "datarobot-genai"
7
- version = "0.2.4"
7
+ version = "0.2.6"
8
8
  description = "Generic helpers for GenAI"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10, <3.13"
@@ -35,6 +35,8 @@ Homepage = "https://github.com/datarobot-oss/datarobot-genai"
35
35
  [project.entry-points.'nat.plugins']
36
36
  datarobot_llm_providers = "datarobot_genai.nat.datarobot_llm_providers"
37
37
  datarobot_llm_clients = "datarobot_genai.nat.datarobot_llm_clients"
38
+ datarobot_mcp_client = "datarobot_genai.nat.datarobot_mcp_client"
39
+ datarobot_auth_provider = "datarobot_genai.nat.datarobot_auth_provider"
38
40
 
39
41
  [project.optional-dependencies]
40
42
  crewai = [
@@ -63,6 +65,7 @@ nat = [
63
65
  "nvidia-nat==1.3.0; python_version >= '3.11'",
64
66
  "nvidia-nat-opentelemetry==1.3.0; python_version >= '3.11'",
65
67
  "nvidia-nat-langchain==1.3.0; python_version >= '3.11'",
68
+ "nvidia-nat-mcp==1.3.0; python_version >= '3.11'",
66
69
  "crewai>=1.1.0; python_version >= '3.11'",
67
70
  "llama-index-llms-litellm>=0.4.1,<0.7.0", # Need this to support datarobot-llm plugin
68
71
  "opentelemetry-instrumentation-crewai>=0.40.5,<1.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
- 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
+