datarobot-genai 0.2.26__tar.gz → 0.2.39__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 (128) hide show
  1. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/PKG-INFO +1 -1
  2. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/pyproject.toml +1 -1
  3. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/cli/agent_kernel.py +4 -1
  4. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/__init__.py +2 -2
  5. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/config.py +121 -83
  6. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/exceptions.py +0 -4
  7. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/logging.py +2 -2
  8. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/tool_config.py +17 -9
  9. datarobot_genai-0.2.39/src/datarobot_genai/drmcp/test_utils/clients/anthropic.py +68 -0
  10. datarobot_genai-0.2.26/src/datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py → datarobot_genai-0.2.39/src/datarobot_genai/drmcp/test_utils/clients/base.py +38 -40
  11. datarobot_genai-0.2.39/src/datarobot_genai/drmcp/test_utils/clients/dr_gateway.py +58 -0
  12. datarobot_genai-0.2.39/src/datarobot_genai/drmcp/test_utils/clients/openai.py +68 -0
  13. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/test_utils/mcp_utils_ete.py +20 -0
  14. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/test_utils/test_interactive.py +16 -16
  15. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/test_utils/tool_base_ete.py +69 -2
  16. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/test_utils/utils.py +1 -1
  17. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/clients/gdrive.py +314 -1
  18. datarobot_genai-0.2.39/src/datarobot_genai/drmcp/tools/clients/microsoft_graph.py +650 -0
  19. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/confluence/tools.py +43 -89
  20. datarobot_genai-0.2.39/src/datarobot_genai/drmcp/tools/gdrive/tools.py +404 -0
  21. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/jira/tools.py +15 -33
  22. datarobot_genai-0.2.39/src/datarobot_genai/drmcp/tools/microsoft_graph/__init__.py +13 -0
  23. datarobot_genai-0.2.39/src/datarobot_genai/drmcp/tools/microsoft_graph/tools.py +330 -0
  24. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/predictive/data.py +16 -8
  25. datarobot_genai-0.2.39/src/datarobot_genai/drmcp/tools/predictive/deployment.py +97 -0
  26. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/predictive/deployment_info.py +100 -107
  27. datarobot_genai-0.2.39/src/datarobot_genai/drmcp/tools/predictive/model.py +183 -0
  28. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/predictive/project.py +2 -2
  29. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/predictive/training.py +53 -24
  30. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/nat/datarobot_llm_clients.py +90 -54
  31. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/nat/datarobot_mcp_client.py +47 -15
  32. datarobot_genai-0.2.39/src/datarobot_genai/py.typed +0 -0
  33. datarobot_genai-0.2.26/src/datarobot_genai/drmcp/tools/gdrive/tools.py +0 -177
  34. datarobot_genai-0.2.26/src/datarobot_genai/drmcp/tools/predictive/deployment.py +0 -91
  35. datarobot_genai-0.2.26/src/datarobot_genai/drmcp/tools/predictive/model.py +0 -148
  36. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/.gitignore +0 -0
  37. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/AUTHORS +0 -0
  38. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/LICENSE +0 -0
  39. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/README.md +0 -0
  40. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/__init__.py +0 -0
  41. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/__init__.py +0 -0
  42. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/agents/__init__.py +0 -0
  43. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/agents/base.py +0 -0
  44. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/chat/__init__.py +0 -0
  45. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/chat/auth.py +0 -0
  46. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/chat/client.py +0 -0
  47. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/chat/responses.py +0 -0
  48. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/cli/__init__.py +0 -0
  49. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/cli/agent_environment.py +0 -0
  50. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/custom_model.py +0 -0
  51. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/mcp/__init__.py +0 -0
  52. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/mcp/common.py +0 -0
  53. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/telemetry_agent.py +0 -0
  54. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/utils/__init__.py +0 -0
  55. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/utils/auth.py +0 -0
  56. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/core/utils/urls.py +0 -0
  57. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/crewai/__init__.py +0 -0
  58. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/crewai/agent.py +0 -0
  59. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/crewai/base.py +0 -0
  60. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/crewai/events.py +0 -0
  61. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/crewai/mcp.py +0 -0
  62. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/__init__.py +0 -0
  63. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/auth.py +0 -0
  64. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/clients.py +0 -0
  65. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/config_utils.py +0 -0
  66. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/constants.py +0 -0
  67. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/credentials.py +0 -0
  68. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dr_mcp_server.py +0 -0
  69. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dr_mcp_server_logo.py +0 -0
  70. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_prompts/__init__.py +0 -0
  71. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_prompts/controllers.py +0 -0
  72. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py +0 -0
  73. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_prompts/register.py +0 -0
  74. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_prompts/utils.py +0 -0
  75. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/__init__.py +0 -0
  76. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/__init__.py +0 -0
  77. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/__init__.py +0 -0
  78. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/base.py +0 -0
  79. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/default.py +0 -0
  80. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/drum.py +0 -0
  81. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/config.py +0 -0
  82. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/controllers.py +0 -0
  83. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/metadata.py +0 -0
  84. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/register.py +0 -0
  85. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_agentic_fallback_schema.json +0 -0
  86. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_prediction_fallback_schema.json +0 -0
  87. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/register.py +0 -0
  88. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/dynamic_tools/schema.py +0 -0
  89. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/mcp_instance.py +0 -0
  90. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/memory_management/__init__.py +0 -0
  91. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/memory_management/manager.py +0 -0
  92. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/memory_management/memory_tools.py +0 -0
  93. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/routes.py +0 -0
  94. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/routes_utils.py +0 -0
  95. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/server_life_cycle.py +0 -0
  96. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/telemetry.py +0 -0
  97. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/tool_filter.py +0 -0
  98. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/core/utils.py +0 -0
  99. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/server.py +0 -0
  100. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/test_utils/__init__.py +0 -0
  101. {datarobot_genai-0.2.26/src/datarobot_genai/drmcp/tools/gdrive → datarobot_genai-0.2.39/src/datarobot_genai/drmcp/test_utils/clients}/__init__.py +0 -0
  102. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/test_utils/elicitation_test_tool.py +0 -0
  103. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/test_utils/integration_mcp_server.py +0 -0
  104. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/test_utils/mcp_utils_integration.py +0 -0
  105. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/__init__.py +0 -0
  106. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/clients/__init__.py +0 -0
  107. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/clients/atlassian.py +0 -0
  108. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/clients/confluence.py +0 -0
  109. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/clients/jira.py +0 -0
  110. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/clients/s3.py +0 -0
  111. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/confluence/__init__.py +0 -0
  112. {datarobot_genai-0.2.26/src/datarobot_genai/langgraph → datarobot_genai-0.2.39/src/datarobot_genai/drmcp/tools/gdrive}/__init__.py +0 -0
  113. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/jira/__init__.py +0 -0
  114. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/predictive/__init__.py +0 -0
  115. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/predictive/predict.py +0 -0
  116. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/drmcp/tools/predictive/predict_realtime.py +0 -0
  117. {datarobot_genai-0.2.26/src/datarobot_genai/nat → datarobot_genai-0.2.39/src/datarobot_genai/langgraph}/__init__.py +0 -0
  118. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/langgraph/agent.py +0 -0
  119. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/langgraph/mcp.py +0 -0
  120. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/llama_index/__init__.py +0 -0
  121. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/llama_index/agent.py +0 -0
  122. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/llama_index/base.py +0 -0
  123. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/llama_index/mcp.py +0 -0
  124. /datarobot_genai-0.2.26/src/datarobot_genai/py.typed → /datarobot_genai-0.2.39/src/datarobot_genai/nat/__init__.py +0 -0
  125. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/nat/agent.py +0 -0
  126. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/nat/datarobot_auth_provider.py +0 -0
  127. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/nat/datarobot_llm_providers.py +0 -0
  128. {datarobot_genai-0.2.26 → datarobot_genai-0.2.39}/src/datarobot_genai/nat/helpers.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: datarobot-genai
