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

datarobot_genai/drmcp/test_utils/mcp_utils_ete.py

@@ -68,6 +68,26 @@ def get_openai_llm_client_config() -> dict[str, str]:
     return config
 
 
+def get_dr_llm_gateway_client_config() -> dict[str, str]:
+    """Get DataRobot LLM Gateway client configuration."""
+    datarobot_api_token = os.environ.get("DATAROBOT_API_TOKEN")
+    datarobot_endpoint = os.environ.get("DATAROBOT_ENDPOINT")
+    save_llm_responses = os.environ.get("SAVE_LLM_RESPONSES", "false").lower() == "true"
+
+    if not datarobot_api_token:
+        raise ValueError("Missing required environment variable: DATAROBOT_API_TOKEN")
+
+    config: dict[str, str] = {
+        "datarobot_api_token": datarobot_api_token,
+        "save_llm_responses": str(save_llm_responses),
+    }
+
+    if datarobot_endpoint:
+        config["datarobot_endpoint"] = datarobot_endpoint
+
+    return config
+
+
 def get_headers() -> dict[str, str]:
     # When the MCP server is deployed in DataRobot, we have to include the API token in headers for
     # authentication.

datarobot_genai/drmcp/test_utils/test_interactive.py

@@ -40,40 +40,40 @@ from mcp.types import ElicitResult
 
 from datarobot_genai.drmcp import get_dr_mcp_server_url
 from datarobot_genai.drmcp import get_headers
-from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import LLMMCPClient
-from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import LLMResponse
-from datarobot_genai.drmcp.test_utils.openai_llm_mcp_client import ToolCall
+from datarobot_genai.drmcp.test_utils.clients.base import LLMResponse
+from datarobot_genai.drmcp.test_utils.clients.base import ToolCall
+from datarobot_genai.drmcp.test_utils.clients.dr_gateway import DRLLMGatewayMCPClient
 
 # Re-export for backwards compatibility
-__all__ = ["LLMMCPClient", "LLMResponse", "ToolCall", "test_mcp_interactive"]
+__all__ = ["DRLLMGatewayMCPClient", "LLMResponse", "ToolCall", "test_mcp_interactive"]
 
 
 async def test_mcp_interactive() -> None:
     """Test the MCP server interactively with LLM agent."""
     # Check for required environment variables
-    openai_api_key = os.environ.get("OPENAI_API_KEY")
-    if not openai_api_key:
-        print("❌ Error: OPENAI_API_KEY environment variable is required")
+    datarobot_api_token = os.environ.get("DATAROBOT_API_TOKEN")
+    if not datarobot_api_token:
+        print("❌ Error: DATAROBOT_API_TOKEN environment variable is required")
         print("Please set it in your .env file or export it")
         return
 
-    # Optional Azure OpenAI settings
-    openai_api_base = os.environ.get("OPENAI_API_BASE")
-    openai_api_deployment_id = os.environ.get("OPENAI_API_DEPLOYMENT_ID")
-    openai_api_version = os.environ.get("OPENAI_API_VERSION")
+    # Optional DataRobot settings
+    datarobot_endpoint = os.environ.get("DATAROBOT_ENDPOINT")
+    model = os.environ.get("MODEL")
 
     print("🤖 Initializing LLM MCP Client...")
 
     # Initialize the LLM client with elicitation handler
     config = {
-        "openai_api_key": openai_api_key,
-        "openai_api_base": openai_api_base,
-        "openai_api_deployment_id": openai_api_deployment_id,
-        "openai_api_version": openai_api_version,
+        "datarobot_api_token": datarobot_api_token,
         "save_llm_responses": False,
     }
+    if datarobot_endpoint:
+        config["datarobot_endpoint"] = datarobot_endpoint
+    if model:
+        config["model"] = model
 
-    llm_client = LLMMCPClient(str(config))
+    llm_client = DRLLMGatewayMCPClient(str(config))
 
     # Get MCP server URL
     mcp_server_url = get_dr_mcp_server_url()

