datarobot-genai 0.2.31__py3-none-any.whl

datarobot_genai/drmcp/tools/jira/tools.py

@@ -0,0 +1,243 @@
+# 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 typing import Any
+
+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", "search", "issues"})
+async def jira_search_issues(
+    *,
+    jql_query: Annotated[
+        str, "The JQL (Jira Query Language) string used to filter and search for issues."
+    ],
+    max_results: Annotated[int, "Maximum number of issues to return. Default is 50."] = 50,
+) -> ToolResult:
+    """
+    Search for Jira issues using a powerful JQL query string.
+
+    Refer to JQL documentation for advanced query construction:
+    JQL functions: https://support.atlassian.com/jira-service-management-cloud/docs/jql-functions/
+    JQL fields: https://support.atlassian.com/jira-service-management-cloud/docs/jql-fields/
+    JQL keywords: https://support.atlassian.com/jira-service-management-cloud/docs/use-advanced-search-with-jira-query-language-jql/
+    JQL operators: https://support.atlassian.com/jira-service-management-cloud/docs/jql-operators/
+    """
+    if not jql_query:
+        raise ToolError("Argument validation error: 'jql_query' cannot be empty.")
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    async with JiraClient(access_token) as client:
+        issues = await client.search_jira_issues(jql_query=jql_query, max_results=max_results)
+
+    n = len(issues)
+    return ToolResult(
+        content=f"Successfully executed JQL query and retrieved {n} issue(s).",
+        structured_content={"data": [issue.as_flat_dict() for issue in issues], "count": n},
+    )
+
+
+@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 while 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(),
+    )
+
+
+@dr_mcp_tool(tags={"jira", "create", "add", "issue"})
+async def jira_create_issue(
+    *,
+    project_key: Annotated[str, "The key of the project where the issue should be created."],
+    summary: Annotated[str, "A brief summary or title for the new issue."],
+    issue_type: Annotated[str, "The type of issue to create (e.g., 'Task', 'Bug', 'Story')."],
+    description: Annotated[str | None, "Detailed description of the issue."] = None,
+) -> ToolResult:
+    """Create a new Jira issue with mandatory project, summary, and type information."""
+    if not all([project_key, summary, issue_type]):
+        raise ToolError(
+            "Argument validation error: project_key, summary, and issue_type are required fields."
+        )
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    async with JiraClient(access_token) as client:
+        # Maybe we should cache it somehow?
+        # It'll be probably constant through whole mcp server lifecycle...
+        issue_types = await client.get_jira_issue_types(project_key=project_key)
+
+        try:
+            issue_type_id = issue_types[issue_type]
+        except KeyError:
+            possible_issue_types = ",".join(issue_types)
+            raise ToolError(
+                f"Unexpected issue type `{issue_type}`. Possible values are {possible_issue_types}."
+            )
+
+    try:
+        async with JiraClient(access_token) as client:
+            issue_key = await client.create_jira_issue(
+                project_key=project_key,
+                summary=summary,
+                issue_type_id=issue_type_id,
+                description=description,
+            )
+    except Exception as e:
+        logger.error(f"Unexpected error while creating Jira issue: {e}")
+        raise ToolError(f"An unexpected error occurred while creating Jira issue: {str(e)}")
+
+    return ToolResult(
+        content=f"Successfully created issue '{issue_key}'.",
+        structured_content={"newIssueKey": issue_key, "projectKey": project_key},
+    )
+
+
+@dr_mcp_tool(tags={"jira", "update", "edit", "issue"})
+async def jira_update_issue(
+    *,
+    issue_key: Annotated[str, "The key (ID) of the Jira issue to retrieve, e.g., 'PROJ-123'."],
+    fields_to_update: Annotated[
+        dict[str, Any],
+        "A dictionary of field names and their new values (e.g., {'summary': 'New content'}).",
+    ],
+) -> ToolResult:
+    """
+    Modify descriptive fields or custom fields on an existing Jira issue using its key.
+    If you want to update issue status you should use `jira_transition_issue` tool instead.
+
+    Some fields needs very specific schema to allow update.
+    You should follow jira rest api guidance.
+    Good example is description field:
+        "description": {
+            "type": "text",
+            "version": 1,
+            "text": [
+                {
+                    "type": "paragraph",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": "[HERE YOU PUT REAL DESCRIPTION]"
+                        }
+                    ]
+                }
+            ]
+        }
+    """
+    if not issue_key:
+        raise ToolError("Argument validation error: 'issue_key' cannot be empty.")
+    if not fields_to_update or not isinstance(fields_to_update, dict):
+        raise ToolError(
+            "Argument validation error: 'fields_to_update' must be a non-empty dictionary."
+        )
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with JiraClient(access_token) as client:
+            updated_fields = await client.update_jira_issue(
+                issue_key=issue_key, fields=fields_to_update
+            )
+    except Exception as e:
+        logger.error(f"Unexpected error while updating Jira issue: {e}")
+        raise ToolError(f"An unexpected error occurred while updating Jira issue: {str(e)}")
+
+    updated_fields_str = ",".join(updated_fields)
+    return ToolResult(
+        content=f"Successfully updated issue '{issue_key}'. Fields modified: {updated_fields_str}.",
+        structured_content={"updatedIssueKey": issue_key, "fields": updated_fields},
+    )
+
+
+@dr_mcp_tool(tags={"jira", "update", "transition", "issue"})
+async def jira_transition_issue(
+    *,
+    issue_key: Annotated[str, "The key (ID) of the Jira issue to transition, e.g. 'PROJ-123'."],
+    transition_name: Annotated[
+        str, "The exact name of the target status/transition (e.g., 'In Progress')."
+    ],
+) -> ToolResult:
+    """
+    Move a Jira issue through its defined workflow to a new status.
+    This leverages Jira's workflow engine directly.
+    """
+    if not all([issue_key, transition_name]):
+        raise ToolError("Argument validation error: issue_key and transition name/ID are required.")
+
+    access_token = await get_atlassian_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    async with JiraClient(access_token) as client:
+        available_transitions = await client.get_available_jira_transitions(issue_key=issue_key)
+
+        try:
+            transition_id = available_transitions[transition_name]
+        except KeyError:
+            available_transitions_str = ",".join(available_transitions)
+            raise ToolError(
+                f"Unexpected transition name `{transition_name}`. "
+                f"Possible values are {available_transitions_str}."
+            )
+
+    try:
+        async with JiraClient(access_token) as client:
+            await client.transition_jira_issue(issue_key=issue_key, transition_id=transition_id)
+    except Exception as e:
+        logger.error(f"Unexpected error while transitioning Jira issue: {e}")
+        raise ToolError(f"An unexpected error occurred while transitioning Jira issue: {str(e)}")
+
+    return ToolResult(
+        content=f"Successfully transitioned issue '{issue_key}' to status '{transition_name}'.",
+        structured_content={
+            "transitionedIssueKey": issue_key,
+            "newStatusName": transition_name,
+            "newStatusId": transition_id,
+        },
+    )