3
- Version: 0.2.26
3
+ Version: 0.2.39
4
4
  Summary: Generic helpers for GenAI
5
5
  Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
6
6
  Author: DataRobot, Inc.
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "datarobot-genai"
7
- version = "0.2.26"
7
+ version = "0.2.39"
8
8
  description = "Generic helpers for GenAI"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10, <3.13"
@@ -107,7 +107,10 @@ class AgentKernel:
107
107
  stream: bool = False,
108
108
  config: Any | None = None,
109
109
  ) -> ChatCompletion | Stream[ChatCompletionChunk]:
110
- chat_api_url = config.agent_endpoint if config else self.base_url
110
+ if config is not None:
111
+ chat_api_url = f"http://localhost:{config.local_dev_port}"
112
+ else:
113
+ chat_api_url = self.base_url
111
114
  print(chat_api_url)
112
115
 
113
116
  return self._do_chat_completion(chat_api_url, user_prompt, completion_json, stream=stream)
@@ -19,11 +19,11 @@ A reusable library for building Model Context Protocol (MCP) servers with DataRo
19
19
  """
20
20
 
21
21
  # Export main server components
22
+ from datarobot_genai.drmcp.test_utils.clients.openai import OpenAILLMMCPClient
22
23
  from datarobot_genai.drmcp.test_utils.mcp_utils_ete import ete_test_mcp_session
23
24
  from datarobot_genai.drmcp.test_utils.mcp_utils_ete import get_dr_mcp_server_url
24
25
  from datarobot_genai.drmcp.test_utils.mcp_utils_ete import get_headers
25
26
  from datarobot_genai.drmcp.test_utils.mcp_utils_integration import integration_test_mcp_session
26
- from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import LLMMCPClient
27
27
  from datarobot_genai.drmcp.test_utils.tool_base_ete import ETETestExpectations
28
28
  from datarobot_genai.drmcp.test_utils.tool_base_ete import ToolBaseE2E
29
29
  from datarobot_genai.drmcp.test_utils.tool_base_ete import ToolCallTestExpectations
@@ -70,7 +70,7 @@ __all__ = [
70
70
  "get_dr_mcp_server_url",
71
71
  "get_headers",
72
72
  "ete_test_mcp_session",
73
- "LLMMCPClient",
73
+ "OpenAILLMMCPClient",
74
74
  "ETETestExpectations",
75
75
  "ToolBaseE2E",
76
76
  "ToolCallTestExpectations",
@@ -31,6 +31,124 @@ from .constants import DEFAULT_DATAROBOT_ENDPOINT
31
31
  from .constants import RUNTIME_PARAM_ENV_VAR_NAME_PREFIX
32
32
 
33
33
 
34
+ class MCPToolConfig(BaseSettings):
35
+ """Tool configuration for MCP server."""
36
+
37
+ enable_predictive_tools: bool = Field(
38
+ default=True,
39
+ validation_alias=AliasChoices(
40
+ RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_PREDICTIVE_TOOLS",
41
+ "ENABLE_PREDICTIVE_TOOLS",
42
+ ),
43
+ description="Enable/disable predictive tools",
44
+ )
45
+
46
+ enable_jira_tools: bool = Field(
47
+ default=False,
48
+ validation_alias=AliasChoices(
49
+ RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_JIRA_TOOLS",
50
+ "ENABLE_JIRA_TOOLS",
51
+ ),
52
+ description="Enable/disable Jira tools",
53
+ )
54
+
55
+ enable_confluence_tools: bool = Field(
56
+ default=False,
57
+ validation_alias=AliasChoices(
58
+ RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_CONFLUENCE_TOOLS",
59
+ "ENABLE_CONFLUENCE_TOOLS",
60
+ ),
61
+ description="Enable/disable Confluence tools",
62
+ )
63
+
64
+ enable_gdrive_tools: bool = Field(
65
+ default=False,
66
+ validation_alias=AliasChoices(
67
+ RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_GDRIVE_TOOLS",
68
+ "ENABLE_GDRIVE_TOOLS",
69
+ ),
70
+ description="Enable/disable GDrive tools",
71
+ )
72
+
73
+ enable_microsoft_graph_tools: bool = Field(
74
+ default=False,
75
+ validation_alias=AliasChoices(
76
+ RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_MICROSOFT_GRAPH_TOOLS",
77
+ "ENABLE_MICROSOFT_GRAPH_TOOLS",
78
+ ),
79
+ description="Enable/disable Sharepoint tools",
80
+ )
81
+
82
+ is_atlassian_oauth_provider_configured: bool = Field(
83
+ default=False,
84
+ validation_alias=AliasChoices(
85
+ RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_ATLASSIAN_OAUTH_PROVIDER_CONFIGURED",
86
+ "IS_ATLASSIAN_OAUTH_PROVIDER_CONFIGURED",
87
+ ),
88
+ description="Whether Atlassian OAuth provider is configured for Atlassian integration",
89
+ )
90
+
91
+ @property
92
+ def is_atlassian_oauth_configured(self) -> bool:
93
+ """Check if Atlassian OAuth is configured via provider flag or environment variables."""
94
+ return self.is_atlassian_oauth_provider_configured or bool(
95
+ os.getenv("ATLASSIAN_CLIENT_ID") and os.getenv("ATLASSIAN_CLIENT_SECRET")
96
+ )
97
+
98
+ is_google_oauth_provider_configured: bool = Field(
99
+ default=False,
100
+ validation_alias=AliasChoices(
101
+ RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_GOOGLE_OAUTH_PROVIDER_CONFIGURED",
102
+ "IS_GOOGLE_OAUTH_PROVIDER_CONFIGURED",
103
+ ),
104
+ description="Whether Google OAuth provider is configured for Google integration",
105
+ )
106
+
107
+ @property
108
+ def is_google_oauth_configured(self) -> bool:
109
+ return self.is_google_oauth_provider_configured or bool(
110
+ os.getenv("GOOGLE_CLIENT_ID") and os.getenv("GOOGLE_CLIENT_SECRET")
111
+ )
112
+
113
+ is_microsoft_oauth_provider_configured: bool = Field(
114
+ default=False,
115
+ validation_alias=AliasChoices(
116
+ RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_MICROSOFT_OAUTH_PROVIDER_CONFIGURED",
117
+ "IS_MICROSOFT_OAUTH_PROVIDER_CONFIGURED",
118
+ ),
119
+ description="Whether Microsoft OAuth provider is configured for Microsoft integration",
120
+ )
121
+
122
+ @property
123
+ def is_microsoft_oauth_configured(self) -> bool:
124
+ return self.is_microsoft_oauth_provider_configured or bool(
125
+ os.getenv("MICROSOFT_CLIENT_ID") and os.getenv("MICROSOFT_CLIENT_SECRET")
126
+ )
127
+
128
+ @field_validator(
129
+ "enable_predictive_tools",
130
+ "enable_jira_tools",
131
+ "enable_confluence_tools",
132
+ "enable_gdrive_tools",
133
+ "enable_microsoft_graph_tools",
134
+ "is_atlassian_oauth_provider_configured",
135
+ "is_google_oauth_provider_configured",
136
+ "is_microsoft_oauth_provider_configured",
137
+ mode="before",
138
+ )
139
+ @classmethod
140
+ def validate_runtime_params(cls, v: Any) -> Any:
141
+ """Validate runtime parameters."""
142
+ return extract_datarobot_runtime_param_payload(v)
143
+
144
+ model_config = SettingsConfigDict(
145
+ env_file=".env",
146
+ case_sensitive=False,
147
+ env_file_encoding="utf-8",
148
+ extra="ignore",
149
+ )
150
+
151
+
34
152
  class MCPServerConfig(BaseSettings):
35
153
  """MCP Server configuration using pydantic settings."""
36
154
 
@@ -188,86 +306,11 @@ class MCPServerConfig(BaseSettings):
188
306
  ),
189
307
  description="Enable/disable memory management",
190
308
  )
191
- enable_predictive_tools: bool = Field(
192
- default=True,
193
- validation_alias=AliasChoices(
194
- RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_PREDICTIVE_TOOLS",
195
- "ENABLE_PREDICTIVE_TOOLS",
196
- ),
197
- description="Enable/disable predictive tools",
198
- )
199
309
 
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",
310
+ tool_config: MCPToolConfig = Field(
311
+ default_factory=MCPToolConfig,
312
+ description="Tool configuration",
208
313
  )
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
-
248
- # Gdrive tools
249
- enable_gdrive_tools: bool = Field(
250
- default=False,
251
- validation_alias=AliasChoices(
252
- RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_GDRIVE_TOOLS",
253
- "ENABLE_GDRIVE_TOOLS",
254
- ),
255
- description="Enable/disable GDrive tools",
256
- )
257
- is_gdrive_oauth_provider_configured: bool = Field(
258
- default=False,
259
- validation_alias=AliasChoices(
260
- RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
261
- "IS_GDRIVE_OAUTH_PROVIDER_CONFIGURED",
262
- ),
263
- description="Whether GDrive OAuth provider is configured for GDrive integration",
264
- )
265
-
266
- @property
267
- def is_gdrive_oauth_configured(self) -> bool:
268
- return self.is_gdrive_oauth_provider_configured or bool(
269
- os.getenv("GDRIVE_CLIENT_ID") and os.getenv("GDRIVE_CLIENT_SECRET")
270
- )
271
314
 
272
315
  @field_validator(
273
316
  "otel_attributes",
@@ -291,11 +334,6 @@ class MCPServerConfig(BaseSettings):
291
334
  "mcp_server_register_dynamic_tools_on_startup",
292
335
  "tool_registration_duplicate_behavior",
293
336
  "mcp_server_register_dynamic_prompts_on_startup",
294
- "enable_predictive_tools",
295
- "enable_jira_tools",
296
- "is_jira_oauth_provider_configured",
297
- "enable_confluence_tools",
298
- "is_confluence_oauth_provider_configured",
299
337
  mode="before",
300
338
  )
301
339
  @classmethod
@@ -19,7 +19,3 @@ class DynamicToolRegistrationError(Exception):
19
19
 
20
20
  class DynamicPromptRegistrationError(Exception):
21
21
  """Exception raised for errors in the dynamic prompt registration process."""
22
-
23
-
24
- class MCPError(Exception):
25
- """Base class for MCP errors."""
@@ -20,7 +20,7 @@ from collections.abc import Callable
20
20
  from typing import Any
21
21
  from typing import TypeVar
22
22
 
23
- from .exceptions import MCPError
23
+ from fastmcp.exceptions import ToolError
24
24
 
25
25
  # Secret patterns to redact from logs
26
26
  SECRET_PATTERNS = [
@@ -93,6 +93,6 @@ def log_execution(func: F) -> F:
93
93
  return result
94
94
  except Exception as e:
95
95
  error_msg = _log_error(logger, func.__name__, e, args=args, kwargs=kwargs)
96
- raise MCPError(error_msg)
96
+ raise ToolError(error_msg)
97
97
 
98
98
  return wrapper # type: ignore[return-value]
@@ -30,6 +30,7 @@ class ToolType(str, Enum):
30
30
  JIRA = "jira"
31
31
  CONFLUENCE = "confluence"
32
32
  GDRIVE = "gdrive"
33
+ MICROSOFT_GRAPH = "microsoft_graph"
33
34
 
34
35
 
35
36
  class ToolConfig(TypedDict):
@@ -39,7 +40,7 @@ class ToolConfig(TypedDict):
39
40
  oauth_check: Callable[["MCPServerConfig"], bool] | None
40
41
  directory: str
41
42
  package_prefix: str
42
- config_field_name: str # Name of the config field (e.g., "enable_predictive_tools")
43
+ config_field_name: str
43
44
 
44
45
 
45
46
  # Tool configuration registry
@@ -53,25 +54,32 @@ TOOL_CONFIGS: dict[ToolType, ToolConfig] = {
53
54
  ),
54
55
  ToolType.JIRA: ToolConfig(
55
56
  name="jira",
56
- oauth_check=lambda config: config.is_jira_oauth_configured,
57
+ oauth_check=lambda config: config.tool_config.is_atlassian_oauth_configured,
57
58
  directory="jira",
58
59
  package_prefix="datarobot_genai.drmcp.tools.jira",
59
60
  config_field_name="enable_jira_tools",
60
61
  ),
61
62
  ToolType.CONFLUENCE: ToolConfig(
62
63
  name="confluence",
63
- oauth_check=lambda config: config.is_confluence_oauth_configured,
64
+ oauth_check=lambda config: config.tool_config.is_atlassian_oauth_configured,
64
65
  directory="confluence",
65
66
  package_prefix="datarobot_genai.drmcp.tools.confluence",
66
67
  config_field_name="enable_confluence_tools",
67
68
  ),
68
69
  ToolType.GDRIVE: ToolConfig(
69
70
  name="gdrive",
70
- oauth_check=lambda config: config.is_gdrive_oauth_configured,
71
+ oauth_check=lambda config: config.tool_config.is_google_oauth_configured,
71
72
  directory="gdrive",
72
73
  package_prefix="datarobot_genai.drmcp.tools.gdrive",
73
74
  config_field_name="enable_gdrive_tools",
74
75
  ),
76
+ ToolType.MICROSOFT_GRAPH: ToolConfig(
77
+ name="microsoft_graph",
78
+ oauth_check=lambda config: config.tool_config.is_microsoft_oauth_configured,
79
+ directory="microsoft_graph",
80
+ package_prefix="datarobot_genai.drmcp.tools.microsoft_graph",
81
+ config_field_name="enable_microsoft_graph_tools",
82
+ ),
75
83
  }
76
84
 
77
85
 
@@ -92,12 +100,12 @@ def is_tool_enabled(tool_type: ToolType, config: "MCPServerConfig") -> bool:
92
100
  -------
93
101
  True if the tool is enabled, False otherwise
94
102
  """
