datarobot-genai 0.2.3__py3-none-any.whl → 0.2.11__py3-none-any.whl

datarobot_genai/drmcp/core/config.py

@@ -197,6 +197,54 @@ class MCPServerConfig(BaseSettings):
         description="Enable/disable predictive tools",
     )
 
+    # Jira tools
+    enable_jira_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_JIRA_TOOLS",
+            "ENABLE_JIRA_TOOLS",
+        ),
+        description="Enable/disable Jira tools",
+    )
+    is_jira_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
+            "IS_JIRA_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Jira OAuth provider is configured for Jira integration",
+    )
+
+    @property
+    def is_jira_oauth_configured(self) -> bool:
+        return self.is_jira_oauth_provider_configured or bool(
+            os.getenv("JIRA_CLIENT_ID") and os.getenv("JIRA_CLIENT_SECRET")
+        )
+
+    # Confluence tools
+    enable_confluence_tools: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "ENABLE_CONFLUENCE_TOOLS",
+            "ENABLE_CONFLUENCE_TOOLS",
+        ),
+        description="Enable/disable Confluence tools",
+    )
+    is_confluence_oauth_provider_configured: bool = Field(
+        default=False,
+        validation_alias=AliasChoices(
+            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
+            "IS_CONFLUENCE_OAUTH_PROVIDER_CONFIGURED",
+        ),
+        description="Whether Confluence OAuth provider is configured for Confluence integration",
+    )
+
+    @property
+    def is_confluence_oauth_configured(self) -> bool:
+        return self.is_confluence_oauth_provider_configured or bool(
+            os.getenv("CONFLUENCE_CLIENT_ID") and os.getenv("CONFLUENCE_CLIENT_SECRET")
+        )
+
     @field_validator(
         "otel_attributes",
         mode="before",
@@ -220,6 +268,10 @@ class MCPServerConfig(BaseSettings):
         "tool_registration_duplicate_behavior",
         "mcp_server_register_dynamic_prompts_on_startup",
         "enable_predictive_tools",
+        "enable_jira_tools",
+        "is_jira_oauth_provider_configured",
+        "enable_confluence_tools",
+        "is_confluence_oauth_provider_configured",
         mode="before",
     )
     @classmethod

datarobot_genai/drmcp/core/dr_mcp_server.py

@@ -40,6 +40,8 @@ from .routes_utils import prefix_mount_path
 from .server_life_cycle import BaseServerLifecycle
 from .telemetry import OtelASGIMiddleware
 from .telemetry import initialize_telemetry
