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

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

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

@@ -0,0 +1,52 @@
+# 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 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.jira import JiraClient
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(tags={"jira", "read", "get", "issue"})
+async def jira_get_issue(
+    *, issue_key: Annotated[str, "The key (ID) of the Jira issue to retrieve, e.g., 'PROJ-123'."]
+) -> ToolResult:
+    """Retrieve all fields and details for a single Jira issue by its key."""
+    if not issue_key:
+        raise ToolError("Argument validation error: 'issue_key' cannot be empty.")
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with JiraClient(access_token) as client:
+            issue = await client.get_jira_issue(issue_key)
+    except Exception as e:
+        logger.error(f"Unexpected error getting Jira issue: {e}")
+        raise ToolError(
+            f"An unexpected error occurred while getting Jira issue '{issue_key}': {str(e)}"
+        )
+
+    return ToolResult(
+        content=f"Successfully retrieved details for issue '{issue_key}'.",
+        structured_content=issue.as_flat_dict(),
+    )

datarobot_genai/drmcp/tools/predictive/predict.py

@@ -22,10 +22,10 @@ from fastmcp.resources import HttpResource
 from fastmcp.resources import ResourceManager
 
 from datarobot_genai.drmcp.core.clients import get_credentials
-from datarobot_genai.drmcp.core.clients import get_s3_bucket_info
 from datarobot_genai.drmcp.core.clients import get_sdk_client
 from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
 from datarobot_genai.drmcp.core.utils import generate_presigned_url
+from datarobot_genai.drmcp.tools.clients.s3 import get_s3_bucket_info
 
 logger = logging.getLogger(__name__)
 

datarobot_genai/drmcp/tools/predictive/predict_realtime.py

@@ -23,11 +23,11 @@ from datarobot_predict import TimeSeriesType
 from datarobot_predict.deployment import predict as dr_predict
 from pydantic import BaseModel
 
-from datarobot_genai.drmcp.core.clients import get_s3_bucket_info
 from datarobot_genai.drmcp.core.clients import get_sdk_client
 from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
 from datarobot_genai.drmcp.core.utils import PredictionResponse
 from datarobot_genai.drmcp.core.utils import predictions_result_response
+from datarobot_genai.drmcp.tools.clients.s3 import get_s3_bucket_info
 
 logger = logging.getLogger(__name__)
 

datarobot_genai/nat/agent.py

@@ -123,6 +123,8 @@ class NatAgent(BaseAgent[None]):
         model: str | None = None,
         verbose: bool | str | None = True,
         timeout: int | None = 90,
+        authorization_context: dict[str, Any] | None = None,
+        forwarded_headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -131,6 +133,8 @@ class NatAgent(BaseAgent[None]):
             model=model,
             verbose=verbose,
             timeout=timeout,
+            authorization_context=authorization_context,
+            forwarded_headers=forwarded_headers,
             **kwargs,
         )
         self.workflow_path = workflow_path

datarobot_genai/nat/datarobot_auth_provider.py

@@ -0,0 +1,110 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# 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.
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from nat.authentication.api_key.api_key_auth_provider import APIKeyAuthProvider
+from nat.authentication.api_key.api_key_auth_provider_config import APIKeyAuthProviderConfig
+from nat.authentication.interfaces import AuthProviderBase
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_auth_provider
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.data_models.authentication import AuthResult
+from nat.data_models.authentication import HeaderCred
+from pydantic import Field
+
+from datarobot_genai.core.mcp.common import MCPConfig
+
+
+class Config(DataRobotAppFrameworkBaseSettings):
+    """
+    Finds variables in the priority order of: env
+    variables (including Runtime Parameters), .env, file_secrets, then
+    Pulumi output variables.
+    """
+
+    datarobot_api_token: str | None = None
+
+
+config = Config()
+
+
+class DataRobotAPIKeyAuthProviderConfig(APIKeyAuthProviderConfig, name="datarobot_api_key"):  # type: ignore[call-arg]
+    raw_key: str = Field(
+        description=(
+            "Raw API token or credential to be injected into the request parameter. "
+            "Used for 'bearer','x-api-key','custom', and other schemes. "
+        ),
+        default=config.datarobot_api_token,
+    )
+    default_user_id: str | None = Field(default="default-user", description="Default user ID")
+    allow_default_user_id_for_tool_calls: bool = Field(
+        default=True, description="Allow default user ID for tool calls"
+    )
+
+
+@register_auth_provider(config_type=DataRobotAPIKeyAuthProviderConfig)
+async def datarobot_api_key_client(
+    config: DataRobotAPIKeyAuthProviderConfig, builder: Builder
+) -> AsyncGenerator[APIKeyAuthProvider]:
+    yield APIKeyAuthProvider(config=config)
+
+
+mcp_config = MCPConfig().server_config
+
+
+class DataRobotMCPAuthProviderConfig(AuthProviderBaseConfig, name="datarobot_mcp_auth"):  # type: ignore[call-arg]
+    headers: dict[str, str] | None = Field(
+        description=("Headers to be used for authentication. "),
+        default=mcp_config["headers"] if mcp_config else None,
+    )
+    default_user_id: str | None = Field(default="default-user", description="Default user ID")
+    allow_default_user_id_for_tool_calls: bool = Field(
+        default=True, description="Allow default user ID for tool calls"
+    )
+
+
+class DataRobotMCPAuthProvider(AuthProviderBase[DataRobotMCPAuthProviderConfig]):
+    def __init__(
+        self, config: DataRobotMCPAuthProviderConfig, config_name: str | None = None
+    ) -> None:
+        assert isinstance(config, DataRobotMCPAuthProviderConfig), (
+            "Config is not DataRobotMCPAuthProviderConfig"
+        )
+        super().__init__(config)
+
+    async def authenticate(self, user_id: str | None = None, **kwargs: Any) -> AuthResult | None:
+        """
+        Authenticate the user using the API key credentials.
+
+        Args:
+            user_id (str): The user ID to authenticate.
+
+        Returns
+        -------
+            AuthenticatedContext: The authenticated context containing headers
+        """
+        return AuthResult(
+            credentials=[
+                HeaderCred(name=name, value=value) for name, value in self.config.headers.items()
+            ]
+        )
+
+
+@register_auth_provider(config_type=DataRobotMCPAuthProviderConfig)
+async def datarobot_mcp_auth_provider(
+    config: DataRobotMCPAuthProviderConfig, builder: Builder
+) -> AsyncGenerator[DataRobotMCPAuthProvider]:
+    yield DataRobotMCPAuthProvider(config=config)