95
- tool_config = TOOL_CONFIGS[tool_type]
96
- enable_config_name = tool_config["config_field_name"]
97
- is_enabled = getattr(config, enable_config_name)
103
+ tool_config_registry = TOOL_CONFIGS[tool_type]
104
+ enable_config_name = tool_config_registry["config_field_name"]
105
+ is_enabled = getattr(config.tool_config, enable_config_name)
98
106
 
99
107
  # If tool is enabled, check OAuth requirements if needed
100
- if is_enabled and tool_config["oauth_check"] is not None:
101
- return tool_config["oauth_check"](config)
108
+ if is_enabled and tool_config_registry["oauth_check"] is not None:
109
+ return tool_config_registry["oauth_check"](config)
102
110
 
103
111
  return is_enabled
@@ -0,0 +1,68 @@
1
+ # Copyright 2026 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
+ """Anthropic LLM MCP Client implementation (example).
16
+
17
+ This is an example implementation showing how easy it is to add a new LLM provider.
18
+ Anthropic's API is OpenAI-compatible, so we can use the OpenAI SDK with their endpoint.
19
+ """
20
+
21
+ import openai
22
+
23
+ from .base import BaseLLMMCPClient
24
+
25
+
26
+ class AnthropicMCPClient(BaseLLMMCPClient):
27
+ """
28
+ Client for interacting with LLMs via MCP using Anthropic Claude.
29
+
30
+ Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
31
+ Tools using FastMCP's built-in elicitation will work automatically.
32
+
33
+ Example:
34
+ ```python
35
+ config = {
36
+ "anthropic_api_key": "sk-ant-...",
37
+ "model": "claude-3-5-sonnet-20241022",
38
+ }
39
+ client = AnthropicMCPClient(str(config))
40
+ ```
41
+ """
42
+
43
+ def __init__(
44
+ self,
45
+ config: str | dict,
46
+ ):
47
+ """
48
+ Initialize the LLM MCP client.
49
+
50
+ Args:
51
+ config: Configuration string or dict with:
52
+ - anthropic_api_key: Anthropic API key
53
+ - model: Model name (default: "claude-3-5-sonnet-20241022")
54
+ - save_llm_responses: Whether to save responses (default: True)
55
+ """
56
+ super().__init__(config)
57
+
58
+ def _create_llm_client(self, config_dict: dict) -> tuple[openai.OpenAI, str]:
59
+ """Create the LLM client for Anthropic (OpenAI-compatible endpoint)."""
60
+ anthropic_api_key = config_dict.get("anthropic_api_key")
61
+ model = config_dict.get("model", "claude-3-5-sonnet-20241022")
62
+
63
+ # Anthropic provides an OpenAI-compatible endpoint
64
+ client = openai.OpenAI(
65
+ api_key=anthropic_api_key,
66
+ base_url="https://api.anthropic.com/v1",
67
+ )
68
+ return client, model
@@ -1,4 +1,4 @@
1
- # Copyright 2025 DataRobot, Inc.
1
+ # Copyright 2026 DataRobot, Inc.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -12,7 +12,11 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
+ """Base classes for LLM MCP clients."""
16
+
15
17
  import json
