datarobot-genai 0.2.23__tar.gz → 0.2.26__tar.gz

{datarobot_genai-0.2.23 → datarobot_genai-0.2.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.23
+Version: 0.2.26
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.

{datarobot_genai-0.2.23 → datarobot_genai-0.2.26}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.23"
+version = "0.2.26"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"

{datarobot_genai-0.2.23 → datarobot_genai-0.2.26}/src/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-0.2.23 → datarobot_genai-0.2.26}/src/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-0.2.23 → datarobot_genai-0.2.26}/src/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-0.2.26/src/datarobot_genai/drmcp/tools/predictive/project.py

@@ -0,0 +1,90 @@
+# 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
+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
+
+logger = logging.getLogger(__name__)
+
+
+@dr_mcp_tool(tags={"predictive", "project", "read", "management", "list"})
+async def list_projects() -> ToolResult:
+    """List all DataRobot projects for the authenticated user."""
+    client = get_sdk_client()
+    projects = client.Project.list()
+    projects = {p.id: p.project_name for p in projects}
+
+    return ToolResult(
+        content=(
+            json.dumps(projects, indent=2)
+            if projects
+            else json.dumps({"message": "No projects found."}, indent=2)
+        ),
+        structured_content=projects,
+    )
+
+
+@dr_mcp_tool(tags={"predictive", "project", "read", "data", "info"})
+async def get_project_dataset_by_name(
+    *,
+    project_id: Annotated[str, "The ID of the DataRobot project."] | None = None,
+    dataset_name: Annotated[str, "The name of the dataset to find (e.g., 'training', 'holdout')."]
+    | None = None,
+) -> ToolError | ToolResult:
+    """Get a dataset ID by name for a given project.
+
+    The dataset ID and the dataset type (source or prediction) as a string, or an error message.
+    """
+    if not project_id:
+        return ToolError("Project ID is required.")
+    if not dataset_name:
+        return ToolError("Dataset name is required.")
+
+    client = get_sdk_client()
+    project = client.Project.get(project_id)
+    all_datasets = []
+    source_dataset = project.get_dataset()
+    if source_dataset:
+        all_datasets.append({"type": "source", "dataset": source_dataset})
+    prediction_datasets = project.get_datasets()
+    if prediction_datasets:
+        all_datasets.extend([{"type": "prediction", "dataset": ds} for ds in prediction_datasets])
+    for ds in all_datasets:
+        if dataset_name.lower() in ds["dataset"].name.lower():
+            return ToolResult(
+                content=(
+                    json.dumps(
+                        {
+                            "dataset_id": ds["dataset"].id,
+                            "dataset_type": ds["type"],
+                        },
+                        indent=2,
+                    )
+                ),
+                structured_content={
+                    "dataset_id": ds["dataset"].id,
+                    "dataset_type": ds["type"],
+                },
+            )
+    return ToolResult(
+        content=f"Dataset with name containing '{dataset_name}' not found in project {project_id}.",
+        structured_content={},
+    )