datarobot-genai 0.2.31__tar.gz → 0.2.32__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/PKG-INFO +1 -1
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/pyproject.toml +1 -1
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/exceptions.py +0 -4
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/logging.py +2 -2
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/clients/base.py +1 -1
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/clients/gdrive.py +91 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/gdrive/tools.py +100 -1
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/predictive/data.py +5 -5
- datarobot_genai-0.2.32/src/datarobot_genai/drmcp/tools/predictive/model.py +183 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/predictive/project.py +2 -2
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/predictive/training.py +14 -14
- datarobot_genai-0.2.31/src/datarobot_genai/drmcp/tools/predictive/model.py +0 -148
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/.gitignore +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/AUTHORS +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/LICENSE +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/README.md +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/agents/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/agents/base.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/chat/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/chat/auth.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/chat/client.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/chat/responses.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/cli/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/cli/agent_environment.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/cli/agent_kernel.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/custom_model.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/mcp/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/mcp/common.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/telemetry_agent.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/utils/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/utils/auth.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/utils/urls.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/crewai/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/crewai/agent.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/crewai/base.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/crewai/events.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/crewai/mcp.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/auth.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/clients.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/config.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/config_utils.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/constants.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/credentials.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dr_mcp_server.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dr_mcp_server_logo.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_prompts/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_prompts/controllers.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_prompts/dr_lib.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_prompts/register.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_prompts/utils.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/base.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/default.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/adapters/drum.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/config.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/controllers.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/metadata.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/register.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_agentic_fallback_schema.json +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/deployment/schemas/drum_prediction_fallback_schema.json +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/register.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dynamic_tools/schema.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/mcp_instance.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/memory_management/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/memory_management/manager.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/memory_management/memory_tools.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/routes.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/routes_utils.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/server_life_cycle.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/telemetry.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/tool_config.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/tool_filter.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/utils.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/server.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/clients/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/clients/anthropic.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/clients/dr_gateway.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/clients/openai.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/elicitation_test_tool.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/integration_mcp_server.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/mcp_utils_ete.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/mcp_utils_integration.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/test_interactive.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/tool_base_ete.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/utils.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/clients/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/clients/atlassian.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/clients/confluence.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/clients/jira.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/clients/microsoft_graph.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/clients/s3.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/confluence/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/confluence/tools.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/gdrive/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/jira/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/jira/tools.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/microsoft_graph/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/microsoft_graph/tools.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/predictive/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/predictive/deployment.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/predictive/deployment_info.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/predictive/predict.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/predictive/predict_realtime.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/langgraph/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/langgraph/agent.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/langgraph/mcp.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/llama_index/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/llama_index/agent.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/llama_index/base.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/llama_index/mcp.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/nat/__init__.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/nat/agent.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/nat/datarobot_auth_provider.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/nat/datarobot_llm_clients.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/nat/datarobot_llm_providers.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/nat/datarobot_mcp_client.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/nat/helpers.py +0 -0
- {datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/py.typed +0 -0
|
@@ -20,7 +20,7 @@ from collections.abc import Callable
|
|
|
20
20
|
from typing import Any
|
|
21
21
|
from typing import TypeVar
|
|
22
22
|
|
|
23
|
-
from .exceptions import
|
|
23
|
+
from fastmcp.exceptions import ToolError
|
|
24
24
|
|
|
25
25
|
# Secret patterns to redact from logs
|
|
26
26
|
SECRET_PATTERNS = [
|
|
@@ -93,6 +93,6 @@ def log_execution(func: F) -> F:
|
|
|
93
93
|
return result
|
|
94
94
|
except Exception as e:
|
|
95
95
|
error_msg = _log_error(logger, func.__name__, e, args=args, kwargs=kwargs)
|
|
96
|
-
raise
|
|
96
|
+
raise ToolError(error_msg)
|
|
97
97
|
|
|
98
98
|
return wrapper # type: ignore[return-value]
|
|
@@ -27,7 +27,7 @@ from mcp.types import ListToolsResult
|
|
|
27
27
|
from mcp.types import TextContent
|
|
28
28
|
from openai.types.chat.chat_completion import ChatCompletion
|
|
29
29
|
|
|
30
|
-
from
|
|
30
|
+
from datarobot_genai.drmcp.test_utils.utils import save_response_to_file
|
|
31
31
|
|
|
32
32
|
|
|
33
33
|
class ToolCall:
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/clients/gdrive.py
RENAMED
|
@@ -20,6 +20,7 @@ import logging
|
|
|
20
20
|
import uuid
|
|
21
21
|
from typing import Annotated
|
|
22
22
|
from typing import Any
|
|
23
|
+
from typing import Literal
|
|
23
24
|
|
|
24
25
|
import httpx
|
|
25
26
|
from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
|
|
@@ -821,6 +822,96 @@ class GoogleDriveClient:
|
|
|
821
822
|
headers={"Content-Type": f"multipart/related; boundary={boundary}"},
|
|
822
823
|
)
|
|
823
824
|
|
|
825
|
+
async def manage_access(
|
|
826
|
+
self,
|
|
827
|
+
*,
|
|
828
|
+
file_id: str,
|
|
829
|
+
action: Literal["add", "update", "remove"],
|
|
830
|
+
role: Literal["reader", "commenter", "writer", "fileOrganizer", "organizer", "owner"]
|
|
831
|
+
| None = None,
|
|
832
|
+
email_address: str | None = None,
|
|
833
|
+
permission_id: str | None = None,
|
|
834
|
+
transfer_ownership: bool = False,
|
|
835
|
+
) -> str:
|
|
836
|
+
"""Manage access permissions for a Google Drive file or folder.
|
|
837
|
+
|
|
838
|
+
Adds, updates, or removes sharing permissions on an existing Google Drive
|
|
839
|
+
file or folder using the Google Drive Permissions API.
|
|
840
|
+
|
|
841
|
+
This method supports granting access to users or groups, changing access
|
|
842
|
+
roles, and revoking permissions. Ownership transfer is supported for files
|
|
843
|
+
in "My Drive" when explicitly requested.
|
|
844
|
+
|
|
845
|
+
Args:
|
|
846
|
+
file_id: The ID of the Google Drive file or folder whose permissions
|
|
847
|
+
are being managed.
|
|
848
|
+
action: The permission operation to perform.
|
|
849
|
+
role: The access role to assign or update. Valid values include
|
|
850
|
+
Required for "add" and "update" actions.
|
|
851
|
+
email_address: The email address of the user or group to grant access to.
|
|
852
|
+
Required for the "add" action.
|
|
853
|
+
permission_id: The ID of the permission to update or remove.
|
|
854
|
+
Required for "update" and "remove" actions.
|
|
855
|
+
transfer_ownership: Whether to transfer ownership of the file.
|
|
856
|
+
Only applicable when action="update" and role="owner".
|
|
857
|
+
|
|
858
|
+
Returns
|
|
859
|
+
-------
|
|
860
|
+
Permission id.
|
|
861
|
+
For "add" its newly added permission.
|
|
862
|
+
For "update"/"remove" its previous permission.
|
|
863
|
+
|
|
864
|
+
Raises
|
|
865
|
+
------
|
|
866
|
+
GoogleDriveError: If the permission operation fails (invalid arguments,
|
|
867
|
+
insufficient permissions, resource not found, ownership transfer
|
|
868
|
+
not allowed, rate limited, etc.).
|
|
869
|
+
"""
|
|
870
|
+
if not file_id.strip():
|
|
871
|
+
raise GoogleDriveError("Argument validation error: 'file_id' cannot be empty.")
|
|
872
|
+
|
|
873
|
+
if action == "add" and not email_address:
|
|
874
|
+
raise GoogleDriveError("'email_address' is required for action 'add'.")
|
|
875
|
+
|
|
876
|
+
if action in ("update", "remove") and not permission_id:
|
|
877
|
+
raise GoogleDriveError("'permission_id' is required for action 'update' or 'remove'.")
|
|
878
|
+
|
|
879
|
+
if action != "remove" and not role:
|
|
880
|
+
raise GoogleDriveError("'role' is required for action 'add' or 'update'.")
|
|
881
|
+
|
|
882
|
+
if action == "add":
|
|
883
|
+
response = await self._client.post(
|
|
884
|
+
url=f"/{file_id}/permissions",
|
|
885
|
+
json={
|
|
886
|
+
"type": "user",
|
|
887
|
+
"role": role,
|
|
888
|
+
"emailAddress": email_address,
|
|
889
|
+
},
|
|
890
|
+
params={"sendNotificationEmail": False, "supportsAllDrives": True},
|
|
891
|
+
)
|
|
892
|
+
|
|
893
|
+
elif action == "update":
|
|
894
|
+
response = await self._client.patch(
|
|
895
|
+
url=f"/{file_id}/permissions/{permission_id}",
|
|
896
|
+
json={"role": role},
|
|
897
|
+
params={"transferOwnership": transfer_ownership, "supportsAllDrives": True},
|
|
898
|
+
)
|
|
899
|
+
|
|
900
|
+
elif action == "remove":
|
|
901
|
+
response = await self._client.delete(url=f"/{file_id}/permissions/{permission_id}")
|
|
902
|
+
|
|
903
|
+
else:
|
|
904
|
+
raise GoogleDriveError(f"Invalid action '{action}'")
|
|
905
|
+
|
|
906
|
+
if response.status_code not in (200, 201, 204):
|
|
907
|
+
raise GoogleDriveError(f"Drive API error {response.status_code}: {response.text}")
|
|
908
|
+
|
|
909
|
+
if action == "add":
|
|
910
|
+
return response.json()["id"]
|
|
911
|
+
|
|
912
|
+
# Cannot be null here because of above validators
|
|
913
|
+
return permission_id # type: ignore
|
|
914
|
+
|
|
824
915
|
async def __aenter__(self) -> "GoogleDriveClient":
|
|
825
916
|
"""Async context manager entry."""
|
|
826
917
|
return self
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/gdrive/tools.py
RENAMED
|
@@ -16,6 +16,7 @@
|
|
|
16
16
|
|
|
17
17
|
import logging
|
|
18
18
|
from typing import Annotated
|
|
19
|
+
from typing import Literal
|
|
19
20
|
|
|
20
21
|
from fastmcp.exceptions import ToolError
|
|
21
22
|
from fastmcp.tools.tool import ToolResult
|
|
@@ -33,7 +34,9 @@ from datarobot_genai.drmcp.tools.clients.gdrive import get_gdrive_access_token
|
|
|
33
34
|
logger = logging.getLogger(__name__)
|
|
34
35
|
|
|
35
36
|
|
|
36
|
-
@dr_mcp_tool(
|
|
37
|
+
@dr_mcp_tool(
|
|
38
|
+
tags={"google", "gdrive", "list", "search", "files", "find", "contents"}, enabled=False
|
|
39
|
+
)
|
|
37
40
|
async def gdrive_find_contents(
|
|
38
41
|
*,
|
|
39
42
|
page_size: Annotated[
|
|
@@ -345,3 +348,99 @@ async def gdrive_update_metadata(
|
|
|
345
348
|
content=f"Successfully updated file '{updated_file.name}': {changes_description}.",
|
|
346
349
|
structured_content=updated_file.as_flat_dict(),
|
|
347
350
|
)
|
|
351
|
+
|
|
352
|
+
|
|
353
|
+
@dr_mcp_tool(tags={"google", "gdrive", "manage", "access", "acl"})
|
|
354
|
+
async def gdrive_manage_access(
|
|
355
|
+
*,
|
|
356
|
+
file_id: Annotated[str, "The ID of the file or folder."],
|
|
357
|
+
action: Annotated[Literal["add", "update", "remove"], "The operation to perform."],
|
|
358
|
+
role: Annotated[
|
|
359
|
+
Literal["reader", "commenter", "writer", "fileOrganizer", "organizer", "owner"] | None,
|
|
360
|
+
"The access level.",
|
|
361
|
+
] = None,
|
|
362
|
+
email_address: Annotated[
|
|
363
|
+
str | None, "The email of the user or group (required for 'add')."
|
|
364
|
+
] = None,
|
|
365
|
+
permission_id: Annotated[
|
|
366
|
+
str | None, "The specific permission ID (required for 'update' or 'remove')."
|
|
367
|
+
] = None,
|
|
368
|
+
transfer_ownership: Annotated[
|
|
369
|
+
bool, "Whether to transfer ownership (only for 'update' to 'owner' role)."
|
|
370
|
+
] = False,
|
|
371
|
+
) -> ToolResult:
|
|
372
|
+
"""
|
|
373
|
+
Consolidated tool for sharing files and managing permissions.
|
|
374
|
+
Pushes all logic to the Google Drive API permissions resource (create, update, delete).
|
|
375
|
+
|
|
376
|
+
Usage:
|
|
377
|
+
- Add role: gdrive_manage_access(
|
|
378
|
+
file_id="SomeFileId",
|
|
379
|
+
action="add",
|
|
380
|
+
role="reader",
|
|
381
|
+
email_address="dummy@user.com"
|
|
382
|
+
)
|
|
383
|
+
- Update role: gdrive_manage_access(
|
|
384
|
+
file_id="SomeFileId",
|
|
385
|
+
action="update",
|
|
386
|
+
role="reader",
|
|
387
|
+
permission_id="SomePermissionId"
|
|
388
|
+
)
|
|
389
|
+
- Remove permission: gdrive_manage_access(
|
|
390
|
+
file_id="SomeFileId",
|
|
391
|
+
action="remove",
|
|
392
|
+
permission_id="SomePermissionId"
|
|
393
|
+
)
|
|
394
|
+
"""
|
|
395
|
+
if not file_id or not file_id.strip():
|
|
396
|
+
raise ToolError("Argument validation error: 'file_id' cannot be empty.")
|
|
397
|
+
|
|
398
|
+
if action == "add" and not email_address:
|
|
399
|
+
raise ToolError("'email_address' is required for action 'add'.")
|
|
400
|
+
|
|
401
|
+
if action in ("update", "remove") and not permission_id:
|
|
402
|
+
raise ToolError("'permission_id' is required for action 'update' or 'remove'.")
|
|
403
|
+
|
|
404
|
+
if action != "remove" and not role:
|
|
405
|
+
raise ToolError("'role' is required for action 'add' or 'update'.")
|
|
406
|
+
|
|
407
|
+
access_token = await get_gdrive_access_token()
|
|
408
|
+
if isinstance(access_token, ToolError):
|
|
409
|
+
raise access_token
|
|
410
|
+
|
|
411
|
+
try:
|
|
412
|
+
async with GoogleDriveClient(access_token) as client:
|
|
413
|
+
permission_id = await client.manage_access(
|
|
414
|
+
file_id=file_id,
|
|
415
|
+
action=action,
|
|
416
|
+
role=role,
|
|
417
|
+
email_address=email_address,
|
|
418
|
+
permission_id=permission_id,
|
|
419
|
+
transfer_ownership=transfer_ownership,
|
|
420
|
+
)
|
|
421
|
+
except GoogleDriveError as e:
|
|
422
|
+
logger.error(f"Google Drive permission operation failed: {e}")
|
|
423
|
+
raise ToolError(str(e))
|
|
424
|
+
except Exception as e:
|
|
425
|
+
logger.error(f"Unexpected error changing permissions for Google Drive file {file_id}: {e}")
|
|
426
|
+
raise ToolError(
|
|
427
|
+
f"Unexpected error changing permissions for Google Drive file {file_id}: {str(e)}"
|
|
428
|
+
)
|
|
429
|
+
|
|
430
|
+
# Build response
|
|
431
|
+
structured_content = {"affectedFileId": file_id}
|
|
432
|
+
if action == "add":
|
|
433
|
+
content = (
|
|
434
|
+
f"Successfully added role '{role}' for '{email_address}' for gdrive file '{file_id}'. "
|
|
435
|
+
f"New permission id '{permission_id}'."
|
|
436
|
+
)
|
|
437
|
+
structured_content["newPermissionId"] = permission_id
|
|
438
|
+
elif action == "update":
|
|
439
|
+
content = (
|
|
440
|
+
f"Successfully updated role '{role}' (permission '{permission_id}') "
|
|
441
|
+
f"for gdrive file '{file_id}'."
|
|
442
|
+
)
|
|
443
|
+
else: # action == "remove":
|
|
444
|
+
content = f"Successfully removed permission '{permission_id}' for gdrive file '{file_id}'."
|
|
445
|
+
|
|
446
|
+
return ToolResult(content=content, structured_content=structured_content)
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/predictive/data.py
RENAMED
|
@@ -35,9 +35,9 @@ async def upload_dataset_to_ai_catalog(
|
|
|
35
35
|
) -> ToolError | ToolResult:
|
|
36
36
|
"""Upload a dataset to the DataRobot AI Catalog / Data Registry."""
|
|
37
37
|
if not file_path and not file_url:
|
|
38
|
-
|
|
38
|
+
raise ToolError("Either file_path or file_url must be provided.")
|
|
39
39
|
if file_path and file_url:
|
|
40
|
-
|
|
40
|
+
raise ToolError("Please provide either file_path or file_url, not both.")
|
|
41
41
|
|
|
42
42
|
# Get client
|
|
43
43
|
client = get_sdk_client()
|
|
@@ -47,17 +47,17 @@ async def upload_dataset_to_ai_catalog(
|
|
|
47
47
|
# Does file exist?
|
|
48
48
|
if not os.path.exists(file_path):
|
|
49
49
|
logger.error("File not found: %s", file_path)
|
|
50
|
-
|
|
50
|
+
raise ToolError(f"File not found: {file_path}")
|
|
51
51
|
catalog_item = client.Dataset.create_from_file(file_path)
|
|
52
52
|
else:
|
|
53
53
|
# Does URL exist?
|
|
54
54
|
if file_url is None or not is_valid_url(file_url):
|
|
55
55
|
logger.error("Invalid file URL: %s", file_url)
|
|
56
|
-
|
|
56
|
+
raise ToolError(f"Invalid file URL: {file_url}")
|
|
57
57
|
catalog_item = client.Dataset.create_from_url(file_url)
|
|
58
58
|
|
|
59
59
|
if not catalog_item:
|
|
60
|
-
|
|
60
|
+
raise ToolError("Failed to upload dataset.")
|
|
61
61
|
|
|
62
62
|
return ToolResult(
|
|
63
63
|
content=f"Successfully uploaded dataset: {catalog_item.id}",
|
|
@@ -0,0 +1,183 @@
|
|
|
1
|
+
# Copyright 2025 DataRobot, Inc.
|
|
2
|
+
#
|
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
# you may not use this file except in compliance with the License.
|
|
5
|
+
# You may obtain a copy of the License at
|
|
6
|
+
#
|
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
#
|
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
# See the License for the specific language governing permissions and
|
|
13
|
+
# limitations under the License.
|
|
14
|
+
|
|
15
|
+
import json
|
|
16
|
+
import logging
|
|
17
|
+
from typing import Annotated
|
|
18
|
+
from typing import Any
|
|
19
|
+
|
|
20
|
+
from datarobot.models.model import Model
|
|
21
|
+
from fastmcp.exceptions import ToolError
|
|
22
|
+
from fastmcp.tools.tool import ToolResult
|
|
23
|
+
|
|
24
|
+
from datarobot_genai.drmcp.core.clients import get_sdk_client
|
|
25
|
+
from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
|
|
26
|
+
|
|
27
|
+
logger = logging.getLogger(__name__)
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def model_to_dict(model: Any) -> dict[str, Any]:
|
|
31
|
+
"""Convert a DataRobot Model object to a dictionary."""
|
|
32
|
+
try:
|
|
33
|
+
return {
|
|
34
|
+
"id": model.id,
|
|
35
|
+
"model_type": model.model_type,
|
|
36
|
+
"metrics": model.metrics,
|
|
37
|
+
}
|
|
38
|
+
except AttributeError as e:
|
|
39
|
+
logger.warning(f"Failed to access some model attributes: {e}")
|
|
40
|
+
# Return minimal information if some attributes are not accessible
|
|
41
|
+
return {
|
|
42
|
+
"id": getattr(model, "id", "unknown"),
|
|
43
|
+
"model_type": getattr(model, "model_type", "unknown"),
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
class ModelEncoder(json.JSONEncoder):
|
|
48
|
+
"""Custom JSON encoder for DataRobot Model objects."""
|
|
49
|
+
|
|
50
|
+
def default(self, obj: Any) -> Any:
|
|
51
|
+
if isinstance(obj, Model):
|
|
52
|
+
return model_to_dict(obj)
|
|
53
|
+
return super().default(obj)
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
@dr_mcp_tool(tags={"predictive", "model", "read", "management", "info"})
|
|
57
|
+
async def get_best_model(
|
|
58
|
+
*,
|
|
59
|
+
project_id: Annotated[str, "The DataRobot project ID"] | None = None,
|
|
60
|
+
metric: Annotated[str, "The metric to use for best model selection (e.g., 'AUC', 'LogLoss')"]
|
|
61
|
+
| None = None,
|
|
62
|
+
) -> ToolError | ToolResult:
|
|
63
|
+
"""Get the best model for a DataRobot project, optionally by a specific metric."""
|
|
64
|
+
if not project_id:
|
|
65
|
+
raise ToolError("Project ID must be provided")
|
|
66
|
+
|
|
67
|
+
client = get_sdk_client()
|
|
68
|
+
project = client.Project.get(project_id)
|
|
69
|
+
if not project:
|
|
70
|
+
raise ToolError(f"Project with ID {project_id} not found.")
|
|
71
|
+
|
|
72
|
+
leaderboard = project.get_models()
|
|
73
|
+
if not leaderboard:
|
|
74
|
+
raise ToolError("No models found for this project.")
|
|
75
|
+
|
|
76
|
+
if metric:
|
|
77
|
+
reverse_sort = metric.upper() in [
|
|
78
|
+
"AUC",
|
|
79
|
+
"ACCURACY",
|
|
80
|
+
"F1",
|
|
81
|
+
"PRECISION",
|
|
82
|
+
"RECALL",
|
|
83
|
+
]
|
|
84
|
+
leaderboard = sorted(
|
|
85
|
+
leaderboard,
|
|
86
|
+
key=lambda m: m.metrics.get(metric, {}).get(
|
|
87
|
+
"validation", float("-inf") if reverse_sort else float("inf")
|
|
88
|
+
),
|
|
89
|
+
reverse=reverse_sort,
|
|
90
|
+
)
|
|
91
|
+
logger.info(f"Sorted models by metric: {metric}")
|
|
92
|
+
|
|
93
|
+
best_model = leaderboard[0]
|
|
94
|
+
logger.info(f"Found best model {best_model.id} for project {project_id}")
|
|
95
|
+
|
|
96
|
+
metric_info = ""
|
|
97
|
+
metric_value = None
|
|
98
|
+
|
|
99
|
+
if metric and best_model.metrics and metric in best_model.metrics:
|
|
100
|
+
metric_value = best_model.metrics[metric].get("validation")
|
|
101
|
+
if metric_value is not None:
|
|
102
|
+
metric_info = f" with {metric}: {metric_value:.2f}"
|
|
103
|
+
|
|
104
|
+
# Include full metrics in the response
|
|
105
|
+
best_model_dict = model_to_dict(best_model)
|
|
106
|
+
best_model_dict["metric"] = metric
|
|
107
|
+
best_model_dict["metric_value"] = metric_value
|
|
108
|
+
|
|
109
|
+
# Format metrics for human-readable content
|
|
110
|
+
metrics_text = ""
|
|
111
|
+
if best_model.metrics:
|
|
112
|
+
metrics_list = []
|
|
113
|
+
for metric_name, metric_data in best_model.metrics.items():
|
|
114
|
+
if isinstance(metric_data, dict) and "validation" in metric_data:
|
|
115
|
+
val = metric_data["validation"]
|
|
116
|
+
if val is not None:
|
|
117
|
+
metrics_list.append(f"{metric_name}: {val:.4f}")
|
|
118
|
+
if metrics_list:
|
|
119
|
+
metrics_text = "\nPerformance metrics:\n" + "\n".join(f" - {m}" for m in metrics_list)
|
|
120
|
+
|
|
121
|
+
return ToolResult(
|
|
122
|
+
content=f"Best model: {best_model.model_type}{metric_info}{metrics_text}",
|
|
123
|
+
structured_content={
|
|
124
|
+
"project_id": project_id,
|
|
125
|
+
"best_model": best_model_dict,
|
|
126
|
+
},
|
|
127
|
+
)
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
@dr_mcp_tool(tags={"predictive", "model", "read", "scoring", "dataset"})
|
|
131
|
+
async def score_dataset_with_model(
|
|
132
|
+
*,
|
|
133
|
+
project_id: Annotated[str, "The DataRobot project ID"] | None = None,
|
|
134
|
+
model_id: Annotated[str, "The DataRobot model ID"] | None = None,
|
|
135
|
+
dataset_url: Annotated[str, "The dataset URL"] | None = None,
|
|
136
|
+
) -> ToolError | ToolResult:
|
|
137
|
+
"""Score a dataset using a specific DataRobot model."""
|
|
138
|
+
if not project_id:
|
|
139
|
+
raise ToolError("Project ID must be provided")
|
|
140
|
+
if not model_id:
|
|
141
|
+
raise ToolError("Model ID must be provided")
|
|
142
|
+
if not dataset_url:
|
|
143
|
+
raise ToolError("Dataset URL must be provided")
|
|
144
|
+
|
|
145
|
+
client = get_sdk_client()
|
|
146
|
+
project = client.Project.get(project_id)
|
|
147
|
+
model = client.Model.get(project, model_id)
|
|
148
|
+
job = model.score(dataset_url)
|
|
149
|
+
|
|
150
|
+
return ToolResult(
|
|
151
|
+
content=f"Scoring job started: {job.id}",
|
|
152
|
+
structured_content={
|
|
153
|
+
"scoring_job_id": job.id,
|
|
154
|
+
"project_id": project_id,
|
|
155
|
+
"model_id": model_id,
|
|
156
|
+
"dataset_url": dataset_url,
|
|
157
|
+
},
|
|
158
|
+
)
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
@dr_mcp_tool(tags={"predictive", "model", "read", "management", "list"})
|
|
162
|
+
async def list_models(
|
|
163
|
+
*,
|
|
164
|
+
project_id: Annotated[str, "The DataRobot project ID"] | None = None,
|
|
165
|
+
) -> ToolError | ToolResult:
|
|
166
|
+
"""List all models in a project."""
|
|
167
|
+
if not project_id:
|
|
168
|
+
raise ToolError("Project ID must be provided")
|
|
169
|
+
|
|
170
|
+
client = get_sdk_client()
|
|
171
|
+
project = client.Project.get(project_id)
|
|
172
|
+
models = project.get_models()
|
|
173
|
+
|
|
174
|
+
return ToolResult(
|
|
175
|
+
content=(
|
|
176
|
+
f"Found {len(models)} models in project {project_id}, here are the details:\n"
|
|
177
|
+
f"{json.dumps(models, indent=2, cls=ModelEncoder)}"
|
|
178
|
+
),
|
|
179
|
+
structured_content={
|
|
180
|
+
"project_id": project_id,
|
|
181
|
+
"models": [model_to_dict(model) for model in models],
|
|
182
|
+
},
|
|
183
|
+
)
|
|
@@ -54,9 +54,9 @@ async def get_project_dataset_by_name(
|
|
|
54
54
|
The dataset ID and the dataset type (source or prediction) as a string, or an error message.
|
|
55
55
|
"""
|
|
56
56
|
if not project_id:
|
|
57
|
-
|
|
57
|
+
raise ToolError("Project ID is required.")
|
|
58
58
|
if not dataset_name:
|
|
59
|
-
|
|
59
|
+
raise ToolError("Dataset name is required.")
|
|
60
60
|
|
|
61
61
|
client = get_sdk_client()
|
|
62
62
|
project = client.Project.get(project_id)
|
|
@@ -63,7 +63,7 @@ async def analyze_dataset(
|
|
|
63
63
|
) -> ToolError | ToolResult:
|
|
64
64
|
"""Analyze a dataset to understand its structure and potential use cases."""
|
|
65
65
|
if not dataset_id:
|
|
66
|
-
|
|
66
|
+
raise ToolError("Dataset ID must be provided")
|
|
67
67
|
|
|
68
68
|
client = get_sdk_client()
|
|
69
69
|
dataset = client.Dataset.get(dataset_id)
|
|
@@ -116,7 +116,7 @@ async def suggest_use_cases(
|
|
|
116
116
|
) -> ToolError | ToolResult:
|
|
117
117
|
"""Analyze a dataset and suggest potential machine learning use cases."""
|
|
118
118
|
if not dataset_id:
|
|
119
|
-
|
|
119
|
+
raise ToolError("Dataset ID must be provided")
|
|
120
120
|
|
|
121
121
|
client = get_sdk_client()
|
|
122
122
|
dataset = client.Dataset.get(dataset_id)
|
|
@@ -148,7 +148,7 @@ async def get_exploratory_insights(
|
|
|
148
148
|
) -> ToolError | ToolResult:
|
|
149
149
|
"""Generate exploratory data insights for a dataset."""
|
|
150
150
|
if not dataset_id:
|
|
151
|
-
|
|
151
|
+
raise ToolError("Dataset ID must be provided")
|
|
152
152
|
|
|
153
153
|
client = get_sdk_client()
|
|
154
154
|
dataset = client.Dataset.get(dataset_id)
|
|
@@ -481,9 +481,9 @@ async def start_autopilot(
|
|
|
481
481
|
|
|
482
482
|
if not project_id:
|
|
483
483
|
if not dataset_url and not dataset_id:
|
|
484
|
-
|
|
484
|
+
raise ToolError("Either dataset_url or dataset_id must be provided")
|
|
485
485
|
if dataset_url and dataset_id:
|
|
486
|
-
|
|
486
|
+
raise ToolError("Please provide either dataset_url or dataset_id, not both")
|
|
487
487
|
|
|
488
488
|
if dataset_url:
|
|
489
489
|
dataset = client.Dataset.create_from_url(dataset_url)
|
|
@@ -497,7 +497,7 @@ async def start_autopilot(
|
|
|
497
497
|
project = client.Project.get(project_id)
|
|
498
498
|
|
|
499
499
|
if not target:
|
|
500
|
-
|
|
500
|
+
raise ToolError("Target variable must be specified")
|
|
501
501
|
|
|
502
502
|
try:
|
|
503
503
|
# Start modeling
|
|
@@ -517,7 +517,7 @@ async def start_autopilot(
|
|
|
517
517
|
)
|
|
518
518
|
|
|
519
519
|
except Exception as e:
|
|
520
|
-
|
|
520
|
+
raise ToolError(
|
|
521
521
|
content=json.dumps(
|
|
522
522
|
{
|
|
523
523
|
"error": f"Failed to start Autopilot: {str(e)}",
|
|
@@ -546,9 +546,9 @@ async def get_model_roc_curve(
|
|
|
546
546
|
) -> ToolError | ToolResult:
|
|
547
547
|
"""Get detailed ROC curve for a specific model."""
|
|
548
548
|
if not project_id:
|
|
549
|
-
|
|
549
|
+
raise ToolError("Project ID must be provided")
|
|
550
550
|
if not model_id:
|
|
551
|
-
|
|
551
|
+
raise ToolError("Model ID must be provided")
|
|
552
552
|
|
|
553
553
|
client = get_sdk_client()
|
|
554
554
|
project = client.Project.get(project_id)
|
|
@@ -587,7 +587,7 @@ async def get_model_roc_curve(
|
|
|
587
587
|
structured_content={"data": roc_data},
|
|
588
588
|
)
|
|
589
589
|
except Exception as e:
|
|
590
|
-
|
|
590
|
+
raise ToolError(f"Failed to get ROC curve: {str(e)}")
|
|
591
591
|
|
|
592
592
|
|
|
593
593
|
@dr_mcp_tool(tags={"predictive", "training", "read", "model", "evaluation"})
|
|
@@ -598,9 +598,9 @@ async def get_model_feature_impact(
|
|
|
598
598
|
) -> ToolError | ToolResult:
|
|
599
599
|
"""Get detailed feature impact for a specific model."""
|
|
600
600
|
if not project_id:
|
|
601
|
-
|
|
601
|
+
raise ToolError("Project ID must be provided")
|
|
602
602
|
if not model_id:
|
|
603
|
-
|
|
603
|
+
raise ToolError("Model ID must be provided")
|
|
604
604
|
|
|
605
605
|
client = get_sdk_client()
|
|
606
606
|
project = client.Project.get(project_id)
|
|
@@ -631,9 +631,9 @@ async def get_model_lift_chart(
|
|
|
631
631
|
) -> ToolError | ToolResult:
|
|
632
632
|
"""Get detailed lift chart for a specific model."""
|
|
633
633
|
if not project_id:
|
|
634
|
-
|
|
634
|
+
raise ToolError("Project ID must be provided")
|
|
635
635
|
if not model_id:
|
|
636
|
-
|
|
636
|
+
raise ToolError("Model ID must be provided")
|
|
637
637
|
|
|
638
638
|
client = get_sdk_client()
|
|
639
639
|
project = client.Project.get(project_id)
|
|
@@ -1,148 +0,0 @@
|
|
|
1
|
-
# Copyright 2025 DataRobot, Inc.
|
|
2
|
-
#
|
|
3
|
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
-
# you may not use this file except in compliance with the License.
|
|
5
|
-
# You may obtain a copy of the License at
|
|
6
|
-
#
|
|
7
|
-
# http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
-
#
|
|
9
|
-
# Unless required by applicable law or agreed to in writing, software
|
|
10
|
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
-
# See the License for the specific language governing permissions and
|
|
13
|
-
# limitations under the License.
|
|
14
|
-
|
|
15
|
-
import json
|
|
16
|
-
import logging
|
|
17
|
-
from typing import Any
|
|
18
|
-
|
|
19
|
-
from datarobot.models.model import Model
|
|
20
|
-
|
|
21
|
-
from datarobot_genai.drmcp.core.clients import get_sdk_client
|
|
22
|
-
from datarobot_genai.drmcp.core.mcp_instance import dr_mcp_tool
|
|
23
|
-
|
|
24
|
-
logger = logging.getLogger(__name__)
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
def model_to_dict(model: Any) -> dict[str, Any]:
|
|
28
|
-
"""Convert a DataRobot Model object to a dictionary."""
|
|
29
|
-
try:
|
|
30
|
-
return {
|
|
31
|
-
"id": model.id,
|
|
32
|
-
"model_type": model.model_type,
|
|
33
|
-
"metrics": model.metrics,
|
|
34
|
-
}
|
|
35
|
-
except AttributeError as e:
|
|
36
|
-
logger.warning(f"Failed to access some model attributes: {e}")
|
|
37
|
-
# Return minimal information if some attributes are not accessible
|
|
38
|
-
return {
|
|
39
|
-
"id": getattr(model, "id", "unknown"),
|
|
40
|
-
"model_type": getattr(model, "model_type", "unknown"),
|
|
41
|
-
}
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
class ModelEncoder(json.JSONEncoder):
|
|
45
|
-
"""Custom JSON encoder for DataRobot Model objects."""
|
|
46
|
-
|
|
47
|
-
def default(self, obj: Any) -> Any:
|
|
48
|
-
if isinstance(obj, Model):
|
|
49
|
-
return model_to_dict(obj)
|
|
50
|
-
return super().default(obj)
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
@dr_mcp_tool(tags={"model", "management", "info"})
|
|
54
|
-
async def get_best_model(project_id: str, metric: str | None = None) -> str:
|
|
55
|
-
"""
|
|
56
|
-
Get the best model for a DataRobot project, optionally by a specific metric.
|
|
57
|
-
|
|
58
|
-
Args:
|
|
59
|
-
project_id: The ID of the DataRobot project.
|
|
60
|
-
metric: (Optional) The metric to use for best model selection (e.g., 'AUC', 'LogLoss').
|
|
61
|
-
|
|
62
|
-
Returns
|
|
63
|
-
-------
|
|
64
|
-
A formatted string describing the best model.
|
|
65
|
-
|
|
66
|
-
Raises
|
|
67
|
-
------
|
|
68
|
-
Exception: If project not found or no models exist in the project.
|
|
69
|
-
"""
|
|
70
|
-
client = get_sdk_client()
|
|
71
|
-
project = client.Project.get(project_id)
|
|
72
|
-
if not project:
|
|
73
|
-
logger.error(f"Project with ID {project_id} not found")
|
|
74
|
-
raise Exception(f"Project with ID {project_id} not found.")
|
|
75
|
-
|
|
76
|
-
leaderboard = project.get_models()
|
|
77
|
-
if not leaderboard:
|
|
78
|
-
logger.info(f"No models found for project {project_id}")
|
|
79
|
-
raise Exception("No models found for this project.")
|
|
80
|
-
|
|
81
|
-
if metric:
|
|
82
|
-
reverse_sort = metric.upper() in [
|
|
83
|
-
"AUC",
|
|
84
|
-
"ACCURACY",
|
|
85
|
-
"F1",
|
|
86
|
-
"PRECISION",
|
|
87
|
-
"RECALL",
|
|
88
|
-
]
|
|
89
|
-
leaderboard = sorted(
|
|
90
|
-
leaderboard,
|
|
91
|
-
key=lambda m: m.metrics.get(metric, {}).get(
|
|
92
|
-
"validation", float("-inf") if reverse_sort else float("inf")
|
|
93
|
-
),
|
|
94
|
-
reverse=reverse_sort,
|
|
95
|
-
)
|
|
96
|
-
logger.info(f"Sorted models by metric: {metric}")
|
|
97
|
-
|
|
98
|
-
best_model = leaderboard[0]
|
|
99
|
-
logger.info(f"Found best model {best_model.id} for project {project_id}")
|
|
100
|
-
|
|
101
|
-
# Format the response as a human-readable string
|
|
102
|
-
metric_info = ""
|
|
103
|
-
if metric and best_model.metrics and metric in best_model.metrics:
|
|
104
|
-
metric_value = best_model.metrics[metric].get("validation")
|
|
105
|
-
if metric_value is not None:
|
|
106
|
-
metric_info = f" with {metric}: {metric_value:.2f}"
|
|
107
|
-
|
|
108
|
-
return f"Best model: {best_model.model_type}{metric_info}"
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
@dr_mcp_tool(tags={"model", "prediction", "scoring"})
|
|
112
|
-
async def score_dataset_with_model(project_id: str, model_id: str, dataset_url: str) -> str:
|
|
113
|
-
"""
|
|
114
|
-
Score a dataset using a specific DataRobot model.
|
|
115
|
-
|
|
116
|
-
Args:
|
|
117
|
-
project_id: The ID of the DataRobot project.
|
|
118
|
-
model_id: The ID of the DataRobot model to use for scoring.
|
|
119
|
-
dataset_url: The URL to the dataset to score (must be accessible to DataRobot).
|
|
120
|
-
|
|
121
|
-
Returns
|
|
122
|
-
-------
|
|
123
|
-
A string summary of the scoring job or a meaningful error message.
|
|
124
|
-
"""
|
|
125
|
-
client = get_sdk_client()
|
|
126
|
-
project = client.Project.get(project_id)
|
|
127
|
-
model = client.Model.get(project, model_id)
|
|
128
|
-
job = model.score(dataset_url)
|
|
129
|
-
logger.info(f"Started scoring job {job.id} for model {model_id}")
|
|
130
|
-
return f"Scoring job started: {job.id}"
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
@dr_mcp_tool(tags={"model", "management", "list"})
|
|
134
|
-
async def list_models(project_id: str) -> str:
|
|
135
|
-
"""
|
|
136
|
-
List all models in a project.
|
|
137
|
-
|
|
138
|
-
Args:
|
|
139
|
-
project_id: The ID of the DataRobot project.
|
|
140
|
-
|
|
141
|
-
Returns
|
|
142
|
-
-------
|
|
143
|
-
A string summary of the models in the project.
|
|
144
|
-
"""
|
|
145
|
-
client = get_sdk_client()
|
|
146
|
-
project = client.Project.get(project_id)
|
|
147
|
-
models = project.get_models()
|
|
148
|
-
return json.dumps(models, indent=2, cls=ModelEncoder)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/agents/__init__.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/chat/responses.py
RENAMED
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/cli/agent_environment.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/cli/agent_kernel.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/telemetry_agent.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/core/utils/__init__.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/__init__.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/config_utils.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/constants.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/credentials.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/dr_mcp_server.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/mcp_instance.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/routes_utils.py
RENAMED
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/telemetry.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/tool_config.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/core/tool_filter.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/__init__.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/test_utils/utils.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/__init__.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/clients/jira.py
RENAMED
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/clients/s3.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/gdrive/__init__.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/jira/__init__.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/drmcp/tools/jira/tools.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/llama_index/__init__.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/nat/datarobot_auth_provider.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/nat/datarobot_llm_clients.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/nat/datarobot_llm_providers.py
RENAMED
|
File without changes
|
{datarobot_genai-0.2.31 → datarobot_genai-0.2.32}/src/datarobot_genai/nat/datarobot_mcp_client.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|