18
+ from abc import ABC
19
+ from abc import abstractmethod
16
20
  from ast import literal_eval
17
21
  from typing import Any
18
22
 
@@ -23,7 +27,7 @@ from mcp.types import ListToolsResult
23
27
  from mcp.types import TextContent
24
28
  from openai.types.chat.chat_completion import ChatCompletion
25
29
 
26
- from .utils import save_response_to_file
30
+ from datarobot_genai.drmcp.test_utils.utils import save_response_to_file
27
31
 
28
32
 
29
33
  class ToolCall:
@@ -44,9 +48,9 @@ class LLMResponse:
44
48
  self.tool_results = tool_results
45
49
 
46
50
 
47
- class LLMMCPClient:
51
+ class BaseLLMMCPClient(ABC):
48
52
  """
49
- Client for interacting with LLMs via MCP.
53
+ Base class for LLM MCP clients.
50
54
 
51
55
  Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
52
56
  Tools using FastMCP's built-in elicitation will work automatically.
@@ -54,54 +58,48 @@ class LLMMCPClient:
54
58
 
55
59
  def __init__(
56
60
  self,
57
- config: str,
61
+ config: str | dict,
58
62
  ):
59
63
  """
60
64
  Initialize the LLM MCP client.
61
65
 
62
66
  Args:
