datarobot-genai 0.2.17__py3-none-any.whl → 0.2.18__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,89 @@
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
+ """Test tool for elicitation testing.
16
+
17
+ This module registers a test tool that can be used to test elicitation support.
18
+ It should be imported in tests that need it.
19
+ """
20
+
21
+ from fastmcp import Context
22
+ from fastmcp.server.context import AcceptedElicitation
23
+ from fastmcp.server.context import CancelledElicitation
24
+ from fastmcp.server.context import DeclinedElicitation
25
+
26
+ from datarobot_genai.drmcp.core.mcp_instance import mcp
27
+
28
+
29
+ @mcp.tool(
30
+ name="get_user_greeting",
31
+ description=(
32
+ "Get a personalized greeting for a user. "
33
+ "Requires a username - if not provided, will request it via elicitation."
34
+ ),
35
+ tags={"test", "elicitation"},
36
+ )
37
+ async def get_user_greeting(ctx: Context, username: str | None = None) -> dict:
38
+ """
39
+ Get a personalized greeting for a user.
40
+
41
+ This tool demonstrates FastMCP's built-in elicitation by requiring a username parameter.
42
+ If username is not provided, it uses ctx.elicit() to request it from the user.
43
+
44
+ Args:
45
+ ctx: FastMCP context (automatically injected)
46
+ username: The username to greet. If None, elicitation will be triggered.
47
+
48
+ Returns
49
+ -------
50
+ Dictionary with greeting message or error if elicitation was declined/cancelled
51
+ """
52
+ if not username:
53
+ # Use elicitation to request username from the client
54
+ try:
55
+ result = await ctx.elicit(
56
+ message="Username is required to generate a personalized greeting",
57
+ response_type=str,
58
+ )
59
+
60
+ if isinstance(result, AcceptedElicitation):
61
+ username = result.data
62
+ elif isinstance(result, DeclinedElicitation):
63
+ return {
64
+ "status": "error",
65
+ "error": "Username declined by user",
66
+ "message": "Cannot generate greeting without username",
67
+ }
68
+ elif isinstance(result, CancelledElicitation):
69
+ return {
70
+ "status": "error",
71
+ "error": "Operation cancelled",
72
+ "message": "Greeting request was cancelled",
73
+ }
74
+ except Exception:
75
+ # Elicitation not supported by client - return graceful skip
76
+ return {
77
+ "status": "skipped",
78
+ "message": (
79
+ "Elicitation not supported by client. "
80
+ "Username parameter is required when client does not support elicitation."
81
+ ),
82
+ "elicitation_supported": False,
83
+ }
84
+
85
+ return {
86
+ "status": "success",
87
+ "message": f"Hello, {username}! Welcome to the DataRobot MCP server.",
88
+ "username": username,
89
+ }
@@ -26,6 +26,13 @@ from typing import Any
26
26
 
27
27
  from datarobot_genai.drmcp import create_mcp_server
28
28
 
29
+ # Import elicitation test tool to register it with the MCP server
30
+ try:
31
+ from datarobot_genai.drmcp.test_utils import elicitation_test_tool # noqa: F401
32
+ except ImportError:
33
+ # Test utils not available (e.g., running in production)
34
+ pass
35
+
29
36
  # Import user components (will be used conditionally)
30
37
  try:
31
38
  from app.core.server_lifecycle import ServerLifecycle # type: ignore # noqa: F401
@@ -15,6 +15,7 @@ import asyncio
15
15
  import os
16
16
  from collections.abc import AsyncGenerator
17
17
  from contextlib import asynccontextmanager
18
+ from typing import Any
18
19
 
19
20
  import aiohttp
20
21
  from aiohttp import ClientSession as HttpClientSession
@@ -78,6 +79,7 @@ def get_headers() -> dict[str, str]:
78
79
  @asynccontextmanager
79
80
  async def ete_test_mcp_session(
80
81
  additional_headers: dict[str, str] | None = None,
82
+ elicitation_callback: Any | None = None,
81
83
  ) -> AsyncGenerator[ClientSession, None]:
82
84
  """Create an MCP session for each test.
83
85
 
@@ -85,6 +87,10 @@ async def ete_test_mcp_session(
85
87
  ----------
86
88
  additional_headers : dict[str, str], optional
87
89
  Additional headers to include in the MCP session (e.g., auth headers for testing).
90
+ elicitation_callback : callable, optional
91
+ Callback function to handle elicitation requests from the server.
92
+ The callback should have signature:
93
+ async def callback(context, params: ElicitRequestParams) -> ElicitResult
88
94
  """
89
95
  try:
90
96
  headers = get_headers()
@@ -96,7 +102,9 @@ async def ete_test_mcp_session(
96
102
  write_stream,
97
103
  _,
98
104
  ):
99
- async with ClientSession(read_stream, write_stream) as session:
105
+ async with ClientSession(
106
+ read_stream, write_stream, elicitation_callback=elicitation_callback
107
+ ) as session:
100
108
  await asyncio.wait_for(session.initialize(), timeout=5)
101
109
  yield session
102
110
  except asyncio.TimeoutError:
@@ -17,6 +17,7 @@ import contextlib
17
17
  import os
18
18
  from collections.abc import AsyncGenerator
19
19
  from pathlib import Path
20
+ from typing import Any
20
21
 
21
22
  from mcp import ClientSession
22
23
  from mcp.client.stdio import StdioServerParameters
@@ -34,7 +35,12 @@ def integration_test_mcp_server_params() -> StdioServerParameters:
34
35
  or "https://test.datarobot.com/api/v2",
35
36
  "MCP_SERVER_LOG_LEVEL": os.environ.get("MCP_SERVER_LOG_LEVEL") or "WARNING",
36
37
  "APP_LOG_LEVEL": os.environ.get("APP_LOG_LEVEL") or "WARNING",
37
- "OTEL_ENABLED": os.environ.get("OTEL_ENABLED") or "false",
38
+ # Disable all OTEL telemetry for integration tests
39
+ "OTEL_ENABLED": "false",
40
+ "OTEL_SDK_DISABLED": "true",
41
+ "OTEL_TRACES_EXPORTER": "none",
42
+ "OTEL_LOGS_EXPORTER": "none",
43
+ "OTEL_METRICS_EXPORTER": "none",
38
44
  "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP": os.environ.get(
39
45
  "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP"
40
46
  )
@@ -64,7 +70,9 @@ def integration_test_mcp_server_params() -> StdioServerParameters:
64
70
 
65
71
  @contextlib.asynccontextmanager
66
72
  async def integration_test_mcp_session(
67
- server_params: StdioServerParameters | None = None, timeout: int = 30
73
+ server_params: StdioServerParameters | None = None,
74
+ timeout: int = 30,
75
+ elicitation_callback: Any | None = None,
68
76
  ) -> AsyncGenerator[ClientSession, None]:
69
77
  """
70
78
  Create and connect a client for the MCP server as a context manager.
@@ -72,6 +80,7 @@ async def integration_test_mcp_session(
72
80
  Args:
73
81
  server_params: Parameters for configuring the server connection
74
82
  timeout: Timeout
83
+ elicitation_callback: Optional callback for handling elicitation requests
75
84
 
76
85
  Yields
77
86
  ------
@@ -86,8 +95,12 @@ async def integration_test_mcp_session(
86
95
 
87
96
  try:
88
97
  async with stdio_client(server_params) as (read_stream, write_stream):
89
- async with ClientSession(read_stream, write_stream) as session:
90
- await asyncio.wait_for(session.initialize(), timeout=timeout)
98
+ async with ClientSession(
99
+ read_stream, write_stream, elicitation_callback=elicitation_callback
100
+ ) as session:
101
+ init_result = await asyncio.wait_for(session.initialize(), timeout=timeout)
102
+ # Store the init result on the session for tests that need to inspect capabilities
103
+ session._init_result = init_result # type: ignore[attr-defined]
91
104
  yield session
92
105
 
93
106
  except asyncio.TimeoutError:
@@ -13,6 +13,7 @@
13
13
  # limitations under the License.
14
14
 
15
15
  import json
16
+ from ast import literal_eval
16
17
  from typing import Any
17
18
 
18
19
  import openai
@@ -44,12 +45,39 @@ class LLMResponse:
44
45
 
45
46
 
46
47
  class LLMMCPClient:
47
- """Client for interacting with LLMs via MCP."""
48
+ """
49
+ Client for interacting with LLMs via MCP.
48
50
 
