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

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()

datarobot_genai/drmcp/tools/clients/jira.py

@@ -0,0 +1,147 @@
+# 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.
+
+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 _IssuePerson(BaseModel):
+    email_address: str = Field(alias="emailAddress")
+
+
+class _IssueStatus(BaseModel):
+    name: str
+
+
+class _IssueFields(BaseModel):
+    summary: str
+    status: _IssueStatus
+    reporter: _IssuePerson
+    assignee: _IssuePerson
+    created: str
+    updated: str
+
+
+class Issue(BaseModel):
+    id: str
+    key: str
+    fields: _IssueFields
+
+    def as_flat_dict(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "key": self.key,
+            "summary": self.fields.summary,
+            "reporterEmailAddress": self.fields.reporter.email_address,
+            "assigneeEmailAddress": self.fields.assignee.email_address,
+            "created": self.fields.created,
+            "updated": self.fields.updated,
+            "status": self.fields.status.name,
+        }
+
+
+class JiraClient:
+    """
+    Client for interacting with Jira API using OAuth access token.
+
+    At the moment of creating this client, official Jira SDK is not supporting async.
+    """
+
+    def __init__(self, access_token: str) -> None:
+        """
+        Initialize Jira 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 Jira instance.
+
+        According to Atlassian OAuth 2.0 documentation, API calls should use:
+        https://api.atlassian.com/ex/jira/{cloudId}/rest/api/3/...
+
+        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="jira")
+        return self._cloud_id
+
+    async def get_jira_issue(self, issue_key: str) -> Issue:
+        """
+        Get a Jira issue by its key.
+
+        Args:
+            issue_key: The key of the Jira issue, e.g., 'PROJ-123'
+
+        Returns
+        -------
+            Jira issue
+
+        Raises
+        ------
+            httpx.HTTPStatusError: If the API request fails
+        """
+        cloud_id = await self._get_cloud_id()
+        url = f"{ATLASSIAN_API_BASE}/ex/jira/{cloud_id}/rest/api/3/issue/{issue_key}"
+
+        response = await self._client.get(
+            url, params={"fields": "id,key,summary,status,reporter,assignee,created,updated"}
+        )
+
+        if response.status_code == HTTPStatus.NOT_FOUND:
+            raise ValueError(f"{issue_key} not found")
+
+        response.raise_for_status()
+        issue = Issue(**response.json())
+        return issue
+
+    async def __aenter__(self) -> "JiraClient":
+        """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()

datarobot_genai/drmcp/tools/clients/s3.py

@@ -0,0 +1,28 @@
+# 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.
+
+import logging
+
+from datarobot_genai.drmcp.core.credentials import get_credentials
+
+logger = logging.getLogger(__name__)
+
+
+def get_s3_bucket_info() -> dict[str, str]:
+    """Get S3 bucket configuration."""
+    credentials = get_credentials()
+    return {
+        "bucket": credentials.aws_predictions_s3_bucket,
+        "prefix": credentials.aws_predictions_s3_prefix,
+    }

datarobot_genai/drmcp/tools/confluence/__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/confluence/tools.py

@@ -0,0 +1,81 @@
+# 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.
+
+"""Confluence MCP tools for interacting with Confluence Cloud."""
+
+import logging
+from typing import Annotated
+
+from fastmcp.exceptions import ToolError
+from fastmcp.tools.tool import ToolResult
+
+from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
+from datarobot_genai.drmcp.tools.clients.atlassian import get_atlassian_access_token
+from datarobot_genai.drmcp.tools.clients.confluence import ConfluenceClient
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(tags={"confluence", "read", "get", "page"})
+async def confluence_get_page(
+    *,
+    page_id_or_title: Annotated[str, "The ID or the exact title of the Confluence page."],
+    space_key: Annotated[
+        str | None,
+        "Required if identifying the page by title. The space key (e.g., 'PROJ').",
+    ] = None,
+) -> ToolResult | ToolError:
+    """Retrieve the content of a specific Confluence page.
+
+    Use this tool to fetch Confluence pages by their numeric ID or by title.
+    Returns page content in HTML storage format.
+
+    Usage:
+        - By ID: page_id_or_title="856391684"
+        - By title: page_id_or_title="Meeting Notes", space_key="TEAM"
+
+    When using a page title, the space_key parameter is required.
+    """
+    if not page_id_or_title:
+        raise ToolError("Argument validation error: 'page_id_or_title' 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:
+            if page_id_or_title.isdigit():
+                page_response = await client.get_page_by_id(page_id_or_title)
+            else:
+                if not space_key:
+                    raise ToolError(
+                        "Argument validation error: "
+                        "'space_key' is required when identifying a page by title."
+                    )
+                page_response = await client.get_page_by_title(page_id_or_title, space_key)
+    except ValueError as e:
+        logger.error(f"Value error getting Confluence page: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error getting Confluence page: {e}")
+        raise ToolError(
+            f"An unexpected error occurred while getting Confluence page "
+            f"'{page_id_or_title}': {str(e)}"
+        )
+
+    return ToolResult(
+        content=f"Successfully retrieved page '{page_response.title}'.",
+        structured_content=page_response.as_flat_dict(),
+    )