datarobot-genai 0.2.22__py3-none-any.whl → 0.2.26__py3-none-any.whl

datarobot_genai/drmcp/core/dr_mcp_server.py

@@ -31,9 +31,6 @@ from .dynamic_prompts.register import register_prompts_from_datarobot_prompt_man
 from .dynamic_tools.deployment.register import register_tools_of_datarobot_deployments
 from .logging import MCPLogging
 from .mcp_instance import mcp
-from .mcp_server_tools import get_all_available_tags  # noqa # pylint: disable=unused-import
-from .mcp_server_tools import get_tool_info_by_name  # noqa # pylint: disable=unused-import
-from .mcp_server_tools import list_tools_by_tags  # noqa # pylint: disable=unused-import
 from .memory_management.manager import MemoryManager
 from .routes import register_routes
 from .routes_utils import prefix_mount_path

datarobot_genai/drmcp/core/mcp_instance.py

@@ -16,20 +16,18 @@ import logging
 from collections.abc import Callable
 from functools import wraps
 from typing import Any
-from typing import overload
+from typing import TypedDict
 
 from fastmcp import Context
 from fastmcp import FastMCP
 from fastmcp.exceptions import NotFoundError
 from fastmcp.prompts.prompt import Prompt
 from fastmcp.server.dependencies import get_context
-from fastmcp.tools import FunctionTool
 from fastmcp.tools import Tool
-from fastmcp.utilities.types import NotSet
-from fastmcp.utilities.types import NotSetT
 from mcp.types import AnyFunction
 from mcp.types import Tool as MCPTool
 from mcp.types import ToolAnnotations
+from typing_extensions import Unpack
 
 from .config import MCPServerConfig
 from .config import get_config
@@ -120,86 +118,6 @@ class TaggedFastMCP(FastMCP):
                 "In stateless mode, clients will see changes on next request."
             )
 
-    @overload
-    def tool(
-        self,
-        name_or_fn: AnyFunction,
-        *,
-        name: str | None = None,
-        title: str | None = None,
-        description: str | None = None,
-        tags: set[str] | None = None,
-        output_schema: dict[str, Any] | None | NotSetT = NotSet,
-        annotations: ToolAnnotations | dict[str, Any] | None = None,
-        exclude_args: list[str] | None = None,
-        meta: dict[str, Any] | None = None,
-        enabled: bool | None = None,
-    ) -> FunctionTool: ...
-
-    @overload
-    def tool(
-        self,
-        name_or_fn: str | None = None,
-        *,
-        name: str | None = None,
-        title: str | None = None,
-        description: str | None = None,
-        tags: set[str] | None = None,
-        output_schema: dict[str, Any] | None | NotSetT = NotSet,
-        annotations: ToolAnnotations | dict[str, Any] | None = None,
-        exclude_args: list[str] | None = None,
-        meta: dict[str, Any] | None = None,
-        enabled: bool | None = None,
-    ) -> Callable[[AnyFunction], FunctionTool]: ...
-
-    def tool(
-        self,
-        name_or_fn: str | Callable[..., Any] | None = None,
-        *,
-        name: str | None = None,
-        title: str | None = None,
-        description: str | None = None,
-        tags: set[str] | None = None,
-        output_schema: dict[str, Any] | None | NotSetT = NotSet,
-        annotations: ToolAnnotations | dict[str, Any] | None = None,
-        exclude_args: list[str] | None = None,
-        meta: dict[str, Any] | None = None,
-        enabled: bool | None = None,
-        **kwargs: Any,
-    ) -> Callable[[AnyFunction], FunctionTool] | FunctionTool:
-        """
-        Extend tool decorator that supports tags and other annotations, while remaining
-        signature-compatible with FastMCP.tool to avoid recursion issues with partials.
-        """
-        if isinstance(annotations, dict):
-            annotations = ToolAnnotations(**annotations)
-
-        # Ensure tags are available both via native fastmcp `tags` and inside annotations
-        if tags is not None:
-            tags_ = sorted(tags)
-            if annotations is None:
-                annotations = ToolAnnotations()  # type: ignore[call-arg]
-                annotations.tags = tags_  # type: ignore[attr-defined, union-attr]
-            else:
-                # At this point, annotations is ToolAnnotations (not dict)
-                assert isinstance(annotations, ToolAnnotations)
-                annotations.tags = tags_  # type: ignore[attr-defined]
-
-        return super().tool(
-            name_or_fn,
-            name=name,
-            title=title,
-            description=description,
-            tags=tags,
-            output_schema=output_schema
-            if output_schema is not None
-            else kwargs.get("output_schema"),
-            annotations=annotations,
-            exclude_args=exclude_args,
-            meta=meta,
-            enabled=enabled,
-        )
-
     async def list_tools(
         self, tags: list[str] | None = None, match_all: bool = False
     ) -> list[MCPTool]:
@@ -371,16 +289,37 @@ mcp = TaggedFastMCP(
 )
 
 
+class ToolKwargs(TypedDict, total=False):
+    """Keyword arguments passed through to FastMCP's mcp.tool() decorator.
+
+    All parameters are optional and forwarded directly to FastMCP tool registration.
+    See FastMCP documentation for full details on each parameter.
+    """
+
+    name: str | None
+    title: str | None
+    description: str | None
+    icons: list[Any] | None
+    tags: set[str] | None
+    output_schema: dict[str, Any] | None
+    annotations: Any | None
+    exclude_args: list[str] | None
+    meta: dict[str, Any] | None
+    enabled: bool | None
+
+
 def dr_core_mcp_tool(
-    name: str | None = None,
-    description: str | None = None,
-    tags: set[str] | None = None,
+    **kwargs: Unpack[ToolKwargs],
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
-    """Combine decorator that includes mcp.tool() and dr_mcp_extras()."""
+    """Combine decorator that includes mcp.tool() and dr_mcp_extras().
+
+    All keyword arguments are passed through to FastMCP's mcp.tool() decorator.
+    See ToolKwargs for available parameters.
+    """
 
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         instrumented = dr_mcp_extras()(func)
-        mcp.tool(name=name, description=description, tags=tags)(instrumented)
+        mcp.tool(**kwargs)(instrumented)
         return instrumented
 
     return decorator
@@ -413,27 +352,23 @@ async def memory_aware_wrapper(func: Callable[..., Any], *args: Any, **kwargs: A
 
 
 def dr_mcp_tool(
-    name: str | None = None,
-    description: str | None = None,
-    tags: set[str] | None = None,
+    **kwargs: Unpack[ToolKwargs],
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
     """Combine decorator that includes mcp.tool(), dr_mcp_extras(), and capture memory ids from
     the request headers if they exist.
 
-    Args:
-        name: Tool name
-        description: Tool description
-        tags: Optional set of tags to apply to the tool
+    All keyword arguments are passed through to FastMCP's mcp.tool() decorator.
+    See ToolKwargs for available parameters.
     """
 
     def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
         @wraps(func)
-        async def wrapper(*args: Any, **kwargs: Any) -> Any:
-            return await memory_aware_wrapper(func, *args, **kwargs)
+        async def wrapper(*args: Any, **inner_kwargs: Any) -> Any:
+            return await memory_aware_wrapper(func, *args, **inner_kwargs)
 
         # Apply the MCP decorators
         instrumented = dr_mcp_extras()(wrapper)
-        mcp.tool(name=name, description=description, tags=tags)(instrumented)
+        mcp.tool(**kwargs)(instrumented)
         return instrumented
 
     return decorator
@@ -488,11 +423,10 @@ async def register_tools(
     # Apply dr_mcp_extras to the memory-aware function
     wrapped_fn = dr_mcp_extras()(memory_aware_fn)
 
-    # Create annotations with tags, deployment_id if provided
-    annotations = ToolAnnotations()  # type: ignore[call-arg]
-    if tags is not None:
-        annotations.tags = tags  # type: ignore[attr-defined]
+    # Create annotations only when additional metadata is required
+    annotations: ToolAnnotations | None = None  # type: ignore[assignment]
     if deployment_id is not None:
+        annotations = ToolAnnotations()  # type: ignore[call-arg]
         annotations.deployment_id = deployment_id  # type: ignore[attr-defined]
 
     tool = Tool.from_function(

datarobot_genai/drmcp/core/tool_filter.py

@@ -41,7 +41,7 @@ def filter_tools_by_tags(
     filtered_tools = []
 
     for tool in tools:
-        tool_tags = getattr(tool.annotations, "tags", []) if tool.annotations else []
+        tool_tags = get_tool_tags(tool)
 
         if not tool_tags:
             continue
@@ -68,9 +68,18 @@ def get_tool_tags(tool: Tool | MCPTool) -> list[str]:
     -------
         List of tags for the tool
     """
+    # Primary: native FastMCP meta location
+    if hasattr(tool, "meta") and getattr(tool, "meta"):
+        fastmcp_meta = tool.meta.get("_fastmcp", {})
+        meta_tags = fastmcp_meta.get("tags", [])
+        if isinstance(meta_tags, list):
+            return meta_tags
+
+    # Fallback: annotations.tags (for compatibility during transition)
     if tool.annotations and hasattr(tool.annotations, "tags"):
         tags = getattr(tool.annotations, "tags", [])
         return tags if isinstance(tags, list) else []
+
     return []
 
 

datarobot_genai/drmcp/tools/clients/confluence.py

@@ -50,6 +50,7 @@ class ConfluencePage(BaseModel):
     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)")
+    version: int = Field(..., description="Current version number of the page")
 
     def as_flat_dict(self) -> dict[str, Any]:
         """Return a flat dictionary representation of the page."""
@@ -59,6 +60,7 @@ class ConfluencePage(BaseModel):
             "space_id": self.space_id,
             "space_key": self.space_key,
             "body": self.body,
+            "version": self.version,
         }
 
 
@@ -111,7 +113,7 @@ class ConfluenceClient:
     At the moment of creating this client, official Confluence SDK is not supporting async.
     """
 
-    EXPAND_FIELDS = "body.storage,space"
+    EXPAND_FIELDS = "body.storage,space,version"
 
     def __init__(self, access_token: str) -> None:
         """
@@ -164,6 +166,8 @@ class ConfluenceClient:
         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", "")
+        version_data = data.get("version", {})
+        version_number = version_data.get("number", 1) if isinstance(version_data, dict) else 1
 
         return ConfluencePage(
             page_id=str(data.get("id", "")),
@@ -171,6 +175,7 @@ class ConfluenceClient:
             space_id=str(space_id),
             space_key=space_key,
             body=body_content,
+            version=version_number,
         )
 
     async def get_page_by_id(self, page_id: str) -> ConfluencePage:
@@ -339,6 +344,93 @@ class ConfluenceClient:
 
         return self._parse_response(response.json())
 
+    async def update_page(
+        self,
+        page_id: str,
+        new_body_content: str,
+        version_number: int,
+    ) -> ConfluencePage:
+        """
+        Update the content of an existing Confluence page.
+
+        Args:
+            page_id: The ID of the page to update
+            new_body_content: The new content in Confluence Storage Format (XML) or raw text
+            version_number: The current version number of the page (for optimistic locking).
+                           The update will increment this by 1.
+
+        Returns
+        -------
+            ConfluencePage with the updated page data including the new version number
+
+        Raises
+        ------
+            ConfluenceError: If page not found (404), version conflict (409),
+                            permission denied (403), invalid content (400),
+                            or rate limited (429)
+            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/{page_id}"
+
+        try:
+            current_page = await self.get_page_by_id(page_id)
+            title_to_use = current_page.title
+        except ConfluenceError as e:
+            if e.status_code == 404:
+                raise ConfluenceError(
+                    f"Page with ID '{page_id}' not found: cannot fetch existing title",
+                    status_code=404,
+                )
+            raise
+
+        payload: dict[str, Any] = {
+            "type": "page",
+            "title": title_to_use,
+            "body": {
+                "storage": {
+                    "value": new_body_content,
+                    "representation": "storage",
+                }
+            },
+            "version": {
+                "number": version_number + 1,
+            },
+        }
+
+        response = await self._client.put(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.CONFLICT:
+            error_msg = self._extract_error_message(response)
+            raise ConfluenceError(
+                f"Version conflict: the page has been modified since version {version_number}. "
+                f"Please fetch the latest version and retry. Details: {error_msg}",
+                status_code=409,
+            )
+
+        if response.status_code == HTTPStatus.FORBIDDEN:
+            raise ConfluenceError(
+                f"Permission denied: you don't have access to update 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_response(response.json())
+
     def _parse_comment_response(self, data: dict, page_id: str) -> ConfluenceComment:
         """Parse API response into ConfluenceComment."""
         body_content = ""

datarobot_genai/drmcp/tools/clients/gdrive.py

@@ -14,6 +14,7 @@
 
 """Google Drive API Client and utilities for OAuth."""
 
+import io
 import logging
 from typing import Annotated
 from typing import Any
@@ -24,6 +25,7 @@ from fastmcp.exceptions import ToolError
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Field
+from pypdf import PdfReader
 
 from datarobot_genai.drmcp.core.auth import get_access_token
 
@@ -37,6 +39,23 @@ DEFAULT_ORDER = "modifiedTime desc"
 MAX_PAGE_SIZE = 100
 LIMIT = 500
 
+GOOGLE_WORKSPACE_EXPORT_MIMES: dict[str, str] = {
+    "application/vnd.google-apps.document": "text/markdown",
+    "application/vnd.google-apps.spreadsheet": "text/csv",
+    "application/vnd.google-apps.presentation": "text/plain",
+}
+
+BINARY_MIME_PREFIXES = (
+    "image/",
+    "audio/",
+    "video/",
+    "application/zip",
+    "application/octet-stream",
+    "application/vnd.google-apps.drawing",
+)
+
+PDF_MIME_TYPE = "application/pdf"
+
 
 async def get_gdrive_access_token() -> str | ToolError:
     """
@@ -116,6 +135,35 @@ class PaginatedResult(BaseModel):
     next_page_token: str | None = None
 
 
+class GoogleDriveFileContent(BaseModel):
+    """Content retrieved from a Google Drive file."""
+
+    id: str
+    name: str
+    mime_type: str
+    content: str
+    original_mime_type: str
+    was_exported: bool = False
+    size: int | None = None
+    web_view_link: str | None = None
+
+    def as_flat_dict(self) -> dict[str, Any]:
+        """Return a flat dictionary representation of the file content."""
+        result: dict[str, Any] = {
+            "id": self.id,
+            "name": self.name,
+            "mimeType": self.mime_type,
+            "content": self.content,
+            "originalMimeType": self.original_mime_type,
+            "wasExported": self.was_exported,
+        }
+        if self.size is not None:
+            result["size"] = self.size
+        if self.web_view_link is not None:
+            result["webViewLink"] = self.web_view_link
+        return result
+
+
 class GoogleDriveClient:
     """Client for interacting with Google Drive API."""
 
@@ -344,6 +392,213 @@ class GoogleDriveClient:
             logger.debug(f"Auto-formatted query '{query}' to '{formatted_query}'")
         return formatted_query
 
+    @staticmethod
+    def _is_binary_mime_type(mime_type: str) -> bool:
+        """Check if MIME type indicates binary content that's not useful for LLM consumption.
+
+        Args:
+            mime_type: The MIME type to check.
+
+        Returns
+        -------
+            True if the MIME type is considered binary, False otherwise.
+        """
+        return any(mime_type.startswith(prefix) for prefix in BINARY_MIME_PREFIXES)
+
+    async def get_file_metadata(self, file_id: str) -> GoogleDriveFile:
+        """Get file metadata from Google Drive.
+
+        Args:
+            file_id: The ID of the file to get metadata for.
+
+        Returns
+        -------
+            GoogleDriveFile with file metadata.
+
+        Raises
+        ------
+            GoogleDriveError: If the file is not found or access is denied.
+        """
+        params = {"fields": SUPPORTED_FIELDS_STR}
+        response = await self._client.get(f"/{file_id}", params=params)
+
+        if response.status_code == 404:
+            raise GoogleDriveError(f"File with ID '{file_id}' not found.")
+        if response.status_code == 403:
+            raise GoogleDriveError(f"Permission denied: you don't have access to file '{file_id}'.")
+        if response.status_code == 429:
+            raise GoogleDriveError("Rate limit exceeded. Please try again later.")
+
+        response.raise_for_status()
+        return GoogleDriveFile.from_api_response(response.json())
+
+    async def _export_workspace_file(self, file_id: str, export_mime_type: str) -> str:
+        """Export a Google Workspace file to the specified format.
+
+        Args:
+            file_id: The ID of the Google Workspace file.
+            export_mime_type: The MIME type to export to (e.g., 'text/markdown').
+
+        Returns
+        -------
+            The exported content as a string.
+
+        Raises
+        ------
+            GoogleDriveError: If export fails.
+        """
+        response = await self._client.get(
+            f"/{file_id}/export",
+            params={"mimeType": export_mime_type},
+        )
+
+        if response.status_code == 404:
+            raise GoogleDriveError(f"File with ID '{file_id}' not found.")
+        if response.status_code == 403:
+            raise GoogleDriveError(
+                f"Permission denied: you don't have access to export file '{file_id}'."
+            )
+        if response.status_code == 400:
+            raise GoogleDriveError(
+                f"Cannot export file '{file_id}' to format '{export_mime_type}'. "
+                "The file may not support this export format."
+            )
+        if response.status_code == 429:
+            raise GoogleDriveError("Rate limit exceeded. Please try again later.")
+
+        response.raise_for_status()
+        return response.text
+
+    async def _download_file(self, file_id: str) -> str:
+        """Download a regular file's content from Google Drive as text."""
+        content = await self._download_file_bytes(file_id)
+        return content.decode("utf-8")
+
+    async def _download_file_bytes(self, file_id: str) -> bytes:
+        """Download a file's content as bytes from Google Drive.
+
+        Args:
+            file_id: The ID of the file to download.
+
+        Returns
+        -------
+            The file content as bytes.
+
+        Raises
+        ------
+            GoogleDriveError: If download fails.
+        """
+        response = await self._client.get(
+            f"/{file_id}",
+            params={"alt": "media"},
+        )
+
+        if response.status_code == 404:
+            raise GoogleDriveError(f"File with ID '{file_id}' not found.")
+        if response.status_code == 403:
+            raise GoogleDriveError(
+                f"Permission denied: you don't have access to download file '{file_id}'."
+            )
+        if response.status_code == 429:
+            raise GoogleDriveError("Rate limit exceeded. Please try again later.")
+
+        response.raise_for_status()
+        return response.content
+
+    def _extract_text_from_pdf(self, pdf_bytes: bytes) -> str:
+        """Extract text from PDF bytes using pypdf.
+
+        Args:
+            pdf_bytes: The PDF file content as bytes.
+
+        Returns
+        -------
+            Extracted text from the PDF.
+
+        Raises
+        ------
+            GoogleDriveError: If PDF text extraction fails.
+        """
+        try:
+            reader = PdfReader(io.BytesIO(pdf_bytes))
+            text_parts = []
+            for page in reader.pages:
+                page_text = page.extract_text()
+                if page_text:
+                    text_parts.append(page_text)
+            return "\n\n".join(text_parts)
+        except Exception as e:
+            raise GoogleDriveError(f"Failed to extract text from PDF: {e}")
+
+    async def read_file_content(
+        self, file_id: str, target_format: str | None = None
+    ) -> GoogleDriveFileContent:
+        """Read the content of a file from Google Drive.
+
+        Google Workspace files (Docs, Sheets, Slides) are automatically exported to
+        LLM-readable formats:
+        - Google Docs -> Markdown (text/markdown)
+        - Google Sheets -> CSV (text/csv)
+        - Google Slides -> Plain text (text/plain)
+        - PDF files -> Extracted text (text/plain)
+
+        Regular text files are downloaded directly.
+        Binary files (images, videos, etc.) will raise an error.
+
+        Args:
+            file_id: The ID of the file to read.
+            target_format: Optional MIME type to export Google Workspace files to.
+                If not specified, uses sensible defaults. Has no effect on non-Workspace files.
+
+        Returns
+        -------
+            GoogleDriveFileContent with the file content and metadata.
+
+        Raises
+        ------
+            GoogleDriveError: If the file cannot be read (not found, permission denied,
+                             binary file, etc.).
+        """
+        file_metadata = await self.get_file_metadata(file_id)
+        original_mime_type = file_metadata.mime_type
+
+        if self._is_binary_mime_type(original_mime_type):
+            raise GoogleDriveError(
+                f"Binary files are not supported for reading. "
+                f"File '{file_metadata.name}' has MIME type '{original_mime_type}'."
+            )
+
+        if original_mime_type == GOOGLE_DRIVE_FOLDER_MIME:
+            raise GoogleDriveError(
+                f"Cannot read content of a folder. '{file_metadata.name}' is a folder, not a file."
+            )
+
+        was_exported = False
+        if original_mime_type in GOOGLE_WORKSPACE_EXPORT_MIMES:
+            export_mime = target_format or GOOGLE_WORKSPACE_EXPORT_MIMES[original_mime_type]
+            content = await self._export_workspace_file(file_id, export_mime)
+            result_mime_type = export_mime
+            was_exported = True
+        elif original_mime_type == PDF_MIME_TYPE:
+            pdf_bytes = await self._download_file_bytes(file_id)
+            content = self._extract_text_from_pdf(pdf_bytes)
+            result_mime_type = "text/plain"
+            was_exported = True
+        else:
+            content = await self._download_file(file_id)
+            result_mime_type = original_mime_type
+
+        return GoogleDriveFileContent(
+            id=file_metadata.id,
+            name=file_metadata.name,
+            mime_type=result_mime_type,
+            content=content,
+            original_mime_type=original_mime_type,
+            was_exported=was_exported,
+            size=file_metadata.size,
+            web_view_link=file_metadata.web_view_link,
+        )
+
     async def __aenter__(self) -> "GoogleDriveClient":
         """Async context manager entry."""
         return self