49
- def __init__(self, config: str):
50
- """Initialize the LLM MCP client."""
51
+ Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
52
+ Tools using FastMCP's built-in elicitation will work automatically.
53
+ """
54
+
55
+ def __init__(
56
+ self,
57
+ config: str,
58
+ ):
59
+ """
60
+ Initialize the LLM MCP client.
61
+
62
+ 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)
70
+ """
51
71
  # Parse config string to extract parameters
52
- config_dict = eval(config) if isinstance(config, str) else config
72
+ if isinstance(config, str):
73
+ # Try JSON first (safer), fall back to literal_eval for Python dict strings
74
+ try:
75
+ config_dict = json.loads(config)
76
+ except json.JSONDecodeError:
77
+ # Fall back to literal_eval for Python dict literal strings
78
+ config_dict = literal_eval(config)
79
+ else:
80
+ config_dict = config
53
81
 
54
82
  openai_api_key = config_dict.get("openai_api_key")
55
83
  openai_api_base = config_dict.get("openai_api_base")
@@ -93,7 +121,21 @@ class LLMMCPClient:
93
121
  async def _call_mcp_tool(
94
122
  self, tool_name: str, parameters: dict[str, Any], mcp_session: ClientSession
95
123
  ) -> str:
96
- """Call an MCP tool and return the result as a string."""
124
+ """
125
+ Call an MCP tool and return the result as a string.
126
+
127
+ Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
128
+ Tools using FastMCP's built-in elicitation will work automatically.
129
+
130
+ Args:
131
+ tool_name: Name of the tool to call
132
+ parameters: Parameters to pass to the tool
133
+ mcp_session: MCP client session
134
+
135
+ Returns
136
+ -------
137
+ Result text from the tool call
138
+ """
97
139
  result: CallToolResult = await mcp_session.call_tool(tool_name, parameters)
98
140
  content = (
99
141
  result.content[0].text
@@ -177,7 +219,26 @@ class LLMMCPClient:
177
219
  async def process_prompt_with_mcp_support(
178
220
  self, prompt: str, mcp_session: ClientSession, output_file_name: str = ""
179
221
  ) -> LLMResponse:
180
- """Process a prompt with MCP tool support."""
222
+ """
223
+ Process a prompt with MCP tool support and elicitation handling.
224
+
225
+ This method:
226
+ 1. Adds MCP tools to available tools
227
+ 2. Sends prompt to LLM
228
+ 3. Processes tool calls
229
+ 4. Continues until LLM provides final response
230
+
231
+ Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
232
+
233
+ Args:
234
+ prompt: User prompt
235
+ mcp_session: MCP client session
236
+ output_file_name: Optional file name to save response
237
+
238
+ Returns
239
+ -------
240
+ LLMResponse with content, tool calls, and tool results
241
+ """
181
242
  # Add MCP tools to available tools
182
243
  await self._add_mcp_tool_to_available_tools(mcp_session)
183
244
 
@@ -191,8 +252,10 @@ class LLMMCPClient:
191
252
  "content": (
192
253
  "You are a helpful AI assistant that can use tools to help users. "
193
254
  "If you need more information to provide a complete response, you can make "
194
- "multiple tool calls. When dealing with file paths, use them as raw paths "
195
- "without converting to file:// URLs."
255
+ "multiple tool calls or ask the user for more info, but prefer tool calls "
256
+ "when possible. "
257
+ "When dealing with file paths, use them as raw paths without converting "
258
+ "to file:// URLs."
196
259
  ),
197
260
  },
198
261
  {"role": "user", "content": prompt},
@@ -0,0 +1,205 @@
1
+ #!/usr/bin/env python3
2
+
3
+ # Copyright 2025 DataRobot, Inc.
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+
17
+ """Interactive MCP Client Test Script.
18
+
19
+ This script allows you to test arbitrary commands with the MCP server
20
+ using an LLM agent that can decide which tools to call.
21
+
22
+ Supports elicitation - when tools require user input (like authentication tokens),
23
+ the script will prompt you interactively.
24
+ """
25
+
26
+ import asyncio
27
+ import json
28
+ import os
29
+ import sys
30
+ import traceback
31
+ from pathlib import Path
32
+ from typing import Any
33
+
34
+ from dotenv import load_dotenv
35
+ from mcp import ClientSession
36
+ from mcp.client.streamable_http import streamablehttp_client
37
+ from mcp.shared.context import RequestContext
38
+ from mcp.types import ElicitRequestParams
39
+ from mcp.types import ElicitResult
40
+
41
+ from datarobot_genai.drmcp import get_dr_mcp_server_url
42
+ from datarobot_genai.drmcp import get_headers
43
+ from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import LLMMCPClient
44
+ from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import LLMResponse
45
+ from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import ToolCall
46
+
47
+ # Re-export for backwards compatibility
48
+ __all__ = ["LLMMCPClient", "LLMResponse", "ToolCall", "test_mcp_interactive"]
49
+
50
+
51
+ async def test_mcp_interactive() -> None:
52
+ """Test the MCP server interactively with LLM agent."""
53
+ # Check for required environment variables
54
+ openai_api_key = os.environ.get("OPENAI_API_KEY")
55
+ if not openai_api_key:
56
+ print("❌ Error: OPENAI_API_KEY environment variable is required")
57
+ print("Please set it in your .env file or export it")
58
+ return
59
+
60
+ # Optional Azure OpenAI settings
61
+ openai_api_base = os.environ.get("OPENAI_API_BASE")
62
+ openai_api_deployment_id = os.environ.get("OPENAI_API_DEPLOYMENT_ID")
63
+ openai_api_version = os.environ.get("OPENAI_API_VERSION")
64
+
65
+ print("🤖 Initializing LLM MCP Client...")
66
+
67
+ # Initialize the LLM client with elicitation handler
68
+ config = {
69
+ "openai_api_key": openai_api_key,
70
+ "openai_api_base": openai_api_base,
71
+ "openai_api_deployment_id": openai_api_deployment_id,
72
+ "openai_api_version": openai_api_version,
73
+ "save_llm_responses": False,
74
+ }
75
+
76
+ llm_client = LLMMCPClient(str(config))
77
+
78
+ # Get MCP server URL
79
+ mcp_server_url = get_dr_mcp_server_url()
80
+ if not mcp_server_url:
81
+ print("❌ Error: MCP server URL is not configured")
82
+ print("Please set DR_MCP_SERVER_URL environment variable or run: task test-interactive")
83
+ return
84
+
85
+ print(f"🔗 Connecting to MCP server at: {mcp_server_url}")
86
+
87
+ # Elicitation handler: prompt user for required values
88
+ async def elicitation_handler(
89
+ context: RequestContext[ClientSession, Any], params: ElicitRequestParams
90
+ ) -> ElicitResult:
91
+ print(f"\n📋 Elicitation Request: {params.message}")
92
+ if params.requestedSchema:
93
+ print(f" Schema: {params.requestedSchema}")
94
+
95
+ while True:
96
+ try:
97
+ response = input(" Enter value (or 'decline'/'cancel'): ").strip()
98
+ except (EOFError, KeyboardInterrupt):
99
+ return ElicitResult(action="cancel")
100
+
101
+ if response.lower() == "decline":
102
+ return ElicitResult(action="decline")
103
+ if response.lower() == "cancel":
104
+ return ElicitResult(action="cancel")
105
+ if response:
106
+ return ElicitResult(action="accept", content={"value": response})
107
+ print(" Please enter a value or 'decline'/'cancel'")
108
+
109
+ try:
110
+ async with streamablehttp_client(
111
+ url=mcp_server_url,
112
+ headers=get_headers(),
113
+ ) as (read_stream, write_stream, _):
114
+ async with ClientSession(
115
+ read_stream,
116
+ write_stream,
117
+ elicitation_callback=elicitation_handler,
118
+ ) as session:
119
+ await session.initialize()
120
+
121
+ print("✅ Connected to MCP server!")
122
+ print("📋 Available tools:")
123
+
124
+ tools_result = await session.list_tools()
125
+ for i, tool in enumerate(tools_result.tools, 1):
126
+ print(f" {i}. {tool.name}: {tool.description}")
127
+
128
+ print("\n" + "=" * 60)
129
+ print("🎯 Interactive Testing Mode")
130
+ print("=" * 60)
131
+ print("Type your questions/commands. The AI will decide which tools to use.")
132
+ print("If a tool requires additional information, you will be prompted.")
133
+ print("Type 'quit' or 'exit' to stop.")
134
+ print()
135
+
136
+ while True:
137
+ try:
138
+ user_input = input("🤔 You: ").strip()
139
+
140
+ if user_input.lower() in ["quit", "exit", "q"]:
141
+ print("👋 Goodbye!")
142
+ break
143
+
144
+ if not user_input:
145
+ continue
146
+ except (EOFError, KeyboardInterrupt):
147
+ print("\n👋 Goodbye!")
148
+ break
149
+
150
+ print("🤖 AI is thinking...")
151
+
152
+ response = await llm_client.process_prompt_with_mcp_support(
153
+ prompt=user_input,
154
+ mcp_session=session,
155
+ )
156
+
157
+ print("\n🤖 AI Response:")
158
+ print("-" * 40)
159
+ print(response.content)
160
+
161
+ if response.tool_calls:
162
+ print("\n🔧 Tools Used:")
163
+ for i, tool_call in enumerate(response.tool_calls, 1):
164
+ print(f" {i}. {tool_call.tool_name}")
165
+ print(f" Parameters: {tool_call.parameters}")
166
+ print(f" Reasoning: {tool_call.reasoning}")
167
+
168
+ if i <= len(response.tool_results):
169
+ result = response.tool_results[i - 1]
170
+ try:
171
+ result_data = json.loads(result)
172
+ if result_data.get("status") == "error":
173
+ error_msg = result_data.get("error", "Unknown error")
174
+ print(f" ❌ Error: {error_msg}")
175
+ elif result_data.get("status") == "success":
176
+ print(" ✅ Success")
177
+ except json.JSONDecodeError:
178
+ if len(result) > 100:
179
+ print(f" Result: {result[:100]}...")
180
+ else:
181
+ print(f" Result: {result}")
182
+
183
+ print("\n" + "=" * 60)
184
+ except Exception as e:
185
+ print(f"❌ Connection Error: {e}")
186
+ print(f" Server URL: {mcp_server_url}")
187
+ traceback.print_exc()
188
+ return
189
+
190
+
191
+ if __name__ == "__main__":
192
+ # Ensure we're in the right directory
193
+ if not Path("src").exists():
194
+ print("❌ Error: Please run this script from the project root")
195
+ sys.exit(1)
196
+
197
+ # Load environment variables from .env file
198
+ print("📄 Loading environment variables...")
199
+ load_dotenv()
200
+
201
+ print("🚀 Starting Interactive MCP Client Test")
202
+ print("Make sure the MCP server is running with: task drmcp-dev")
203
+ print()
204
+
205
+ asyncio.run(test_mcp_interactive())
@@ -116,28 +116,30 @@ class ToolBaseE2E:
116
116
  f"Should have called {test_expectations.tool_calls_expected[i].name} tool, but "
117
117
  f"got: {tool_call.tool_name}"
118
118
  )
119
- assert tool_call.parameters == test_expectations.tool_calls_expected[i].parameters, (
120
- f"Should have called {tool_call.tool_name} tool with the correct parameters, but "
121
- f"got: {tool_call.parameters}"
122
- )
123
- if test_expectations.tool_calls_expected[i].result != SHOULD_NOT_BE_EMPTY:
124
- expected_result = test_expectations.tool_calls_expected[i].result
125
- if isinstance(expected_result, str):
126
- assert expected_result in response.tool_results[i], (
127
- f"Should have called {tool_call.tool_name} tool with the correct result, "
128
- f"but got: {response.tool_results[i]}"
129
- )
119
+ assert (
120
+ tool_call.parameters == test_expectations.tool_calls_expected[i].parameters
121
+ ), (
122
+ f"Should have called {tool_call.tool_name} tool with the correct parameters, "
123
+ f"but got: {tool_call.parameters}"
124
+ )
125
+ if test_expectations.tool_calls_expected[i].result != SHOULD_NOT_BE_EMPTY:
126
+ expected_result = test_expectations.tool_calls_expected[i].result
127
+ if isinstance(expected_result, str):
128
+ assert expected_result in response.tool_results[i], (
129
+ f"Should have called {tool_call.tool_name} tool with the correct "
130
+ f"result, but got: {response.tool_results[i]}"
131
+ )
132
+ else:
133
+ actual_result = json.loads(response.tool_results[i])
134
+ assert _check_dict_has_keys(expected_result, actual_result), (
135
+ f"Should have called {tool_call.tool_name} tool with the correct "
136
+ f"result structure, but got: {response.tool_results[i]}"
137
+ )
130
138
  else:
131
- actual_result = json.loads(response.tool_results[i])
132
- assert _check_dict_has_keys(expected_result, actual_result), (
133
- f"Should have called {tool_call.tool_name} tool with the correct result "
134
- f"structure, but got: {response.tool_results[i]}"
139
+ assert len(response.tool_results[i]) > 0, (
140
+ f"Should have called {tool_call.tool_name} tool with non-empty result, but "
141
+ f"got: {response.tool_results[i]}"
135
142
  )
136
- else:
137
- assert len(response.tool_results[i]) > 0, (
138
- f"Should have called {tool_call.tool_name} tool with non-empty result, but "
139
- f"got: {response.tool_results[i]}"
140
- )
141
143
 
142
144
  # Verify LLM provided comprehensive response
143
145
  assert len(response.content) > 100, "LLM should provide detailed response"
@@ -14,6 +14,10 @@
14
14
 
15
15
  import logging
16
16
  import os
17
+ from typing import Annotated
18
+
19
+ from fastmcp.exceptions import ToolError
20
+ from fastmcp.tools.tool import ToolResult
17
21
 
18
22
  from datarobot_genai.drmcp.core.clients import get_sdk_client
19
23
  from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
@@ -21,44 +25,48 @@ from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
21
25
  logger = logging.getLogger(__name__)
22
26
 
23
27
 
24
- @dr_mcp_tool(tags={"data", "management", "upload"})
25
- async def upload_dataset_to_ai_catalog(file_path: str) -> str:
26
- """
27
- Upload a dataset to the DataRobot AI Catalog.
28
-
29
- Args:
30
- file_path: Path to the file to upload.
31
-
32
- Returns
33
- -------
34
- A string summary of the upload result.
35
- """
36
- client = get_sdk_client()
28
+ @dr_mcp_tool(tags={"predictive", "data", "write", "upload", "catalog"})
29
+ async def upload_dataset_to_ai_catalog(
30
+ file_path: Annotated[str, "The path to the dataset file to upload."],
31
+ ) -> ToolResult:
32
+ """Upload a dataset to the DataRobot AI Catalog."""
37
33
  if not os.path.exists(file_path):
38
34
  logger.error(f"File not found: {file_path}")
39
- return f"File not found: {file_path}"
35
+ raise ToolError(f"File not found: {file_path}")
36
+
37
+ client = get_sdk_client()
40
38
  catalog_item = client.Dataset.create_from_file(file_path)
41
- logger.info(f"Successfully uploaded dataset: {catalog_item.id}")
42
- return f"AI Catalog ID: {catalog_item.id}"
43
39
 
40
+ return ToolResult(
41
+ content=f"Successfully uploaded dataset: {catalog_item.id}.",
42
+ structured_content={
43
+ "id": catalog_item.id,
44
+ "name": catalog_item.name,
45
+ "status": catalog_item.status,
46
+ },
47
+ )
44
48
 
45
- @dr_mcp_tool(tags={"data", "management", "list"})
46
- async def list_ai_catalog_items() -> str:
47
- """
48
- List all AI Catalog items (datasets) for the authenticated user.
49
49
 