datarobot_genai/drmcp/test_utils/tool_base_ete.py

@@ -17,7 +17,7 @@ from typing import Any
 
 from pydantic import BaseModel
 
-from .openai_llm_mcp_client import LLMResponse
+from .clients.base import LLMResponse
 
 
 class ToolCallTestExpectations(BaseModel):
@@ -39,6 +39,54 @@ class ETETestExpectations(BaseModel):
 SHOULD_NOT_BE_EMPTY = "SHOULD_NOT_BE_EMPTY"
 
 
+def _extract_structured_content(tool_result: str) -> Any:
+    r"""
+    Extract and parse structured content from tool result string.
+
+    Tool results are formatted as:
+    "Content: {content}\nStructured content: {structured_content}"
+
+    Structured content can be:
+    1. A JSON object with a "result" key: {"result": "..."} or {"result": "{...}"}
+    2. A direct JSON object: {"key": "value", ...}
+    3. Empty or missing
+
+    Args:
+        tool_result: The tool result string
+
+    Returns
+    -------
+        Parsed structured content, or None if not available
+    """
+    # Early returns for invalid inputs
+    if not tool_result or "Structured content: " not in tool_result:
+        return None
+
+    structured_part = tool_result.split("Structured content: ", 1)[1].strip()
+    # Parse JSON, return None on failure or empty structured_part
+    if not structured_part:
+        return None
+    try:
+        structured_data = json.loads(structured_part)
+    except json.JSONDecodeError:
+        return None
+
+    # If structured data has a "result" key, extract and parse that
+    if isinstance(structured_data, dict) and "result" in structured_data:
+        result_value = structured_data["result"]
+        # If result is a JSON string (starts with { or [), try to parse it
+        if isinstance(result_value, str) and result_value.strip().startswith(("{", "[")):
+            try:
+                parsed_result = json.loads(result_value)
+            except json.JSONDecodeError:
+                parsed_result = result_value  # Return string as-is if parsing fails
+            return parsed_result
+        return result_value  # Return result value directly
+
+    # If it's a direct JSON object (not wrapped in {"result": ...}), return it as-is
+    return structured_data
+
+
 def _check_dict_has_keys(
     expected: dict[str, Any],
     actual: dict[str, Any] | list[dict[str, Any]],
@@ -130,7 +178,26 @@ class ToolBaseE2E:
                             f"result, but got: {response.tool_results[i]}"
                         )
                     else:
-                        actual_result = json.loads(response.tool_results[i])
+                        actual_result = _extract_structured_content(response.tool_results[i])
+                        if actual_result is None:
+                            # Fallback: try to parse the entire tool result as JSON
+                            try:
+                                actual_result = json.loads(response.tool_results[i])
+                            except json.JSONDecodeError:
+                                # If that fails, try to extract content part
+                                if "Content: " in response.tool_results[i]:
+                                    content_part = response.tool_results[i].split("Content: ", 1)[1]
+                                    if "\nStructured content: " in content_part:
+                                        content_part = content_part.split(
+                                            "\nStructured content: ", 1
+                                        )[0]
+                                    try:
+                                        actual_result = json.loads(content_part.strip())
+                                    except json.JSONDecodeError:
+                                        raise AssertionError(
+                                            f"Could not parse tool result for "
+                                            f"{tool_call.tool_name}: {response.tool_results[i]}"
+                                        )
                         assert _check_dict_has_keys(expected_result, actual_result), (
                             f"Should have called {tool_call.tool_name} tool with the correct "
                             f"result structure, but got: {response.tool_results[i]}"

datarobot_genai/drmcp/test_utils/utils.py

@@ -19,7 +19,7 @@ from typing import TYPE_CHECKING
 from typing import Any
 
 if TYPE_CHECKING:
-    from .openai_llm_mcp_client import LLMResponse
+    from .clients.base import LLMResponse
 
 from dotenv import load_dotenv
 

datarobot_genai/drmcp/tools/clients/gdrive.py

@@ -15,9 +15,12 @@
 """Google Drive API Client and utilities for OAuth."""
 
 import io
+import json
 import logging
+import uuid
 from typing import Annotated
 from typing import Any
+from typing import Literal
 
 import httpx
 from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
@@ -31,7 +34,17 @@ from datarobot_genai.drmcp.core.auth import get_access_token
 
 logger = logging.getLogger(__name__)
 
-SUPPORTED_FIELDS = {"id", "name", "size", "mimeType", "webViewLink", "createdTime", "modifiedTime"}
+SUPPORTED_FIELDS = {
+    "id",
+    "name",
+    "size",
+    "mimeType",
+    "webViewLink",
+    "createdTime",
+    "modifiedTime",
+    "starred",
+    "trashed",
+}
 SUPPORTED_FIELDS_STR = ",".join(SUPPORTED_FIELDS)
 DEFAULT_FIELDS = f"nextPageToken,files({SUPPORTED_FIELDS_STR})"
 GOOGLE_DRIVE_FOLDER_MIME = "application/vnd.google-apps.folder"
@@ -45,6 +58,12 @@ GOOGLE_WORKSPACE_EXPORT_MIMES: dict[str, str] = {
     "application/vnd.google-apps.presentation": "text/plain",
 }
 
+# MIME type mappings for content conversion during upload to Google Workspace formats
+UPLOAD_CONTENT_TYPES: dict[str, str] = {
+    "application/vnd.google-apps.document": "text/plain",
+    "application/vnd.google-apps.spreadsheet": "text/csv",
+}
+
 BINARY_MIME_PREFIXES = (
     "image/",
     "audio/",
@@ -111,6 +130,8 @@ class GoogleDriveFile(BaseModel):
     web_view_link: Annotated[str | None, Field(alias="webViewLink")] = None
     created_time: Annotated[str | None, Field(alias="createdTime")] = None
     modified_time: Annotated[str | None, Field(alias="modifiedTime")] = None
+    starred: bool | None = None
+    trashed: bool | None = None
 
     model_config = ConfigDict(populate_by_name=True)
 
@@ -125,8 +146,31 @@ class GoogleDriveFile(BaseModel):
             web_view_link=data.get("webViewLink"),
             created_time=data.get("createdTime"),
             modified_time=data.get("modifiedTime"),
+            starred=data.get("starred"),
+            trashed=data.get("trashed"),
         )
 
+    def as_flat_dict(self) -> dict[str, Any]:
+        """Return a flat dictionary representation of the file."""
+        result: dict[str, Any] = {
+            "id": self.id,
+            "name": self.name,
+            "mimeType": self.mime_type,
+        }
+        if self.size is not None:
+            result["size"] = self.size
+        if self.web_view_link is not None:
+            result["webViewLink"] = self.web_view_link
+        if self.created_time is not None:
+            result["createdTime"] = self.created_time
+        if self.modified_time is not None:
+            result["modifiedTime"] = self.modified_time
+        if self.starred is not None:
+            result["starred"] = self.starred
+        if self.trashed is not None:
+            result["trashed"] = self.trashed
+        return result
+
 
 class PaginatedResult(BaseModel):
     """Result of a paginated API call."""
@@ -432,6 +476,66 @@ class GoogleDriveClient:
         response.raise_for_status()
         return GoogleDriveFile.from_api_response(response.json())
 
+    async def update_file_metadata(
+        self,
+        file_id: str,
+        new_name: str | None = None,
+        starred: bool | None = None,
+        trashed: bool | None = None,
+    ) -> GoogleDriveFile:
+        """Update file metadata in Google Drive.
+
+        Args:
+            file_id: The ID of the file to update.
+            new_name: A new name to rename the file. Must not be empty or whitespace.
+            starred: Set to True to star the file or False to unstar it.
+            trashed: Set to True to trash the file or False to restore it.
+
+        Returns
+        -------
+            GoogleDriveFile with updated metadata.
+
+        Raises
+        ------
+            GoogleDriveError: If no update fields are provided, file is not found,
+                             access is denied, or the request is invalid.
+        """
+        if new_name is None and starred is None and trashed is None:
+            raise GoogleDriveError(
+                "At least one of new_name, starred, or trashed must be provided."
+            )
+
+        if new_name is not None and not new_name.strip():
+            raise GoogleDriveError("new_name cannot be empty or whitespace.")
+
+        body: dict[str, Any] = {}
+        if new_name is not None:
+            body["name"] = new_name
+        if starred is not None:
+            body["starred"] = starred
+        if trashed is not None:
+            body["trashed"] = trashed
+
+        response = await self._client.patch(
+            f"/{file_id}",
+            json=body,
+            params={"fields": SUPPORTED_FIELDS_STR, "supportsAllDrives": "true"},
+        )
+
+        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 permission to update file '{file_id}'."
+            )
+        if response.status_code == 400:
+            raise GoogleDriveError("Bad request: invalid parameters for file update.")
+        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.
 