datarobot_genai/drmcp/tools/microsoft_graph/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2026 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/microsoft_graph/tools.py

@@ -0,0 +1,198 @@
+# Copyright 2026 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.
+
+"""Microsoft Graph MCP tools for searching SharePoint and OneDrive content."""
+
+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.microsoft_graph import MicrosoftGraphClient
+from datarobot_genai.drmcp.tools.clients.microsoft_graph import MicrosoftGraphError
+from datarobot_genai.drmcp.tools.clients.microsoft_graph import get_microsoft_graph_access_token
+from datarobot_genai.drmcp.tools.clients.microsoft_graph import validate_site_url
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(
+    tags={
+        "microsoft",
+        "graph api",
+        "sharepoint",
+        "drive",
+        "list",
+        "search",
+        "files",
+        "find",
+        "contents",
+    }
+)
+async def microsoft_graph_search_content(
+    *,
+    search_query: Annotated[str, "The search string to find files, folders, or list items."],
+    site_url: Annotated[
+        str | None,
+        "Optional SharePoint site URL to scope the search "
+        "(e.g., https://tenant.sharepoint.com/sites/sitename). "
+        "If not provided, searches across all accessible sites.",
+    ] = None,
+    site_id: Annotated[
+        str | None,
+        "Optional ID of the site to scope the search. If provided, takes precedence over site_url.",
+    ] = None,
+    from_offset: Annotated[
+        int,
+        "The zero-based index of the first result to return. Use this for pagination. "
+        "Default: 0 (start from the beginning). To get the next page, increment by the size "
+        "value (e.g., first page: from=0 size=250, second page: from=250 size=250, "
+        "third page: from=500 size=250).",
+    ] = 0,
+    size: Annotated[
+        int,
+        "Maximum number of results to return in this request. Default is 250, max is 250. "
+        "The LLM should control pagination by making multiple calls with different 'from' values.",
+    ] = 250,
+    entity_types: Annotated[
+        list[str] | None,
+        "Optional list of entity types to search. Valid values: 'driveItem', 'listItem', "
+        "'site', 'list', 'drive'. Default: ['driveItem', 'listItem']. "
+        "Multiple types can be specified.",
+    ] = None,
+    filters: Annotated[
+        list[str] | None,
+        "Optional list of KQL filter expressions to refine search results "
+        "(e.g., ['fileType:docx', 'size>1000']).",
+    ] = None,
+    include_hidden_content: Annotated[
+        bool,
+        "Whether to include hidden content in search results. Only works with delegated "
+        "permissions, not application permissions. Default: False.",
+    ] = False,
+    region: Annotated[
+        str | None,
+        "Optional region code for application permissions (e.g., 'NAM', 'EUR', 'APC'). "
+        "Required when using application permissions to search SharePoint content in "
+        "specific regions.",
+    ] = None,
+) -> ToolResult | ToolError:
+    """
+    Search for SharePoint and OneDrive content using Microsoft Graph Search API.
+
+    Search Scope:
+    - When site_url or site_id is provided: searches within the specified SharePoint site
+    - When neither is provided: searches across all accessible SharePoint sites and OneDrive
+
+    Supported Entity Types:
+    - driveItem: Files and folders in document libraries and OneDrive
+    - listItem: Items in SharePoint lists
+    - site: SharePoint sites
+    - list: SharePoint lists
+    - drive: Document libraries/drives
+
+    Filtering:
+    - Filters use KQL (Keyword Query Language) syntax
+    - Multiple filters are combined with AND operators
+    - Examples: ['fileType:docx', 'size>1000', 'lastModifiedTime>2024-01-01']
+    - Filters are applied in addition to the search query
+
+    Pagination:
+    - Controlled via from_offset (zero-based index) and size parameters
+    - Maximum size per request: 250 results
+    - To paginate: increment from_offset by size value for each subsequent page
+    - Example pagination sequence:
+      * Page 1: from_offset=0, size=250 (returns results 0-249)
+      * Page 2: from_offset=250, size=250 (returns results 250-499)
+      * Page 3: from_offset=500, size=250 (returns results 500-749)
+
+    API Reference:
+    - Endpoint: POST /search/query
+    - Documentation: https://learn.microsoft.com/en-us/graph/api/search-query
+    - Search concepts: https://learn.microsoft.com/en-us/graph/search-concept-files
+
+    Permissions:
+    - Requires Sites.Read.All or Sites.Search.All permission
+    - include_hidden_content only works with delegated permissions
+    - region parameter is required for application permissions in multi-region environments
+    """
+    if not search_query:
+        raise ToolError("Argument validation error: 'search_query' cannot be empty.")
+
+    # Validate site_url if provided
+    if site_url:
+        validation_error = validate_site_url(site_url)
+        if validation_error:
+            raise ToolError(validation_error)
+
+    access_token = await get_microsoft_graph_access_token()
+    if isinstance(access_token, ToolError):
+        raise access_token
+
+    try:
+        async with MicrosoftGraphClient(access_token=access_token, site_url=site_url) as client:
+            items = await client.search_content(
+                search_query=search_query,
+                site_id=site_id,
+                from_offset=from_offset,
+                size=size,
+                entity_types=entity_types,
+                filters=filters,
+                include_hidden_content=include_hidden_content,
+                region=region,
+            )
+    except MicrosoftGraphError as e:
+        logger.error(f"Microsoft Graph error searching content: {e}")
+        raise ToolError(str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error searching Microsoft Graph content: {e}", exc_info=True)
+        raise ToolError(
+            f"An unexpected error occurred while searching Microsoft Graph content: {str(e)}"
+        )
+
+    results = []
+    for item in items:
+        result_dict = {
+            "id": item.id,  # Unique ID of the file, folder, or list item
+            "name": item.name,
+            "webUrl": item.web_url,
+            "size": item.size,
+            "createdDateTime": item.created_datetime,
+            "lastModifiedDateTime": item.last_modified_datetime,
+            "isFolder": item.is_folder,
+            "mimeType": item.mime_type,
+            # Document library/drive ID (driveId in Microsoft Graph API)
+            "documentLibraryId": item.drive_id,
+            "parentFolderId": item.parent_folder_id,  # Parent folder ID
+        }
+        results.append(result_dict)
+
+    n = len(results)
+    return ToolResult(
+        content=(
+            f"Successfully searched Microsoft Graph and retrieved {n} result(s) for "
+            f"'{search_query}' (from={from_offset}, size={size})."
+        ),
+        structured_content={
+            "query": search_query,
+            "siteUrl": site_url,
+            "siteId": site_id,
+            "from": from_offset,
+            "size": size,
+            "results": results,
+            "count": n,
+        },
+    )

datarobot_genai/drmcp/tools/predictive/__init__.py

@@ -0,0 +1,27 @@
+# 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 MCP tools for interacting with DataRobot platform.
+
+This package contains pre-built tools for DataRobot operations:
+- data: Upload and manage datasets
+- project: Create and manage projects
+- model: List and interact with models
+- training: Train models with autopilot
+- deployment: Deploy and manage models
+- predict: Make batch predictions
+- predict_realtime: Make real-time predictions
+- deployment_info: Get deployment information
+"""

datarobot_genai/drmcp/tools/predictive/data.py

@@ -0,0 +1,133 @@
+# 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 json
+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
+from datarobot_genai.drmcp.core.utils import is_valid_url
+
+logger = logging.getLogger(__name__)
+
+
+@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."] | None = None,
+    file_url: Annotated[str, "The URL to the dataset file to upload."] | None = None,
+) -> ToolError | ToolResult:
+    """Upload a dataset to the DataRobot AI Catalog / Data Registry."""
+    if not file_path and not file_url:
+        return ToolError("Either file_path or file_url must be provided.")
+    if file_path and file_url:
+        return ToolError("Please provide either file_path or file_url, not both.")
+
+    # Get client
+    client = get_sdk_client()
+    catalog_item = None
+    # If file path is provided, create dataset from file.
+    if file_path:
+        # Does file exist?
+        if not os.path.exists(file_path):
+            logger.error("File not found: %s", file_path)
+            return ToolError(f"File not found: {file_path}")
+        catalog_item = client.Dataset.create_from_file(file_path)
+    else:
+        # Does URL exist?
+        if file_url is None or not is_valid_url(file_url):
+            logger.error("Invalid file URL: %s", file_url)
+            return ToolError(f"Invalid file URL: {file_url}")
+        catalog_item = client.Dataset.create_from_url(file_url)
+
+    if not catalog_item:
+        return ToolError("Failed to upload dataset.")
+
+    return ToolResult(
+        content=f"Successfully uploaded dataset: {catalog_item.id}",
+        structured_content={
+            "dataset_id": catalog_item.id,
+            "dataset_version_id": catalog_item.version_id,
+            "dataset_name": catalog_item.name,
+        },
+    )
+
+
+@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 ToolResult(
+            content="No AI Catalog items found.",
+            structured_content={"datasets": []},
+        )
+
+    datasets_dict = {ds.id: ds.name for ds in datasets}
+    datasets_count = len(datasets)
+
+    return ToolResult(
+        content=(
+            f"Found {datasets_count} AI Catalog items, here are the details:\n"
+            f"{json.dumps(datasets_dict, indent=2)}"
+        ),
+        structured_content={
+            "datasets": datasets_dict,
+            "count": datasets_count,
+        },
+    )
+
+
+# from fastmcp import Context
+
+# from datarobot_genai.drmcp.core.memory_management import MemoryManager, get_memory_manager
+
+
+# @dr_mcp_tool()
+# async def list_ai_catalog_items(
+#     ctx: Context, agent_id: str = None, storage_id: str = None
+# ) -> str:
+#     """
+#     List all AI Catalog items (datasets) for the authenticated user.
+
+#     Returns:
+#         a resource id that can be used to retrieve the list of AI Catalog items using the
+#         get_resource tool
+#     """
+#     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)
+
+#     if MemoryManager.is_initialized():
+#         resource_id = await get_memory_manager().store_resource(
+#             data=result,
+#             memory_storage_id=storage_id,
+#             agent_identifier=agent_id,
+#         )
+#     else:
+#         raise ValueError("MemoryManager is not initialized")
+
+#     logger.info(f"Found {len(datasets)} AI Catalog items")
+#     return resource_id