50
- Returns
51
- -------
52
- A string summary of the AI Catalog items with their IDs and names.
53
- """
50
+ @dr_mcp_tool(tags={"predictive", "data", "read", "list", "catalog"})
51
+ async def list_ai_catalog_items() -> ToolResult:
52
+ """List all AI Catalog items (datasets) for the authenticated user."""
54
53
  client = get_sdk_client()
55
54
  datasets = client.Dataset.list()
55
+
56
56
  if not datasets:
57
57
  logger.info("No AI Catalog items found")
58
- return "No AI Catalog items found."
59
- result = "\n".join(f"{ds.id}: {ds.name}" for ds in datasets)
60
- logger.info(f"Found {len(datasets)} AI Catalog items")
61
- return result
58
+ return ToolResult(
59
+ content="No AI Catalog items found.",
60
+ structured_content={"datasets": []},
61
+ )
62
+
63
+ return ToolResult(
64
+ content=f"Found {len(datasets)} AI Catalog items.",
65
+ structured_content={
66
+ "datasets": [{"id": ds.id, "name": ds.name} for ds in datasets],
67
+ "count": len(datasets),
68
+ },
69
+ )
62
70
 
63
71
 
64
72
  # from fastmcp import Context
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: datarobot-genai
3
- Version: 0.2.17
3
+ Version: 0.2.18
4
4
  Summary: Generic helpers for GenAI
5
5
  Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
6
6
  Author: DataRobot, Inc.
@@ -67,11 +67,13 @@ datarobot_genai/drmcp/core/memory_management/__init__.py,sha256=y4yapzp3KnFMzSR6
67
67
  datarobot_genai/drmcp/core/memory_management/manager.py,sha256=gmc_SQs12YQFMWl2UbfWR40QmLV9XuCnwPZgQwKWrbA,30552
68
68
  datarobot_genai/drmcp/core/memory_management/memory_tools.py,sha256=AxzpwOlldmhhDfKZcAxaGs7Xih2SCe0XbQuXX5nQczI,6397
69
69
  datarobot_genai/drmcp/test_utils/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
70
- datarobot_genai/drmcp/test_utils/integration_mcp_server.py,sha256=MdoR7r3m9uT7crodyhY69yhkrM7Thpe__BBD9lB_2oA,3328
71
- datarobot_genai/drmcp/test_utils/mcp_utils_ete.py,sha256=rgZkPF26YCHX2FGppWE4v22l_NQ3kLSPSUimO0tD4nM,4402
72
- datarobot_genai/drmcp/test_utils/mcp_utils_integration.py,sha256=0sU29Khal0CelnHBDInyTRiuPKrFFbTbIomOoUbyMhs,3271
73
- datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py,sha256=TvTkDBcHscLDmqge9NhHxVo1ABtb0n4NmmG2318mQHU,9088
74
- datarobot_genai/drmcp/test_utils/tool_base_ete.py,sha256=-mKHBkGkyOKQCVS2LHFhSnRofIqJBbeAPRkwizBDtTg,6104
70
+ datarobot_genai/drmcp/test_utils/elicitation_test_tool.py,sha256=UVKwy39nl3XcVAh6IATcN-cWL2bfrprgRQ7fbK82jTI,3287
71
+ datarobot_genai/drmcp/test_utils/integration_mcp_server.py,sha256=YSk19tbaka_0ziqi7LoXie4SJs-cvi9-H00Go0ZtQWE,3575
72
+ datarobot_genai/drmcp/test_utils/mcp_utils_ete.py,sha256=46rH0fYYmUj7ygf968iRbdSp5u95v23BEw3Ts_c431Y,4788
73
+ datarobot_genai/drmcp/test_utils/mcp_utils_integration.py,sha256=sHA_BWtpgIAFp9IXiNkUeBartBMjLAauqkV9bYtCr-g,3874
74
+ datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py,sha256=YgyqHK09MB-PBaqT34heqvmvYYFtLpzzSJt7xuTJmDg,11224
75
+ datarobot_genai/drmcp/test_utils/test_interactive.py,sha256=guXvR8q2H6VUdmvIjEJcElQJCC6lQ-oTrzbD2EkHeCs,8025
76
+ datarobot_genai/drmcp/test_utils/tool_base_ete.py,sha256=wmI-xcL0rSr56-ZoGNB8np0CZHAU583o8-Kw7fRwtMw,6232
75
77
  datarobot_genai/drmcp/test_utils/utils.py,sha256=esGKFv8aO31-Qg3owayeWp32BYe1CdYOEutjjdbweCw,3048
76
78
  datarobot_genai/drmcp/tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
77
79
  datarobot_genai/drmcp/tools/clients/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
@@ -84,7 +86,7 @@ datarobot_genai/drmcp/tools/confluence/tools.py,sha256=jSF7yXGFqqlMcavkRIY4HbMxb
84
86
  datarobot_genai/drmcp/tools/jira/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
85
87
  datarobot_genai/drmcp/tools/jira/tools.py,sha256=dfkqTU2HH-7n44hX80ODFacKq0p0LOchFcZtIIKFNMM,9687
86
88
  datarobot_genai/drmcp/tools/predictive/__init__.py,sha256=WuOHlNNEpEmcF7gVnhckruJRKU2qtmJLE3E7zoCGLDo,1030
87
- datarobot_genai/drmcp/tools/predictive/data.py,sha256=k4EJxJrl8DYVGVfJ0DM4YTfnZlC_K3OUHZ0eRUzfluI,3165
89
+ datarobot_genai/drmcp/tools/predictive/data.py,sha256=nc2OdkcZs9_L_o466tkaKaNvuhmvqcJKJ3nrA4H8B_c,3488
88
90
  datarobot_genai/drmcp/tools/predictive/deployment.py,sha256=lm02Ayuo11L1hP41fgi3QpR1Eyty-Wc16rM0c8SgliM,3277
89
91
  datarobot_genai/drmcp/tools/predictive/deployment_info.py,sha256=BGEF_dmbxOBJR0n1Tt9TO2-iNTQSBTr-oQUyaxLZ0ZI,15297
90
92
  datarobot_genai/drmcp/tools/predictive/model.py,sha256=Yih5-KedJ-1yupPLXCJsCXOdyWWi9pRvgapXDlgXWJA,4891
@@ -106,9 +108,9 @@ datarobot_genai/nat/datarobot_llm_clients.py,sha256=Yu208Ed_p_4P3HdpuM7fYnKcXtim
106
108
  datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
107
109
  datarobot_genai/nat/datarobot_mcp_client.py,sha256=35FzilxNp4VqwBYI0NsOc91-xZm1C-AzWqrOdDy962A,9612
108
110
  datarobot_genai/nat/helpers.py,sha256=Q7E3ADZdtFfS8E6OQPyw2wgA6laQ58N3bhLj5CBWwJs,3265
109
- datarobot_genai-0.2.17.dist-info/METADATA,sha256=FAAExscPOUIc1_sTNKY3DPTygcVtfrql4fubM9b6nDg,6301
110
- datarobot_genai-0.2.17.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
111
- datarobot_genai-0.2.17.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
112
- datarobot_genai-0.2.17.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
113
- datarobot_genai-0.2.17.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
114
- datarobot_genai-0.2.17.dist-info/RECORD,,
111
+ datarobot_genai-0.2.18.dist-info/METADATA,sha256=VQhojO4HSwkoFiHSGH9AdedWp-JV2RkRH-IewqJa38c,6301
112
+ datarobot_genai-0.2.18.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
113
+ datarobot_genai-0.2.18.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
114
+ datarobot_genai-0.2.18.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
115
+ datarobot_genai-0.2.18.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
116
+ datarobot_genai-0.2.18.dist-info/RECORD,,