+from .tool_config import TOOL_CONFIGS
+from .tool_config import is_tool_enabled
 
 
 def _import_modules_from_dir(
@@ -142,11 +144,12 @@ class DataRobotMCPServer:
 
         # Load static tools modules
         base_dir = os.path.dirname(os.path.dirname(__file__))
-        if self._config.enable_predictive_tools:
-            _import_modules_from_dir(
-                os.path.join(base_dir, "tools", "predictive"),
-                "datarobot_genai.drmcp.tools.predictive",
-            )
+        for tool_type, tool_config in TOOL_CONFIGS.items():
+            if is_tool_enabled(tool_type, self._config):
+                _import_modules_from_dir(
+                    os.path.join(base_dir, "tools", tool_config["directory"]),
+                    tool_config["package_prefix"],
+                )
 
         # Load memory management tools if available
         if self._memory_manager:
@@ -213,6 +216,9 @@ class DataRobotMCPServer:
                 self._logger.info("Registering dynamic prompts from prompt management...")
                 asyncio.run(register_prompts_from_datarobot_prompt_management())
 
+            # Execute pre-server start actions
+            asyncio.run(self._lifecycle.pre_server_start(self._mcp))
+
             # List registered tools, prompts, and resources before starting server
             tools = asyncio.run(self._mcp._list_tools_mcp())
             prompts = asyncio.run(self._mcp._list_prompts_mcp())
@@ -232,9 +238,6 @@ class DataRobotMCPServer:
             for resource in resources:
                 self._logger.info(f" > {resource.name}")
 
-            # Execute pre-server start actions
-            asyncio.run(self._lifecycle.pre_server_start(self._mcp))
-
             # Create event loop for async operations
             loop = asyncio.new_event_loop()
             asyncio.set_event_loop(loop)

datarobot_genai/drmcp/core/tool_config.py

@@ -0,0 +1,95 @@
+# 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.
+
+"""Tool configuration and enablement logic."""
+
+from collections.abc import Callable
+from enum import Enum
+from typing import TYPE_CHECKING
+from typing import TypedDict
+
+if TYPE_CHECKING:
+    from .config import MCPServerConfig
+
+
+class ToolType(str, Enum):
+    """Enumeration of available tool types."""
+
+    PREDICTIVE = "predictive"
+    JIRA = "jira"
+    CONFLUENCE = "confluence"
+
+
+class ToolConfig(TypedDict):
+    """Configuration for a tool type."""
+
+    name: str
+    oauth_check: Callable[["MCPServerConfig"], bool] | None
+    directory: str
+    package_prefix: str
+    config_field_name: str  # Name of the config field (e.g., "enable_predictive_tools")
+
+
+# Tool configuration registry
+TOOL_CONFIGS: dict[ToolType, ToolConfig] = {
+    ToolType.PREDICTIVE: ToolConfig(
+        name="predictive",
+        oauth_check=None,
+        directory="predictive",
+        package_prefix="datarobot_genai.drmcp.tools.predictive",
+        config_field_name="enable_predictive_tools",
+    ),
+    ToolType.JIRA: ToolConfig(
+        name="jira",
+        oauth_check=lambda config: config.is_jira_oauth_configured,
+        directory="jira",
+        package_prefix="datarobot_genai.drmcp.tools.jira",
+        config_field_name="enable_jira_tools",
+    ),
+    ToolType.CONFLUENCE: ToolConfig(
+        name="confluence",
+        oauth_check=lambda config: config.is_confluence_oauth_configured,
+        directory="confluence",
+        package_prefix="datarobot_genai.drmcp.tools.confluence",
+        config_field_name="enable_confluence_tools",
+    ),
+}
+
+
+def get_tool_enable_config_name(tool_type: ToolType) -> str:
+    """Get the configuration field name for enabling a tool."""
+    return TOOL_CONFIGS[tool_type]["config_field_name"]
+
+
+def is_tool_enabled(tool_type: ToolType, config: "MCPServerConfig") -> bool:
+    """
+    Check if a tool is enabled based on configuration.
+
+    Args:
+        tool_type: The type of tool to check
+        config: The server configuration
+
+    Returns
+    -------
+        True if the tool is enabled, False otherwise
+    """
+    tool_config = TOOL_CONFIGS[tool_type]
+    enable_config_name = tool_config["config_field_name"]
+    is_enabled = getattr(config, enable_config_name)
+
+    # If tool is enabled, check OAuth requirements if needed
+    if is_enabled and tool_config["oauth_check"] is not None:
+        return tool_config["oauth_check"](config)
+
+    return is_enabled

datarobot_genai/drmcp/test_utils/openai_llm_mcp_client.py

@@ -95,11 +95,16 @@ class LLMMCPClient:
     ) -> str:
         """Call an MCP tool and return the result as a string."""
         result: CallToolResult = await mcp_session.call_tool(tool_name, parameters)