@@ -599,6 +703,215 @@ class GoogleDriveClient:
             web_view_link=file_metadata.web_view_link,
         )
 
+    async def create_file(
+        self,
+        name: str,
+        mime_type: str,
+        parent_id: str | None = None,
+        initial_content: str | None = None,
+    ) -> GoogleDriveFile:
+        """Create a new file or folder in Google Drive.
+
+        Creates a new file with the specified name and MIME type. Optionally places
+        it in a specific folder and populates it with initial content.
+
+        For Google Workspace files (Docs, Sheets), the Drive API automatically
+        converts plain text content to the appropriate format.
+
+        Args:
+            name: The name for the new file or folder.
+            mime_type: The MIME type of the file (e.g., 'text/plain',
+                'application/vnd.google-apps.document',
+                'application/vnd.google-apps.folder').
+            parent_id: Optional ID of the parent folder. If not specified,
+                the file is created in the root of the user's Drive.
+            initial_content: Optional text content to populate the file.
+                Ignored for folders.
+
+        Returns
+        -------
+            GoogleDriveFile with the created file's metadata.
+
+        Raises
+        ------
+            GoogleDriveError: If file creation fails (permission denied,
+                parent not found, rate limited, etc.).
+        """
+        metadata: dict[str, Any] = {
+            "name": name,
+            "mimeType": mime_type,
+        }
+        if parent_id:
+            metadata["parents"] = [parent_id]
+
+        if mime_type == GOOGLE_DRIVE_FOLDER_MIME or not initial_content:
+            response = await self._client.post(
+                "/",
+                json=metadata,
+                params={"fields": SUPPORTED_FIELDS_STR, "supportsAllDrives": "true"},
+            )
+        else:
+            response = await self._create_file_with_content(
+                metadata=metadata,
+                content=initial_content,
+                target_mime_type=mime_type,
+            )
+
+        if response.status_code == 404:
+            raise GoogleDriveError(
+                f"Parent folder with ID '{parent_id}' not found."
+                if parent_id
+                else "Resource not found."
+            )
+        if response.status_code == 403:
+            raise GoogleDriveError(
+                "Permission denied: you don't have permission to create files in this location."
+            )
+        if response.status_code == 400:
+            raise GoogleDriveError(
+                f"Bad request: invalid parameters for file creation. "
+                f"Check that the MIME type '{mime_type}' is valid."
+            )
+        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 _create_file_with_content(
+        self,
+        metadata: dict[str, Any],
+        content: str,
+        target_mime_type: str,
+    ) -> httpx.Response:
+        """Create a file with content using multipart upload.
+
+        Args:
+            metadata: File metadata dictionary.
+            content: Text content for the file.
+            target_mime_type: The target MIME type for the file.
+
+        Returns
+        -------
+            The HTTP response from the upload.
+        """
+        content_type = UPLOAD_CONTENT_TYPES.get(target_mime_type, "text/plain")
+        boundary = f"===gdrive_boundary_{uuid.uuid4().hex}==="
+        body_parts = [
+            f"--{boundary}",
+            "Content-Type: application/json; charset=UTF-8",
+            "",
+            json.dumps(metadata),
+            f"--{boundary}",
+            f"Content-Type: {content_type}",
+            "",
+            content,
+            f"--{boundary}--",
+        ]
+        body = "\r\n".join(body_parts)
+
+        upload_url = "https://www.googleapis.com/upload/drive/v3/files"
+        return await self._client.post(
+            upload_url,
+            content=body.encode("utf-8"),
+            params={
+                "uploadType": "multipart",
+                "fields": SUPPORTED_FIELDS_STR,
+                "supportsAllDrives": "true",
+            },
+            headers={"Content-Type": f"multipart/related; boundary={boundary}"},
+        )
+
+    async def manage_access(
+        self,
+        *,
+        file_id: str,
+        action: Literal["add", "update", "remove"],
+        role: Literal["reader", "commenter", "writer", "fileOrganizer", "organizer", "owner"]
+        | None = None,
+        email_address: str | None = None,
+        permission_id: str | None = None,
+        transfer_ownership: bool = False,
+    ) -> str:
+        """Manage access permissions for a Google Drive file or folder.
+
+        Adds, updates, or removes sharing permissions on an existing Google Drive
+        file or folder using the Google Drive Permissions API.
+
+        This method supports granting access to users or groups, changing access
+        roles, and revoking permissions. Ownership transfer is supported for files
+        in "My Drive" when explicitly requested.
+
+        Args:
+            file_id: The ID of the Google Drive file or folder whose permissions
+                are being managed.
+            action: The permission operation to perform.
+            role: The access role to assign or update. Valid values include
+                Required for "add" and "update" actions.
+            email_address: The email address of the user or group to grant access to.
+                Required for the "add" action.
+            permission_id: The ID of the permission to update or remove.
+                Required for "update" and "remove" actions.
+            transfer_ownership: Whether to transfer ownership of the file.
+                Only applicable when action="update" and role="owner".
+
+        Returns
+        -------
+            Permission id.
+            For "add" its newly added permission.
+            For "update"/"remove" its previous permission.
+
+        Raises
+        ------
+            GoogleDriveError: If the permission operation fails (invalid arguments,
+                insufficient permissions, resource not found, ownership transfer
+                not allowed, rate limited, etc.).
+        """
+        if not file_id.strip():
+            raise GoogleDriveError("Argument validation error: 'file_id' cannot be empty.")
+
+        if action == "add" and not email_address:
+            raise GoogleDriveError("'email_address' is required for action 'add'.")
+
+        if action in ("update", "remove") and not permission_id:
+            raise GoogleDriveError("'permission_id' is required for action 'update' or 'remove'.")
+
+        if action != "remove" and not role:
+            raise GoogleDriveError("'role' is required for action 'add' or 'update'.")
+
+        if action == "add":
+            response = await self._client.post(
+                url=f"/{file_id}/permissions",
+                json={
+                    "type": "user",
+                    "role": role,
+                    "emailAddress": email_address,
+                },
+                params={"sendNotificationEmail": False, "supportsAllDrives": True},
+            )
+
+        elif action == "update":
+            response = await self._client.patch(
+                url=f"/{file_id}/permissions/{permission_id}",
+                json={"role": role},
+                params={"transferOwnership": transfer_ownership, "supportsAllDrives": True},
+            )
+
+        elif action == "remove":
+            response = await self._client.delete(url=f"/{file_id}/permissions/{permission_id}")
+
+        else:
+            raise GoogleDriveError(f"Invalid action '{action}'")
+
+        if response.status_code not in (200, 201, 204):
+            raise GoogleDriveError(f"Drive API error {response.status_code}: {response.text}")
+
+        if action == "add":
+            return response.json()["id"]
+
+        # Cannot be null here because of above validators
+        return permission_id  # type: ignore
+
     async def __aenter__(self) -> "GoogleDriveClient":
         """Async context manager entry."""
         return self