63
- config: Configuration string or dict with:
64
- - openai_api_key: OpenAI API key
65
- - openai_api_base: Optional Azure OpenAI endpoint
66
- - openai_api_deployment_id: Optional Azure deployment ID
67
- - openai_api_version: Optional Azure API version
68
- - model: Model name (default: "gpt-3.5-turbo")
69
- - save_llm_responses: Whether to save responses (default: True)
67
+ config: Configuration string or dict with provider-specific keys.
70
68
  """
71
- # Parse config string to extract parameters
69
+ config_dict = self._parse_config(config)
70
+ self.openai_client, self.model = self._create_llm_client(config_dict)
71
+ self.save_llm_responses = config_dict.get("save_llm_responses", True)
72
+ self.available_tools: list[dict[str, Any]] = []
73
+ self.available_prompts: list[dict[str, Any]] = []
74
+ self.available_resources: list[dict[str, Any]] = []
75
+
76
+ @staticmethod
77
+ def _parse_config(config: str | dict) -> dict:
78
+ """Parse config string to dict."""
72
79
  if isinstance(config, str):
73
80
  # Try JSON first (safer), fall back to literal_eval for Python dict strings
74
81
  try:
75
- config_dict = json.loads(config)
82
+ return json.loads(config)
76
83
  except json.JSONDecodeError:
77
84
  # Fall back to literal_eval for Python dict literal strings
78
- config_dict = literal_eval(config)
79
- else:
80
- config_dict = config
81
-
82
- openai_api_key = config_dict.get("openai_api_key")
83
- openai_api_base = config_dict.get("openai_api_base")
84
- openai_api_deployment_id = config_dict.get("openai_api_deployment_id")
85
- model = config_dict.get("model", "gpt-3.5-turbo")
86
- save_llm_responses = config_dict.get("save_llm_responses", True)
87
-
88
- if openai_api_base and openai_api_deployment_id:
89
- # Azure OpenAI
90
- self.openai_client = openai.AzureOpenAI(
91
- api_key=openai_api_key,
92
- azure_endpoint=openai_api_base,
93
- api_version=config_dict.get("openai_api_version", "2024-02-15-preview"),
94
- )
95
- self.model = openai_api_deployment_id
96
- else:
97
- # Regular OpenAI
98
- self.openai_client = openai.OpenAI(api_key=openai_api_key) # type: ignore[assignment]
99
- self.model = model
85
+ return literal_eval(config)
86
+ return config
100
87
 
101
- self.save_llm_responses = save_llm_responses
102
- self.available_tools: list[dict[str, Any]] = []
103
- self.available_prompts: list[dict[str, Any]] = []
104
- self.available_resources: list[dict[str, Any]] = []
88
+ @abstractmethod
89
+ def _create_llm_client(
90
+ self, config_dict: dict
91
+ ) -> tuple[openai.OpenAI | openai.AzureOpenAI, str]:
92
+ """
93
+ Create the LLM client.
94
+
95
+ Args:
96
+ config_dict: Parsed configuration dictionary
97
+
98
+ Returns
99
+ -------
100
+ Tuple of (LLM client instance, model name)
101
+ """
102
+ pass
105
103
 
106
104
  async def _add_mcp_tool_to_available_tools(self, mcp_session: ClientSession) -> None:
107
105
  """Add a tool to the available tools."""
@@ -0,0 +1,58 @@
1
+ # Copyright 2026 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
+ """DataRobot LLM Gateway MCP Client implementation."""
16
+
17
+ import openai
18
+
19
+ from .base import BaseLLMMCPClient
20
+
21
+
22
+ class DRLLMGatewayMCPClient(BaseLLMMCPClient):
23
+ """
24
+ Client for interacting with LLMs via MCP using DataRobot LLM Gateway.
25
+
26
+ Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
27
+ Tools using FastMCP's built-in elicitation will work automatically.
28
+ """
29
+
30
+ def __init__(
31
+ self,
32
+ config: str | dict,
33
+ ):
34
+ """
35
+ Initialize the LLM MCP client.
36
+
37
+ Args:
38
+ config: Configuration string or dict with:
39
+ - datarobot_api_token: DataRobot API token
40
+ - datarobot_endpoint: DataRobot endpoint URL (default: "https://app.datarobot.com/api/v2")
41
+ - model: Model name (default: "gpt-4o-mini")
42
+ - save_llm_responses: Whether to save responses (default: True)
43
+ """
44
+ super().__init__(config)
45
+
46
+ def _create_llm_client(self, config_dict: dict) -> tuple[openai.OpenAI, str]:
47
+ """Create the LLM client for DataRobot LLM Gateway."""
48
+ datarobot_api_token = config_dict.get("datarobot_api_token")
49
+ datarobot_endpoint = config_dict.get(
50
+ "datarobot_endpoint", "https://app.datarobot.com/api/v2"
51
+ )
52
+ model = config_dict.get("model", "gpt-4o-mini")
53
+
54
+ # Build gateway URL: {endpoint}/genai/llmgw
55
+ gateway_url = datarobot_endpoint.rstrip("/") + "/genai/llmgw"
56
+
57
+ client = openai.OpenAI(api_key=datarobot_api_token, base_url=gateway_url)
58
+ return client, model