-        return (
+        content = (
             result.content[0].text
             if result.content and isinstance(result.content[0], TextContent)
             else str(result.content)
         )
+        if result.structuredContent is not None:
+            structured_content = json.dumps(result.structuredContent)
+        else:
+            structured_content = ""
+        return f"Content: {content}\nStructured content: {structured_content}"
 
     async def _process_tool_calls(
         self,

datarobot_genai/drmcp/tools/clients/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+

datarobot_genai/drmcp/tools/clients/atlassian.py

@@ -0,0 +1,188 @@
+# 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.
+
+"""Atlassian API client utilities for OAuth and cloud ID management."""
+
+import logging
+from typing import Any
+from typing import Literal
+
+import httpx
+from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
+from fastmcp.exceptions import ToolError
+
+from datarobot_genai.drmcp.core.auth import get_access_token
+
+logger = logging.getLogger(__name__)
+
+# Atlassian Cloud API base URL
+ATLASSIAN_API_BASE = "https://api.atlassian.com"
+
+# API endpoint paths
+OAUTH_ACCESSIBLE_RESOURCES_PATH = "/oauth/token/accessible-resources"
+
+# Supported Atlassian service types
+AtlassianServiceType = Literal["jira", "confluence"]
+
+
+async def get_atlassian_access_token() -> str | ToolError:
+    """
+    Get Atlassian OAuth access token with error handling.
+
+    Returns
+    -------
+        Access token string on success, ToolError on failure
+
+    Example:
+        ```python
+        token = await get_atlassian_access_token()
+        if isinstance(token, ToolError):
+            # Handle error
+            return token
+        # Use token
+        ```
+    """
+    try:
+        access_token = await get_access_token("atlassian")
+        if not access_token:
+            logger.warning("Empty access token received")
+            return ToolError("Received empty access token. Please complete the OAuth flow.")
+        return access_token
+    except OAuthServiceClientErr as e:
+        logger.error(f"OAuth client error: {e}", exc_info=True)
+        return ToolError(
+            "Could not obtain access token for Atlassian. Make sure the OAuth "
+            "permission was granted for the application to act on your behalf."
+        )
+    except Exception as e:
+        logger.error(f"Unexpected error obtaining access token: {e}", exc_info=True)
+        return ToolError("An unexpected error occurred while obtaining access token for Atlassian.")
+
+
+def _find_resource_by_service(
+    resources: list[dict[str, Any]], service_type: str
+) -> dict[str, Any] | None:
+    """
+    Find a resource that matches the specified service type.
+
+    Args:
+        resources: List of accessible resources from Atlassian API
+        service_type: Service type to filter by (e.g., "jira", "confluence")
+
+    Returns
+    -------
+        Resource dictionary if found, None otherwise
+    """
+    service_lower = service_type.lower()
+    for resource in resources:
+        if not resource.get("id"):
+            continue
+        scopes = resource.get("scopes", [])
+        if any(service_lower in scope.lower() for scope in scopes):
+            return resource
+    return None
+
+
+def _find_first_resource_with_id(resources: list[dict[str, Any]]) -> dict[str, Any] | None:
+    """
+    Find the first resource that has an ID.
+
+    Args:
+        resources: List of accessible resources from Atlassian API
+
+    Returns
+    -------
+        Resource dictionary if found, None otherwise
+    """
+    for resource in resources:
+        if resource.get("id"):
+            return resource
+    return None
+
+
+async def get_atlassian_cloud_id(
+    client: httpx.AsyncClient,
+    service_type: AtlassianServiceType | None = None,
+) -> str:
+    """
+    Get the cloud ID for the authenticated Atlassian instance.
+
+    According to Atlassian OAuth 2.0 documentation, API calls should use:
+    https://api.atlassian.com/ex/{service}/{cloudId}/rest/api/3/...
+
+    Args:
+        client: HTTP client with authentication headers configured
+        service_type: Optional service type to filter by (e.g., "jira", "confluence").
+            If provided, will prioritize resources matching this service type.
+
+    Returns
+    -------
+        Cloud ID string for the Atlassian instance
+
+    Raises
+    ------
+        ValueError: If cloud ID cannot be retrieved due to:
+            - No accessible resources found
+            - No cloud ID found in resources
+            - Authentication failure (401)
+            - HTTP request failure
+
+    Example:
+        ```python
+        client = httpx.AsyncClient(headers={"Authorization": f"Bearer {token}"})
+        cloud_id = await get_atlassian_cloud_id(client, service_type="jira")
+        ```
+    """
+    url = f"{ATLASSIAN_API_BASE}{OAUTH_ACCESSIBLE_RESOURCES_PATH}"
+
+    try:
+        response = await client.get(url)
+        response.raise_for_status()
+        resources = response.json()
+
+        if not resources:
+            raise ValueError(
+                "No accessible resources found. Ensure OAuth token has required scopes."
+            )
+
+        # If service_type is specified, try to find matching resource
+        if service_type:
+            resource = _find_resource_by_service(resources, service_type)
+            if resource:
+                cloud_id = resource["id"]
+                logger.debug(f"Using {service_type} cloud ID: {cloud_id}")
+                return cloud_id
+            logger.warning(
+                f"No {service_type} resource found, falling back to first available resource"
+            )
+
+        # Fallback: use the first resource with an ID
+        resource = _find_first_resource_with_id(resources)
+        if resource:
+            cloud_id = resource["id"]
+            logger.debug(f"Using cloud ID (fallback): {cloud_id}")
+            return cloud_id
+
+        raise ValueError("No cloud ID found in accessible resources")
+    except httpx.HTTPStatusError as e:
+        if e.response.status_code == 401:
+            raise ValueError(
+                "Authentication failed. Token may be expired. "
+                "Complete OAuth flow again: GET /oauth/atlassian/authorize"
+            ) from e
+        logger.error(f"HTTP error getting cloud ID: {e.response.status_code}")
+        raise ValueError(f"Failed to get Atlassian cloud ID: HTTP {e.response.status_code}") from e
+    except httpx.RequestError as e:
+        logger.error(f"Request error getting cloud ID: {e}")
+        raise ValueError("Failed to get Atlassian cloud ID: Network error") from e

datarobot_genai/drmcp/tools/clients/confluence.py

@@ -0,0 +1,196 @@
+# 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.
+
+"""Async client for interacting with Confluence Cloud REST API.
+
+At the moment of creating this client, official Confluence SDK is not supporting async.
+"""
+
+import logging
+from http import HTTPStatus
+from typing import Any
+
+import httpx
+from pydantic import BaseModel
+from pydantic import Field
+
+from .atlassian import ATLASSIAN_API_BASE
+from .atlassian import get_atlassian_cloud_id
+
+logger = logging.getLogger(__name__)
+
+
+class ConfluencePage(BaseModel):
+    """Pydantic model for Confluence page."""
+
+    page_id: str = Field(..., description="The unique page ID")
+    title: str = Field(..., description="Page title")
+    space_id: str = Field(..., description="Space ID where the page resides")
+    space_key: str | None = Field(None, description="Space key (if available)")
+    body: str = Field(..., description="Page content in storage format (HTML-like)")
+
+    def as_flat_dict(self) -> dict[str, Any]:
+        """Return a flat dictionary representation of the page."""
+        return {
+            "page_id": self.page_id,
+            "title": self.title,
+            "space_id": self.space_id,
+            "space_key": self.space_key,
+            "body": self.body,
+        }
+
+
+class ConfluenceClient:
+    """
+    Client for interacting with Confluence API using OAuth access token.
+
+    At the moment of creating this client, official Confluence SDK is not supporting async.
+    """
+
+    EXPAND_FIELDS = "body.storage,space"
+
+    def __init__(self, access_token: str) -> None:
+        """
+        Initialize Confluence client with access token.
+
+        Args:
+            access_token: OAuth access token for Atlassian API
+        """
+        self.access_token = access_token
+        self._client = httpx.AsyncClient(
+            headers={
+                "Authorization": f"Bearer {access_token}",
+                "Accept": "application/json",
+                "Content-Type": "application/json",
+            },
+            timeout=30.0,
+        )
+        self._cloud_id: str | None = None
+
+    async def _get_cloud_id(self) -> str:
+        """
+        Get the cloud ID for the authenticated Atlassian Confluence instance.
+
+        According to Atlassian OAuth 2.0 documentation, API calls should use:
+        https://api.atlassian.com/ex/confluence/{cloudId}/wiki/rest/api/...
+
+        Returns
+        -------
+            Cloud ID string
+
+        Raises
+        ------
+            ValueError: If cloud ID cannot be retrieved
+        """
+        if self._cloud_id:
+            return self._cloud_id
+
+        self._cloud_id = await get_atlassian_cloud_id(self._client, service_type="confluence")
+        return self._cloud_id
+
+    def _parse_response(self, data: dict) -> ConfluencePage:
+        """Parse API response into ConfluencePage."""
+        body_content = ""
+        body = data.get("body", {})
+        if isinstance(body, dict):
+            storage = body.get("storage", {})
+            if isinstance(storage, dict):
+                body_content = storage.get("value", "")
+
+        space = data.get("space", {})
+        space_key = space.get("key") if isinstance(space, dict) else None
+        space_id = space.get("id", "") if isinstance(space, dict) else data.get("spaceId", "")
+
+        return ConfluencePage(
+            page_id=str(data.get("id", "")),
+            title=data.get("title", ""),
+            space_id=str(space_id),
+            space_key=space_key,
+            body=body_content,
+        )
+
+    async def get_page_by_id(self, page_id: str) -> ConfluencePage:
+        """
+        Get a Confluence page by its ID.
+
+        Args:
+            page_id: The numeric page ID
+
+        Returns
+        -------
+            ConfluencePage with page data
+
+        Raises
+        ------
+            ValueError: If page is not found
+            httpx.HTTPStatusError: If the API request fails
+        """
+        cloud_id = await self._get_cloud_id()
+        url = f"{ATLASSIAN_API_BASE}/ex/confluence/{cloud_id}/wiki/rest/api/content/{page_id}"
+
+        response = await self._client.get(url, params={"expand": self.EXPAND_FIELDS})
+
+        if response.status_code == HTTPStatus.NOT_FOUND:
+            raise ValueError(f"Page with ID '{page_id}' not found")
+
+        response.raise_for_status()
+        return self._parse_response(response.json())
+
+    async def get_page_by_title(self, title: str, space_key: str) -> ConfluencePage:
+        """
+        Get a Confluence page by its title within a specific space.
+
+        Args:
+            title: The exact page title
+            space_key: The space key where the page resides
+
+        Returns
+        -------
+            ConfluencePage with page data
+
+        Raises
+        ------
+            ValueError: If the page is not found
+            httpx.HTTPStatusError: If the API request fails
+        """
+        cloud_id = await self._get_cloud_id()
+        url = f"{ATLASSIAN_API_BASE}/ex/confluence/{cloud_id}/wiki/rest/api/content"
+
+        response = await self._client.get(
+            url,
+            params={
+                "title": title,
+                "spaceKey": space_key,
+                "expand": self.EXPAND_FIELDS,
+            },
+        )
+        response.raise_for_status()
+
+        data = response.json()
+        results = data.get("results", [])
+
+        if not results:
+            raise ValueError(f"Page with title '{title}' not found in space '{space_key}'")
+
+        return self._parse_response(results[0])
+
+    async def __aenter__(self) -> "ConfluenceClient":
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any
+    ) -> None:
+        """Async context manager exit."""
+        await self._client.aclose()