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

datarobot_genai/drmcp/test_utils/elicitation_test_tool.py

@@ -0,0 +1,89 @@
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test tool for elicitation testing.
+
+This module registers a test tool that can be used to test elicitation support.
+It should be imported in tests that need it.
+"""
+
+from fastmcp import Context
+from fastmcp.server.context import AcceptedElicitation
+from fastmcp.server.context import CancelledElicitation
+from fastmcp.server.context import DeclinedElicitation
+
+from datarobot_genai.drmcp.core.mcp_instance import mcp
+
+
+@mcp.tool(
+    name="get_user_greeting",
+    description=(
+        "Get a personalized greeting for a user. "
+        "Requires a username - if not provided, will request it via elicitation."
+    ),
+    tags={"test", "elicitation"},
+)
+async def get_user_greeting(ctx: Context, username: str | None = None) -> dict:
+    """
+    Get a personalized greeting for a user.
+
+    This tool demonstrates FastMCP's built-in elicitation by requiring a username parameter.
+    If username is not provided, it uses ctx.elicit() to request it from the user.
+
+    Args:
+        ctx: FastMCP context (automatically injected)
+        username: The username to greet. If None, elicitation will be triggered.
+
+    Returns
+    -------
+        Dictionary with greeting message or error if elicitation was declined/cancelled
+    """
+    if not username:
+        # Use elicitation to request username from the client
+        try:
+            result = await ctx.elicit(
+                message="Username is required to generate a personalized greeting",
+                response_type=str,
+            )
+
+            if isinstance(result, AcceptedElicitation):
+                username = result.data
+            elif isinstance(result, DeclinedElicitation):
+                return {
+                    "status": "error",
+                    "error": "Username declined by user",
+                    "message": "Cannot generate greeting without username",
+                }
+            elif isinstance(result, CancelledElicitation):
+                return {
+                    "status": "error",
+                    "error": "Operation cancelled",
+                    "message": "Greeting request was cancelled",
+                }
+        except Exception:
+            # Elicitation not supported by client - return graceful skip
+            return {
+                "status": "skipped",
+                "message": (
+                    "Elicitation not supported by client. "
+                    "Username parameter is required when client does not support elicitation."
+                ),
+                "elicitation_supported": False,
+            }
+
+    return {
+        "status": "success",
+        "message": f"Hello, {username}! Welcome to the DataRobot MCP server.",
+        "username": username,
+    }

datarobot_genai/drmcp/test_utils/integration_mcp_server.py

@@ -26,6 +26,13 @@ from typing import Any
 
 from datarobot_genai.drmcp import create_mcp_server
 
+# Import elicitation test tool to register it with the MCP server
+try:
+    from datarobot_genai.drmcp.test_utils import elicitation_test_tool  # noqa: F401
+except ImportError:
+    # Test utils not available (e.g., running in production)
+    pass
+
 # Import user components (will be used conditionally)
 try:
     from app.core.server_lifecycle import ServerLifecycle  # type: ignore  # noqa: F401

datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

@@ -15,6 +15,7 @@ import asyncio
 import os
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
+from typing import Any
 
 import aiohttp
 from aiohttp import ClientSession as HttpClientSession
@@ -78,6 +79,7 @@ def get_headers() -> dict[str, str]:
 @asynccontextmanager
 async def ete_test_mcp_session(
     additional_headers: dict[str, str] | None = None,
+    elicitation_callback: Any | None = None,
 ) -> AsyncGenerator[ClientSession, None]:
     """Create an MCP session for each test.
 
