datarobot-genai 0.2.26__tar.gz → 0.2.31__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.
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/PKG-INFO +1 -1
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/pyproject.toml +1 -1
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/__init__.py +2 -2
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/config.py +121 -83
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/tool_config.py +17 -9
- datarobot_genai-0.2.31/src/datarobot_genai/drmcp/test_utils/clients/anthropic.py +68 -0
- datarobot_genai-0.2.26/src/datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py → datarobot_genai-0.2.31/src/datarobot_genai/drmcp/test_utils/clients/base.py +38 -40
- datarobot_genai-0.2.31/src/datarobot_genai/drmcp/test_utils/clients/dr_gateway.py +58 -0
- datarobot_genai-0.2.31/src/datarobot_genai/drmcp/test_utils/clients/openai.py +68 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/test_utils/test_interactive.py +16 -16
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/test_utils/tool_base_ete.py +69 -2
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/test_utils/utils.py +1 -1
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/clients/gdrive.py +223 -1
- datarobot_genai-0.2.31/src/datarobot_genai/drmcp/tools/clients/microsoft_graph.py +479 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/gdrive/tools.py +173 -3
- datarobot_genai-0.2.31/src/datarobot_genai/drmcp/tools/microsoft_graph/__init__.py +13 -0
- datarobot_genai-0.2.31/src/datarobot_genai/drmcp/tools/microsoft_graph/tools.py +198 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/predictive/data.py +11 -3
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/predictive/training.py +1 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/nat/datarobot_llm_clients.py +90 -54
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/nat/datarobot_mcp_client.py +47 -15
- datarobot_genai-0.2.31/src/datarobot_genai/py.typed +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/.gitignore +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/AUTHORS +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/LICENSE +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/README.md +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/agents/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/agents/base.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/chat/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/chat/auth.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/chat/client.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/chat/responses.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/cli/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/cli/agent_environment.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/cli/agent_kernel.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/custom_model.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/mcp/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/mcp/common.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/telemetry_agent.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/utils/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/utils/auth.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/core/utils/urls.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/crewai/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/crewai/agent.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/crewai/base.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/crewai/events.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/crewai/mcp.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/auth.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/clients.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/config_utils.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/constants.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/credentials.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dr_mcp_server.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dr_mcp_server_logo.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_prompts/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_prompts/controllers.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_prompts/register.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_prompts/utils.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/base.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/default.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/drum.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/config.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/controllers.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/metadata.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/register.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_agentic_fallback_schema.json +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_prediction_fallback_schema.json +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/register.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/dynamic_tools/schema.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/exceptions.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/logging.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/mcp_instance.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/memory_management/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/memory_management/manager.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/memory_management/memory_tools.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/routes.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/routes_utils.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/server_life_cycle.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/telemetry.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/tool_filter.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/utils.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/server.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/test_utils/__init__.py +0 -0
- {datarobot_genai-0.2.26/src/datarobot_genai/drmcp/tools/gdrive → datarobot_genai-0.2.31/src/datarobot_genai/drmcp/test_utils/clients}/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/test_utils/elicitation_test_tool.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/test_utils/integration_mcp_server.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/test_utils/mcp_utils_ete.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/test_utils/mcp_utils_integration.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/clients/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/clients/atlassian.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/clients/confluence.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/clients/jira.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/clients/s3.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/confluence/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/confluence/tools.py +0 -0
- {datarobot_genai-0.2.26/src/datarobot_genai/langgraph → datarobot_genai-0.2.31/src/datarobot_genai/drmcp/tools/gdrive}/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/jira/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/jira/tools.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/predictive/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/predictive/deployment.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/predictive/deployment_info.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/predictive/model.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/predictive/predict.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/predictive/predict_realtime.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/tools/predictive/project.py +0 -0
- {datarobot_genai-0.2.26/src/datarobot_genai/nat → datarobot_genai-0.2.31/src/datarobot_genai/langgraph}/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/langgraph/agent.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/langgraph/mcp.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/llama_index/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/llama_index/agent.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/llama_index/base.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/llama_index/mcp.py +0 -0
- /datarobot_genai-0.2.26/src/datarobot_genai/py.typed → /datarobot_genai-0.2.31/src/datarobot_genai/nat/__init__.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/nat/agent.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/nat/datarobot_auth_provider.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/nat/datarobot_llm_providers.py +0 -0
- {datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/nat/helpers.py +0 -0
|
@@ -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
|
-
"
|
|
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
|
-
|
|
201
|
-
|
|
202
|
-
|
|
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
|
{datarobot_genai-0.2.26 → datarobot_genai-0.2.31}/src/datarobot_genai/drmcp/core/tool_config.py
RENAMED
|
@@ -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
|
|
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.
|
|
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.
|
|
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.
|
|
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
|
-
|
|
96
|
-
enable_config_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
|
|
101
|
-
return
|
|
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
|
|
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
|
|
30
|
+
from ..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
|
|
51
|
+
class BaseLLMMCPClient(ABC):
|
|
48
52
|
"""
|
|
49
|
-
|
|
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
|
-
|
|
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
|
-
|
|
82
|
+
return json.loads(config)
|
|
76
83
|
except json.JSONDecodeError:
|
|
77
84
|
# Fall back to literal_eval for Python dict literal strings
|
|
78
|
-
|
|
79
|
-
|
|
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
|
-
|
|
102
|
-
|
|
103
|
-
self
|
|
104
|
-
|
|
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
|
|
@@ -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
|
+
"""OpenAI LLM MCP Client implementation."""
|
|
16
|
+
|
|
17
|
+
import openai
|
|
18
|
+
|
|
19
|
+
from .base import BaseLLMMCPClient
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class OpenAILLMMCPClient(BaseLLMMCPClient):
|
|
23
|
+
"""
|
|
24
|
+
Client for interacting with LLMs via MCP using OpenAI or Azure OpenAI.
|
|
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
|
+
- openai_api_key: OpenAI API key
|
|
40
|
+
- openai_api_base: Optional Azure OpenAI endpoint
|
|
41
|
+
- openai_api_deployment_id: Optional Azure deployment ID
|
|
42
|
+
- openai_api_version: Optional Azure API version
|
|
43
|
+
- model: Model name (default: "gpt-3.5-turbo")
|
|
44
|
+
- save_llm_responses: Whether to save responses (default: True)
|
|
45
|
+
"""
|
|
46
|
+
super().__init__(config)
|
|
47
|
+
|
|
48
|
+
def _create_llm_client(
|
|
49
|
+
self, config_dict: dict
|
|
50
|
+
) -> tuple[openai.OpenAI | openai.AzureOpenAI, str]:
|
|
51
|
+
"""Create the LLM client for OpenAI or Azure OpenAI."""
|
|
52
|
+
openai_api_key = config_dict.get("openai_api_key")
|
|
53
|
+
openai_api_base = config_dict.get("openai_api_base")
|
|
54
|
+
openai_api_deployment_id = config_dict.get("openai_api_deployment_id")
|
|
55
|
+
model = config_dict.get("model", "gpt-3.5-turbo")
|
|
56
|
+
|
|
57
|
+
if openai_api_base and openai_api_deployment_id:
|
|
58
|
+
# Azure OpenAI
|
|
59
|
+
client = openai.AzureOpenAI(
|
|
60
|
+
api_key=openai_api_key,
|
|
61
|
+
azure_endpoint=openai_api_base,
|
|
62
|
+
api_version=config_dict.get("openai_api_version", "2024-02-15-preview"),
|
|
63
|
+
)
|
|
64
|
+
return client, openai_api_deployment_id
|
|
65
|
+
else:
|
|
66
|
+
# Regular OpenAI
|
|
67
|
+
client = openai.OpenAI(api_key=openai_api_key) # type: ignore[assignment]
|
|
68
|
+
return client, model
|