@@ -85,6 +87,10 @@ async def ete_test_mcp_session(
     ----------
     additional_headers : dict[str, str], optional
         Additional headers to include in the MCP session (e.g., auth headers for testing).
+    elicitation_callback : callable, optional
+        Callback function to handle elicitation requests from the server.
+        The callback should have signature:
+        async def callback(context, params: ElicitRequestParams) -> ElicitResult
     """
     try:
         headers = get_headers()
@@ -96,7 +102,9 @@ async def ete_test_mcp_session(
             write_stream,
             _,
         ):
-            async with ClientSession(read_stream, write_stream) as session:
+            async with ClientSession(
+                read_stream, write_stream, elicitation_callback=elicitation_callback
+            ) as session:
                 await asyncio.wait_for(session.initialize(), timeout=5)
                 yield session
     except asyncio.TimeoutError:

datarobot_genai/drmcp/test_utils/mcp_utils_integration.py

@@ -17,6 +17,7 @@ import contextlib
 import os
 from collections.abc import AsyncGenerator
 from pathlib import Path
+from typing import Any
 
 from mcp import ClientSession
 from mcp.client.stdio import StdioServerParameters
@@ -34,7 +35,12 @@ def integration_test_mcp_server_params() -> StdioServerParameters:
         or "https://test.datarobot.com/api/v2",
         "MCP_SERVER_LOG_LEVEL": os.environ.get("MCP_SERVER_LOG_LEVEL") or "WARNING",
         "APP_LOG_LEVEL": os.environ.get("APP_LOG_LEVEL") or "WARNING",
-        "OTEL_ENABLED": os.environ.get("OTEL_ENABLED") or "false",
+        # Disable all OTEL telemetry for integration tests
+        "OTEL_ENABLED": "false",
+        "OTEL_SDK_DISABLED": "true",
+        "OTEL_TRACES_EXPORTER": "none",
+        "OTEL_LOGS_EXPORTER": "none",
+        "OTEL_METRICS_EXPORTER": "none",
         "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP": os.environ.get(
             "MCP_SERVER_REGISTER_DYNAMIC_TOOLS_ON_STARTUP"
         )
@@ -64,7 +70,9 @@ def integration_test_mcp_server_params() -> StdioServerParameters:
 
 @contextlib.asynccontextmanager
 async def integration_test_mcp_session(
-    server_params: StdioServerParameters | None = None, timeout: int = 30
+    server_params: StdioServerParameters | None = None,
+    timeout: int = 30,
+    elicitation_callback: Any | None = None,
 ) -> AsyncGenerator[ClientSession, None]:
     """
     Create and connect a client for the MCP server as a context manager.
@@ -72,6 +80,7 @@ async def integration_test_mcp_session(
     Args:
         server_params: Parameters for configuring the server connection
         timeout: Timeout
+        elicitation_callback: Optional callback for handling elicitation requests
 
     Yields
     ------
@@ -86,8 +95,12 @@ async def integration_test_mcp_session(
 
     try:
         async with stdio_client(server_params) as (read_stream, write_stream):
-            async with ClientSession(read_stream, write_stream) as session:
-                await asyncio.wait_for(session.initialize(), timeout=timeout)
+            async with ClientSession(
+                read_stream, write_stream, elicitation_callback=elicitation_callback
+            ) as session:
+                init_result = await asyncio.wait_for(session.initialize(), timeout=timeout)
+                # Store the init result on the session for tests that need to inspect capabilities
+                session._init_result = init_result  # type: ignore[attr-defined]
                 yield session
 
     except asyncio.TimeoutError:

datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+from ast import literal_eval
 from typing import Any
 
 import openai
@@ -44,12 +45,39 @@ class LLMResponse:
 
 
 class LLMMCPClient:
-    """Client for interacting with LLMs via MCP."""
+    """
+    Client for interacting with LLMs via MCP.
 
-    def __init__(self, config: str):
-        """Initialize the LLM MCP client."""
+    Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
+    Tools using FastMCP's built-in elicitation will work automatically.
+    """
+
+    def __init__(
+        self,
+        config: str,
+    ):
+        """
+        Initialize the LLM MCP client.
+
+        Args:
+            config: Configuration string or dict with:
+                - openai_api_key: OpenAI API key
+                - openai_api_base: Optional Azure OpenAI endpoint
+                - openai_api_deployment_id: Optional Azure deployment ID
+                - openai_api_version: Optional Azure API version
+                - model: Model name (default: "gpt-3.5-turbo")
+                - save_llm_responses: Whether to save responses (default: True)
+        """
         # Parse config string to extract parameters
-        config_dict = eval(config) if isinstance(config, str) else config
+        if isinstance(config, str):
+            # Try JSON first (safer), fall back to literal_eval for Python dict strings
+            try:
+                config_dict = json.loads(config)
+            except json.JSONDecodeError:
+                # Fall back to literal_eval for Python dict literal strings
+                config_dict = literal_eval(config)
+        else:
+            config_dict = config
 
         openai_api_key = config_dict.get("openai_api_key")
         openai_api_base = config_dict.get("openai_api_base")
@@ -93,7 +121,21 @@ class LLMMCPClient:
     async def _call_mcp_tool(
         self, tool_name: str, parameters: dict[str, Any], mcp_session: ClientSession
     ) -> str:
-        """Call an MCP tool and return the result as a string."""
+        """
+        Call an MCP tool and return the result as a string.
+
+        Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
+        Tools using FastMCP's built-in elicitation will work automatically.
+
+        Args:
+            tool_name: Name of the tool to call
+            parameters: Parameters to pass to the tool
+            mcp_session: MCP client session
+
+        Returns
+        -------
+            Result text from the tool call
+        """
         result: CallToolResult = await mcp_session.call_tool(tool_name, parameters)
         content = (
             result.content[0].text
@@ -177,7 +219,26 @@ class LLMMCPClient:
     async def process_prompt_with_mcp_support(
         self, prompt: str, mcp_session: ClientSession, output_file_name: str = ""
     ) -> LLMResponse:
-        """Process a prompt with MCP tool support."""
+        """
+        Process a prompt with MCP tool support and elicitation handling.
+
+        This method:
+        1. Adds MCP tools to available tools
+        2. Sends prompt to LLM
+        3. Processes tool calls
+        4. Continues until LLM provides final response
+
+        Note: Elicitation is handled at the protocol level by FastMCP's ctx.elicit().
+
+        Args:
+            prompt: User prompt
+            mcp_session: MCP client session
+            output_file_name: Optional file name to save response
+
+        Returns
+        -------
+            LLMResponse with content, tool calls, and tool results
+        """
         # Add MCP tools to available tools
         await self._add_mcp_tool_to_available_tools(mcp_session)
 
@@ -191,8 +252,10 @@ class LLMMCPClient:
                 "content": (
                     "You are a helpful AI assistant that can use tools to help users. "
                     "If you need more information to provide a complete response, you can make "
-                    "multiple tool calls. When dealing with file paths, use them as raw paths "
-                    "without converting to file:// URLs."
+                    "multiple tool calls or ask the user for more info, but prefer tool calls "
+                    "when possible. "
+                    "When dealing with file paths, use them as raw paths without converting "
+                    "to file:// URLs."
                 ),
             },
             {"role": "user", "content": prompt},

datarobot_genai/drmcp/test_utils/test_interactive.py

@@ -0,0 +1,205 @@
+#!/usr/bin/env python3
+
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Interactive MCP Client Test Script.
+
+This script allows you to test arbitrary commands with the MCP server
+using an LLM agent that can decide which tools to call.
+
+Supports elicitation - when tools require user input (like authentication tokens),
+the script will prompt you interactively.
+"""
+
+import asyncio
+import json
+import os
+import sys
+import traceback
+from pathlib import Path
+from typing import Any
+
+from dotenv import load_dotenv
+from mcp import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
+from mcp.shared.context import RequestContext
+from mcp.types import ElicitRequestParams
+from mcp.types import ElicitResult
+
+from datarobot_genai.drmcp import get_dr_mcp_server_url
+from datarobot_genai.drmcp import get_headers
+from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import LLMMCPClient
+from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import LLMResponse
+from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import ToolCall
+
+# Re-export for backwards compatibility
+__all__ = ["LLMMCPClient", "LLMResponse", "ToolCall", "test_mcp_interactive"]
+
+
+async def test_mcp_interactive() -> None:
+    """Test the MCP server interactively with LLM agent."""
+    # Check for required environment variables
+    openai_api_key = os.environ.get("OPENAI_API_KEY")
+    if not openai_api_key:
+        print("❌ Error: OPENAI_API_KEY environment variable is required")
+        print("Please set it in your .env file or export it")
+        return
+
+    # Optional Azure OpenAI settings
+    openai_api_base = os.environ.get("OPENAI_API_BASE")
+    openai_api_deployment_id = os.environ.get("OPENAI_API_DEPLOYMENT_ID")
+    openai_api_version = os.environ.get("OPENAI_API_VERSION")
+
+    print("🤖 Initializing LLM MCP Client...")
+
+    # Initialize the LLM client with elicitation handler
+    config = {
+        "openai_api_key": openai_api_key,
+        "openai_api_base": openai_api_base,
+        "openai_api_deployment_id": openai_api_deployment_id,
+        "openai_api_version": openai_api_version,
+        "save_llm_responses": False,
+    }
+
+    llm_client = LLMMCPClient(str(config))
+
+    # Get MCP server URL
+    mcp_server_url = get_dr_mcp_server_url()
+    if not mcp_server_url:
+        print("❌ Error: MCP server URL is not configured")
+        print("Please set DR_MCP_SERVER_URL environment variable or run: task test-interactive")
+        return
+
+    print(f"🔗 Connecting to MCP server at: {mcp_server_url}")
+
+    # Elicitation handler: prompt user for required values
+    async def elicitation_handler(
+        context: RequestContext[ClientSession, Any], params: ElicitRequestParams
+    ) -> ElicitResult:
+        print(f"\n📋 Elicitation Request: {params.message}")
+        if params.requestedSchema:
+            print(f"   Schema: {params.requestedSchema}")
+
+        while True:
+            try:
+                response = input("   Enter value (or 'decline'/'cancel'): ").strip()
+            except (EOFError, KeyboardInterrupt):
+                return ElicitResult(action="cancel")
+
+            if response.lower() == "decline":
+                return ElicitResult(action="decline")
+            if response.lower() == "cancel":
+                return ElicitResult(action="cancel")
+            if response:
+                return ElicitResult(action="accept", content={"value": response})
+            print("   Please enter a value or 'decline'/'cancel'")
+
+    try:
+        async with streamablehttp_client(
+            url=mcp_server_url,
+            headers=get_headers(),
+        ) as (read_stream, write_stream, _):
+            async with ClientSession(
+                read_stream,
+                write_stream,
+                elicitation_callback=elicitation_handler,
+            ) as session:
+                await session.initialize()
+
+                print("✅ Connected to MCP server!")
+                print("📋 Available tools:")
+
+                tools_result = await session.list_tools()
+                for i, tool in enumerate(tools_result.tools, 1):
+                    print(f"  {i}. {tool.name}: {tool.description}")
+
+                print("\n" + "=" * 60)
+                print("🎯 Interactive Testing Mode")
+                print("=" * 60)
+                print("Type your questions/commands. The AI will decide which tools to use.")
+                print("If a tool requires additional information, you will be prompted.")
+                print("Type 'quit' or 'exit' to stop.")
+                print()
+
+                while True:
+                    try:
+                        user_input = input("🤔 You: ").strip()
+
+                        if user_input.lower() in ["quit", "exit", "q"]:
+                            print("👋 Goodbye!")
+                            break
+
+                        if not user_input:
+                            continue
+                    except (EOFError, KeyboardInterrupt):
+                        print("\n👋 Goodbye!")
+                        break
+
+                    print("🤖 AI is thinking...")
+
+                    response = await llm_client.process_prompt_with_mcp_support(
+                        prompt=user_input,
+                        mcp_session=session,
+                    )
+
+                    print("\n🤖 AI Response:")
+                    print("-" * 40)
+                    print(response.content)
+
+                    if response.tool_calls:
+                        print("\n🔧 Tools Used:")
+                        for i, tool_call in enumerate(response.tool_calls, 1):
+                            print(f"  {i}. {tool_call.tool_name}")
+                            print(f"     Parameters: {tool_call.parameters}")
+                            print(f"     Reasoning: {tool_call.reasoning}")
+
+                            if i <= len(response.tool_results):
+                                result = response.tool_results[i - 1]
+                                try:
+                                    result_data = json.loads(result)
+                                    if result_data.get("status") == "error":
+                                        error_msg = result_data.get("error", "Unknown error")
+                                        print(f"     ❌ Error: {error_msg}")
+                                    elif result_data.get("status") == "success":
+                                        print("     ✅ Success")
+                                except json.JSONDecodeError:
+                                    if len(result) > 100:
+                                        print(f"     Result: {result[:100]}...")
+                                    else:
+                                        print(f"     Result: {result}")
+
+                    print("\n" + "=" * 60)
+    except Exception as e:
+        print(f"❌ Connection Error: {e}")
+        print(f"   Server URL: {mcp_server_url}")
+        traceback.print_exc()
+        return
+
+
+if __name__ == "__main__":
+    # Ensure we're in the right directory
+    if not Path("src").exists():
+        print("❌ Error: Please run this script from the project root")
+        sys.exit(1)
+
+    # Load environment variables from .env file
+    print("📄 Loading environment variables...")
+    load_dotenv()
+
+    print("🚀 Starting Interactive MCP Client Test")
+    print("Make sure the MCP server is running with: task drmcp-dev")
+    print()
+
+    asyncio.run(test_mcp_interactive())

datarobot_genai/drmcp/test_utils/tool_base_ete.py

@@ -116,28 +116,30 @@ class ToolBaseE2E:
                     f"Should have called {test_expectations.tool_calls_expected[i].name} tool, but "
                     f"got: {tool_call.tool_name}"
                 )
-            assert tool_call.parameters == test_expectations.tool_calls_expected[i].parameters, (
-                f"Should have called {tool_call.tool_name} tool with the correct parameters, but "
-                f"got: {tool_call.parameters}"
-            )
-            if test_expectations.tool_calls_expected[i].result != SHOULD_NOT_BE_EMPTY:
-                expected_result = test_expectations.tool_calls_expected[i].result
-                if isinstance(expected_result, str):
-                    assert expected_result in response.tool_results[i], (
-                        f"Should have called {tool_call.tool_name} tool with the correct result, "
-                        f"but got: {response.tool_results[i]}"
-                    )
+                assert (
+                    tool_call.parameters == test_expectations.tool_calls_expected[i].parameters
+                ), (
+                    f"Should have called {tool_call.tool_name} tool with the correct parameters, "
+                    f"but got: {tool_call.parameters}"
+                )
+                if test_expectations.tool_calls_expected[i].result != SHOULD_NOT_BE_EMPTY:
+                    expected_result = test_expectations.tool_calls_expected[i].result
+                    if isinstance(expected_result, str):
+                        assert expected_result in response.tool_results[i], (
+                            f"Should have called {tool_call.tool_name} tool with the correct "
+                            f"result, but got: {response.tool_results[i]}"
+                        )
+                    else:
+                        actual_result = json.loads(response.tool_results[i])
+                        assert _check_dict_has_keys(expected_result, actual_result), (
+                            f"Should have called {tool_call.tool_name} tool with the correct "
+                            f"result structure, but got: {response.tool_results[i]}"
+                        )
                 else:
-                    actual_result = json.loads(response.tool_results[i])
-                    assert _check_dict_has_keys(expected_result, actual_result), (
-                        f"Should have called {tool_call.tool_name} tool with the correct result "
-                        f"structure, but got: {response.tool_results[i]}"
+                    assert len(response.tool_results[i]) > 0, (
+                        f"Should have called {tool_call.tool_name} tool with non-empty result, but "
+                        f"got: {response.tool_results[i]}"
                     )
-            else:
-                assert len(response.tool_results[i]) > 0, (
-                    f"Should have called {tool_call.tool_name} tool with non-empty result, but "
-                    f"got: {response.tool_results[i]}"
-                )
 
         # Verify LLM provided comprehensive response
         assert len(response.content) > 100, "LLM should provide detailed response"

datarobot_genai/drmcp/tools/clients/confluence.py

@@ -59,6 +59,22 @@ class ConfluencePage(BaseModel):
         }
 
 
+class ConfluenceComment(BaseModel):
+    """Pydantic model for Confluence comment."""
+
+    comment_id: str = Field(..., description="The unique comment ID")
+    page_id: str = Field(..., description="The page ID where the comment was added")
+    body: str = Field(..., description="Comment content in storage format")
+
+    def as_flat_dict(self) -> dict[str, Any]:
+        """Return a flat dictionary representation of the comment."""
+        return {
+            "comment_id": self.comment_id,
+            "page_id": self.page_id,
+            "body": self.body,
+        }
+
+
 class ConfluenceClient:
     """
     Client for interacting with Confluence API using OAuth access token.
@@ -294,6 +310,78 @@ class ConfluenceClient:
 
         return self._parse_response(response.json())
 
+    def _parse_comment_response(self, data: dict, page_id: str) -> ConfluenceComment:
+        """Parse API response into ConfluenceComment."""
+        body_content = ""
+        body = data.get("body", {})
+        if isinstance(body, dict):
+            storage = body.get("storage", {})
+            if isinstance(storage, dict):
+                body_content = storage.get("value", "")
+
+        return ConfluenceComment(
+            comment_id=str(data.get("id", "")),
+            page_id=page_id,
+            body=body_content,
+        )
+
+    async def add_comment(self, page_id: str, comment_body: str) -> ConfluenceComment:
+        """
+        Add a comment to a Confluence page.
+
+        Args:
+            page_id: The numeric page ID where the comment will be added
+            comment_body: The text content of the comment
+
+        Returns
+        -------
+            ConfluenceComment with the created comment data
+
+        Raises
+        ------
+            ConfluenceError: If page not found, permission denied, or invalid content
+            httpx.HTTPStatusError: If the API request fails with unexpected status
+        """
+        cloud_id = await self._get_cloud_id()
+        url = f"{ATLASSIAN_API_BASE}/ex/confluence/{cloud_id}/wiki/rest/api/content"
+
+        payload: dict[str, Any] = {
+            "type": "comment",
+            "container": {"id": page_id, "type": "page"},
+            "body": {
+                "storage": {
+                    "value": comment_body,
+                    "representation": "storage",
+                }
+            },
+        }
+
+        response = await self._client.post(url, json=payload)
+
+        if response.status_code == HTTPStatus.NOT_FOUND:
+            error_msg = self._extract_error_message(response)
+            raise ConfluenceError(
+                f"Page with ID '{page_id}' not found: {error_msg}",
+                status_code=404,
+            )
+
+        if response.status_code == HTTPStatus.FORBIDDEN:
+            raise ConfluenceError(
+                f"Permission denied: you don't have access to add comments to page '{page_id}'",
+                status_code=403,
+            )
+
+        if response.status_code == HTTPStatus.BAD_REQUEST:
+            error_msg = self._extract_error_message(response)
+            raise ConfluenceError(f"Invalid request: {error_msg}", status_code=400)
+
+        if response.status_code == HTTPStatus.TOO_MANY_REQUESTS:
+            raise ConfluenceError("Rate limit exceeded. Please try again later.", status_code=429)
+
+        response.raise_for_status()
+
+        return self._parse_comment_response(response.json(), page_id)
+
     async def __aenter__(self) -> "ConfluenceClient":
         """Async context manager entry."""
         return self

datarobot_genai/drmcp/tools/confluence/tools.py

@@ -138,3 +138,51 @@ async def confluence_create_page(
         content=f"New page '{title}' created successfully in space '{space_key}'.",
         structured_content={"new_page_id": page_response.page_id, "title": page_response.title},
     )
+
+
+@dr_mcp_tool(tags={"confluence", "write", "add", "comment"})
+async def confluence_add_comment(
+    *,
+    page_id: Annotated[str, "The numeric ID of the page where the comment will be added."],
+    comment_body: Annotated[str, "The text content of the comment."],
+) -> ToolResult:
+    """Add a new comment to a specified Confluence page for collaboration.
+
+    Use this tool to add comments to Confluence pages to facilitate collaboration
+    and discussion. Comments are added at the page level.
+
+    Usage:
+        - Add comment: page_id="856391684", comment_body="Great work on this documentation!"
+    """
+    if not page_id:
+        raise ToolError("Argument validation error: 'page_id' cannot be empty.")
+
+    if not comment_body:
+        raise ToolError("Argument validation error: 'comment_body' cannot be empty.")
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with ConfluenceClient(access_token) as client:
+            comment_response = await client.add_comment(
+                page_id=page_id,
+                comment_body=comment_body,
+            )
+    except ConfluenceError as e:
+        logger.error(f"Confluence error adding comment: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error adding comment to Confluence page: {e}")
+        raise ToolError(
+            f"An unexpected error occurred while adding comment to page '{page_id}': {str(e)}"
+        )
+
+    return ToolResult(
+        content=f"Comment added successfully to page ID {page_id}.",
+        structured_content={
+            "comment_id": comment_response.comment_id,
+            "page_id": page_id,
+        },
+    )

datarobot_genai/drmcp/tools/predictive/data.py

@@ -14,6 +14,10 @@
 
 import logging
 import os
+from typing import Annotated
+
+from fastmcp.exceptions import ToolError
+from fastmcp.tools.tool import ToolResult
 
 from datarobot_genai.drmcp.core.clients import get_sdk_client
 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
 logger = logging.getLogger(__name__)
 
 
-@dr_mcp_tool(tags={"data", "management", "upload"})
-async def upload_dataset_to_ai_catalog(file_path: str) -> str:
-    """
-    Upload a dataset to the DataRobot AI Catalog.
-
-    Args:
-        file_path: Path to the file to upload.
-
-    Returns
-    -------
-        A string summary of the upload result.
-    """
-    client = get_sdk_client()
+@dr_mcp_tool(tags={"predictive", "data", "write", "upload", "catalog"})
+async def upload_dataset_to_ai_catalog(
+    file_path: Annotated[str, "The path to the dataset file to upload."],
+) -> ToolResult:
+    """Upload a dataset to the DataRobot AI Catalog."""
     if not os.path.exists(file_path):
         logger.error(f"File not found: {file_path}")
-        return f"File not found: {file_path}"
+        raise ToolError(f"File not found: {file_path}")
+
+    client = get_sdk_client()
     catalog_item = client.Dataset.create_from_file(file_path)
-    logger.info(f"Successfully uploaded dataset: {catalog_item.id}")
-    return f"AI Catalog ID: {catalog_item.id}"
 
+    return ToolResult(
+        content=f"Successfully uploaded dataset: {catalog_item.id}.",
+        structured_content={
+            "id": catalog_item.id,
+            "name": catalog_item.name,
+            "status": catalog_item.status,
+        },
+    )
 
-@dr_mcp_tool(tags={"data", "management", "list"})
-async def list_ai_catalog_items() -> str:
-    """
-    List all AI Catalog items (datasets) for the authenticated user.
 
-    Returns
-    -------
-        A string summary of the AI Catalog items with their IDs and names.
-    """
+@dr_mcp_tool(tags={"predictive", "data", "read", "list", "catalog"})
+async def list_ai_catalog_items() -> ToolResult:
+    """List all AI Catalog items (datasets) for the authenticated user."""
     client = get_sdk_client()
     datasets = client.Dataset.list()
+
     if not datasets:
         logger.info("No AI Catalog items found")
-        return "No AI Catalog items found."
-    result = "\n".join(f"{ds.id}: {ds.name}" for ds in datasets)
-    logger.info(f"Found {len(datasets)} AI Catalog items")
-    return result
+        return ToolResult(
+            content="No AI Catalog items found.",
+            structured_content={"datasets": []},
+        )
+
+    return ToolResult(
+        content=f"Found {len(datasets)} AI Catalog items.",
+        structured_content={
+            "datasets": [{"id": ds.id, "name": ds.name} for ds in datasets],
+            "count": len(datasets),
+        },
+    )
 
 
 # from fastmcp import Context

{datarobot_genai-0.2.16.dist-info → datarobot_genai-0.2.18.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.16
+Version: 0.2.18
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.

{datarobot_genai-0.2.16.dist-info → datarobot_genai-0.2.18.dist-info}/RECORD

@@ -67,24 +67,26 @@ datarobot_genai/drmcp/core/memory_management/__init__.py,sha256=y4yapzp3KnFMzSR6
 datarobot_genai/drmcp/core/memory_management/manager.py,sha256=gmc_SQs12YQFMWl2UbfWR40QmLV9XuCnwPZgQwKWrbA,30552
 datarobot_genai/drmcp/core/memory_management/memory_tools.py,sha256=AxzpwOlldmhhDfKZcAxaGs7Xih2SCe0XbQuXX5nQczI,6397
 datarobot_genai/drmcp/test_utils/__init__.py,sha256=y4yapzp3KnFMzSR6HlNDS4uSuyNT7I1iPBvaCLsS0sU,577
-datarobot_genai/drmcp/test_utils/integration_mcp_server.py,sha256=MdoR7r3m9uT7crodyhY69yhkrM7Thpe__BBD9lB_2oA,3328
-datarobot_genai/drmcp/test_utils/mcp_utils_ete.py,sha256=rgZkPF26YCHX2FGppWE4v22l_NQ3kLSPSUimO0tD4nM,4402
-datarobot_genai/drmcp/test_utils/mcp_utils_integration.py,sha256=0sU29Khal0CelnHBDInyTRiuPKrFFbTbIomOoUbyMhs,3271
-datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py,sha256=TvTkDBcHscLDmqge9NhHxVo1ABtb0n4NmmG2318mQHU,9088
-datarobot_genai/drmcp/test_utils/tool_base_ete.py,sha256=-mKHBkGkyOKQCVS2LHFhSnRofIqJBbeAPRkwizBDtTg,6104
+datarobot_genai/drmcp/test_utils/elicitation_test_tool.py,sha256=UVKwy39nl3XcVAh6IATcN-cWL2bfrprgRQ7fbK82jTI,3287
+datarobot_genai/drmcp/test_utils/integration_mcp_server.py,sha256=YSk19tbaka_0ziqi7LoXie4SJs-cvi9-H00Go0ZtQWE,3575
+datarobot_genai/drmcp/test_utils/mcp_utils_ete.py,sha256=46rH0fYYmUj7ygf968iRbdSp5u95v23BEw3Ts_c431Y,4788
+datarobot_genai/drmcp/test_utils/mcp_utils_integration.py,sha256=sHA_BWtpgIAFp9IXiNkUeBartBMjLAauqkV9bYtCr-g,3874
+datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py,sha256=YgyqHK09MB-PBaqT34heqvmvYYFtLpzzSJt7xuTJmDg,11224
+datarobot_genai/drmcp/test_utils/test_interactive.py,sha256=guXvR8q2H6VUdmvIjEJcElQJCC6lQ-oTrzbD2EkHeCs,8025
+datarobot_genai/drmcp/test_utils/tool_base_ete.py,sha256=wmI-xcL0rSr56-ZoGNB8np0CZHAU583o8-Kw7fRwtMw,6232
 datarobot_genai/drmcp/test_utils/utils.py,sha256=esGKFv8aO31-Qg3owayeWp32BYe1CdYOEutjjdbweCw,3048
 datarobot_genai/drmcp/tools/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/clients/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/clients/atlassian.py,sha256=__M_uz7FrcbKCYRzeMn24DCEYD6OmFx_LuywHCxgXsA,6472
-datarobot_genai/drmcp/tools/clients/confluence.py,sha256=DF6TIGJfR3Lh-D_x66cDNkvOTS8gxL6bVhHRtcP0LKw,10493
+datarobot_genai/drmcp/tools/clients/confluence.py,sha256=gDzy8t5t3b1mwEr-CuZ5BwXXQ52AXke8J_Ra7i_8T1g,13692
 datarobot_genai/drmcp/tools/clients/jira.py,sha256=Rm91JAyrNIqxu66-9rU1YqoRXVnWbEy-Ahvy6f6HlVg,9823
 datarobot_genai/drmcp/tools/clients/s3.py,sha256=GmwzvurFdNfvxOooA8g5S4osRysHYU0S9ypg_177Glg,953
 datarobot_genai/drmcp/tools/confluence/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
-datarobot_genai/drmcp/tools/confluence/tools.py,sha256=iqX7CR57WCXsQxHQCsfPL_Q78QjN9YZv3uIQbTMfYAg,5459
+datarobot_genai/drmcp/tools/confluence/tools.py,sha256=jSF7yXGFqqlMcavkRIY4HbMxb7tCeunA2ST41wa2vGI,7219
 datarobot_genai/drmcp/tools/jira/__init__.py,sha256=0kq9vMkF7EBsS6lkEdiLibmUrghTQqosHbZ5k-V9a5g,578
 datarobot_genai/drmcp/tools/jira/tools.py,sha256=dfkqTU2HH-7n44hX80ODFacKq0p0LOchFcZtIIKFNMM,9687
 datarobot_genai/drmcp/tools/predictive/__init__.py,sha256=WuOHlNNEpEmcF7gVnhckruJRKU2qtmJLE3E7zoCGLDo,1030
-datarobot_genai/drmcp/tools/predictive/data.py,sha256=k4EJxJrl8DYVGVfJ0DM4YTfnZlC_K3OUHZ0eRUzfluI,3165
+datarobot_genai/drmcp/tools/predictive/data.py,sha256=nc2OdkcZs9_L_o466tkaKaNvuhmvqcJKJ3nrA4H8B_c,3488
 datarobot_genai/drmcp/tools/predictive/deployment.py,sha256=lm02Ayuo11L1hP41fgi3QpR1Eyty-Wc16rM0c8SgliM,3277
 datarobot_genai/drmcp/tools/predictive/deployment_info.py,sha256=BGEF_dmbxOBJR0n1Tt9TO2-iNTQSBTr-oQUyaxLZ0ZI,15297
 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
 datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
 datarobot_genai/nat/datarobot_mcp_client.py,sha256=35FzilxNp4VqwBYI0NsOc91-xZm1C-AzWqrOdDy962A,9612
 datarobot_genai/nat/helpers.py,sha256=Q7E3ADZdtFfS8E6OQPyw2wgA6laQ58N3bhLj5CBWwJs,3265
-datarobot_genai-0.2.16.dist-info/METADATA,sha256=lTI4a5PfXD7nIUcmAmzENdRiaIfgdaLoruSwq95NYU8,6301
-datarobot_genai-0.2.16.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.16.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
-datarobot_genai-0.2.16.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.16.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.16.dist-info/RECORD,,
+datarobot_genai-0.2.18.dist-info/METADATA,sha256=VQhojO4HSwkoFiHSGH9AdedWp-JV2RkRH-IewqJa38c,6301
+datarobot_genai-0.2.18.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.2.18.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
+datarobot_genai-0.2.18.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.2.18.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.2.18.dist-info/RECORD,,