datarobot-genai 0.2.31__tar.gz → 0.2.41__tar.gz

{datarobot_genai-0.2.31 → datarobot_genai-0.2.41}/PKG-INFO

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

{datarobot_genai-0.2.31 → datarobot_genai-0.2.41}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.31"
+version = "0.2.41"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
@@ -130,6 +130,7 @@ dev = [
   "pytest-asyncio>=0.23.5",
   "responses>=0.25.8",
   "ipdb>=0.13.13",
+  "aioresponses>=0.7.8",
 ]
 
 

{datarobot_genai-0.2.31 → datarobot_genai-0.2.41}/src/datarobot_genai/core/cli/agent_kernel.py

@@ -107,7 +107,10 @@ class AgentKernel:
         stream: bool = False,
         config: Any | None = None,
     ) -> ChatCompletion | Stream[ChatCompletionChunk]:
-        chat_api_url = config.agent_endpoint if config else self.base_url
+        if config is not None:
+            chat_api_url = f"http://localhost:{config.local_dev_port}"
+        else:
+            chat_api_url = self.base_url
         print(chat_api_url)
 
         return self._do_chat_completion(chat_api_url, user_prompt, completion_json, stream=stream)

{datarobot_genai-0.2.31 → datarobot_genai-0.2.41}/src/datarobot_genai/core/utils/auth.py

@@ -12,13 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
+import os
 import warnings
+from collections.abc import Sequence
 from typing import Any
+from typing import Protocol
 
+import aiohttp
 import jwt
-from datarobot.auth.datarobot.oauth import AsyncOAuth as DatarobotAsyncOAuthClient
+from datarobot.auth.datarobot.oauth import AsyncOAuth as DatarobotOAuthClient
+from datarobot.auth.exceptions import OAuthProviderNotFound
+from datarobot.auth.exceptions import OAuthValidationErr
 from datarobot.auth.identity import Identity
-from datarobot.auth.oauth import AsyncOAuthComponent
+from datarobot.auth.oauth import OAuthToken
 from datarobot.auth.session import AuthCtx
 from datarobot.core.config import DataRobotAppFrameworkBaseSettings
 from datarobot.models.genai.agent.auth import ToolAuth
@@ -183,52 +189,203 @@ class AuthContextHeaderHandler:
             return None
 
 
+# --- OAuth Token Provider Implementation ---
+
+
+class TokenRetriever(Protocol):
+    """Protocol for OAuth token retrievers."""
+
+    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
+        """Filter identities to only those valid for this retriever implementation."""
+        ...
+
+    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
+        """Refresh the access token for the given identity ID.
+
+        Parameters
+        ----------
+        identity_id : str
+            The provider identity ID to refresh the token for.
+
+        Returns
+        -------
+        OAuthToken
+            The refreshed OAuth token.
+        """
+        ...
+
+
+class DatarobotTokenRetriever:
+    """Retrieves OAuth tokens using the DataRobot platform."""
+
+    def __init__(self) -> None:
+        self._client = DatarobotOAuthClient()
+
+    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
+        """Filter oauth2 identities to only those with provider_identity_id.
+
+        The `provider_identity_id` is required in order to identify the provider
+        and retrieve the access token from DataRobot OAuth Providers service.
+        """
+        return [i for i in identities if i.type == "oauth2" and i.provider_identity_id]
+
+    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
+        """Refresh the access token using DataRobot's OAuth client."""
+        return await self._client.refresh_access_token(identity_id=identity.provider_identity_id)
+
+
+class AuthlibTokenRetriever:
+    """Retrieves OAuth tokens from a generic Authlib-based endpoint."""
+
+    def __init__(self, application_endpoint: str) -> None:
+        if not application_endpoint:
+            raise ValueError("AuthlibTokenRetriever requires 'application_endpoint'.")
+        self.application_endpoint = application_endpoint.rstrip("/")
+
+    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
+        """Filter identities to only OAuth2 identities."""
+        return [i for i in identities if i.type == "oauth2" and i.provider_identity_id is None]
+
+    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
+        """Retrieve an OAuth token via an HTTP POST request.
+
+        Parameters
+        ----------
+        identity : Identity
+            The identity to retrieve the token for.
+
+        Returns
+        -------
+        OAuthToken
+            The retrieved OAuth token.
+        """
+        api_token = os.environ.get("DATAROBOT_API_TOKEN")
+        if not api_token:
+            raise ValueError("DATAROBOT_API_TOKEN environment variable is required but not set.")
+
+        token_url = f"{self.application_endpoint}/oauth/token/"
+        headers = {"Authorization": f"Bearer {api_token}"}
+        payload = {"identity_id": identity.id}
+        timeout = aiohttp.ClientTimeout(total=30)
+
+        try:
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+                async with session.post(token_url, headers=headers, json=payload) as response:
+                    response.raise_for_status()
+                    data = await response.json()
+                    logger.debug(f"Retrieved access token from {token_url}")
+                    return OAuthToken(**data)
+        except aiohttp.ClientError as e:
+            logger.error(f"Error retrieving token from {token_url}: {e}")
+            raise
+
+
+class OAuthConfig(BaseModel):
+    """Configuration extracted from AuthCtx metadata for OAuth operations."""
+
+    implementation: str = "datarobot"
+    application_endpoint: str | None = None
+
+    @classmethod
+    def from_auth_ctx(cls, auth_ctx: AuthCtx) -> "OAuthConfig":
+        metadata = auth_ctx.metadata or {}
+        return cls(
+            implementation=metadata.get("oauth_implementation", "datarobot"),
+            application_endpoint=metadata.get("application_endpoint"),
+        )
+
+
+def create_token_retriever(config: OAuthConfig) -> TokenRetriever:
+    """Create a token retriever based on the OAuth configuration.
+
+    Parameters
+    ----------
+    config : OAuthConfig
+        The OAuth configuration specifying implementation type and endpoints.
+
+    Returns
+    -------
+    TokenRetriever
+        The configured token retriever instance.
+    """
+    if config.implementation == "datarobot":
+        return DatarobotTokenRetriever()
+
+    if config.implementation == "authlib":
+        if not config.application_endpoint:
+            raise ValueError("Required 'application_endpoint' not found in metadata.")
+        return AuthlibTokenRetriever(config.application_endpoint)
+
+    raise ValueError(
+        f"Unsupported OAuth implementation: '{config.implementation}'. "
+        f"Supported values: datarobot, authlib."
+    )
+
+
 class AsyncOAuthTokenProvider:
-    """Manages OAuth access tokens using generic OAuth client."""
+    """Provides OAuth tokens for authorized users.
+
+    This class manages OAuth token retrieval for users with multiple identity providers.
+    It uses either DataRobot or Authlib as the OAuth token storage and refresh backend
+    based on the auth context metadata.
+    """
 
     def __init__(self, auth_ctx: AuthCtx) -> None:
+        """Initialize the provider with an authorization context.
+
+        Parameters
+        ----------
+        auth_ctx : AuthCtx
+            The authorization context containing user identities and metadata.
+        """
         self.auth_ctx = auth_ctx
-        self.oauth_client = self._create_oauth_client()
+        config = OAuthConfig.from_auth_ctx(auth_ctx)
+        self._retriever = create_token_retriever(config)
 
     def _get_identity(self, provider_type: str | None) -> Identity:
-        """Retrieve the appropriate identity from the authentication context."""
-        identities = [x for x in self.auth_ctx.identities if x.provider_identity_id is not None]
-
-        if not identities:
-            raise ValueError("No identities found in authorization context.")
+        """Get identity from auth context, filtered by provider_type if specified."""
+        oauth_identities = self._retriever.filter_identities(self.auth_ctx.identities)
+        if not oauth_identities:
+            raise OAuthProviderNotFound("No OAuth provider found.")
 
         if provider_type is None:
-            if len(identities) > 1:
-                raise ValueError(
-                    "Multiple identities found. Please specify 'provider_type' parameter."
+            if len(oauth_identities) > 1:
+                raise OAuthValidationErr(
+                    "Multiple OAuth providers found. Specify 'provider_type' parameter."
                 )
-            return identities[0]
-
-        identity = next((id for id in identities if id.provider_type == provider_type), None)
+            return oauth_identities[0]
 
+        identity = next((i for i in oauth_identities if i.provider_type == provider_type), None)
         if identity is None:
-            raise ValueError(f"No identity found for provider '{provider_type}'.")
-
+            raise OAuthValidationErr(f"No identity found for provider '{provider_type}'.")
         return identity
 
     async def get_token(self, auth_type: ToolAuth, provider_type: str | None = None) -> str:
-        """Get OAuth access token using the specified method."""
+        """Get an OAuth access token for the specified auth type and provider.
+
+        Parameters
+        ----------
+        auth_type : ToolAuth
+            Authentication type (only OBO is supported).
+        provider_type : str, optional
+            The specific provider to use (e.g., 'google'). Required if multiple
+            identities are available.
+
+        Returns
+        -------
+        str
+            The retrieved OAuth access token.
+
+        Raises
+        ------
+        ValueError
+            If the auth type is unsupported or if a suitable identity cannot be found.
+        """
         if auth_type != ToolAuth.OBO:
             raise ValueError(
-                f"Unsupported auth type: {auth_type}. Only {ToolAuth.OBO} is supported."
+                f"Unsupported auth type: {auth_type}. Only OBO (on-behalf-of) is supported."
             )
 
         identity = self._get_identity(provider_type)
-        token_data = await self.oauth_client.refresh_access_token(
-            identity_id=identity.provider_identity_id
-        )
+        token_data = await self._retriever.refresh_access_token(identity)
         return token_data.access_token
-
-    def _create_oauth_client(self) -> AsyncOAuthComponent:
-        """Create either DataRobot or Authlib OAuth client based on
-        authorization context.
-
-        Note: at the moment, only DataRobot OAuth client is supported.
-        """
-        logger.debug("Using DataRobot OAuth client")
-        return DatarobotAsyncOAuthClient()

{datarobot_genai-0.2.31 → datarobot_genai-0.2.41}/src/datarobot_genai/drmcp/core/exceptions.py

@@ -19,7 +19,3 @@ class DynamicToolRegistrationError(Exception):
 
 class DynamicPromptRegistrationError(Exception):
     """Exception raised for errors in the dynamic prompt registration process."""
-
-
-class MCPError(Exception):
-    """Base class for MCP errors."""

{datarobot_genai-0.2.31 → datarobot_genai-0.2.41}/src/datarobot_genai/drmcp/core/logging.py

@@ -20,7 +20,7 @@ from collections.abc import Callable
 from typing import Any
 from typing import TypeVar
 
-from .exceptions import MCPError
+from fastmcp.exceptions import ToolError
 
 # Secret patterns to redact from logs
 SECRET_PATTERNS = [
@@ -93,6 +93,6 @@ def log_execution(func: F) -> F:
             return result
         except Exception as e:
             error_msg = _log_error(logger, func.__name__, e, args=args, kwargs=kwargs)
-            raise MCPError(error_msg)
+            raise ToolError(error_msg)
 
     return wrapper  # type: ignore[return-value]

{datarobot_genai-0.2.31 → datarobot_genai-0.2.41}/src/datarobot_genai/drmcp/test_utils/clients/base.py

@@ -27,7 +27,7 @@ from mcp.types import ListToolsResult
 from mcp.types import TextContent
 from openai.types.chat.chat_completion import ChatCompletion
 
-from ..utils import save_response_to_file
+from datarobot_genai.drmcp.test_utils.utils import save_response_to_file
 
 
 class ToolCall:

{datarobot_genai-0.2.31 → datarobot_genai-0.2.41}/src/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-0.2.31 → datarobot_genai-0.2.41}/src/datarobot_genai/drmcp/tools/clients/gdrive.py

@@ -20,6 +20,7 @@ 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
@@ -821,6 +822,96 @@ class GoogleDriveClient:
             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

{datarobot_genai-0.2.31 → datarobot_genai-0.2.41}/src/datarobot_genai/drmcp/tools/clients/microsoft_graph.py

@@ -16,6 +16,8 @@
 
 import logging
 from typing import Any
+from typing import Literal
+from urllib.parse import quote
 
 import httpx
 from datarobot.auth.datarobot.exceptions import OAuthServiceClientErr
@@ -415,6 +417,175 @@ class MicrosoftGraphClient:
 
         return base_resource
 
+    async def get_personal_drive_id(self) -> str:
+        """Get the current user's personal OneDrive drive ID.
+
+        Returns
+        -------
+            The drive ID string for the user's personal OneDrive.
+
+        Raises
+        ------
+            MicrosoftGraphError: If the drive cannot be retrieved.
+        """
+        try:
+            response = await self._client.get(f"{GRAPH_API_BASE}/me/drive")
+            response.raise_for_status()
+            data = response.json()
+            return data["id"]
+        except httpx.HTTPStatusError as e:
+            status_code = e.response.status_code
+            if status_code == 401:
+                raise MicrosoftGraphError(
+                    "Authentication failed. Access token may be expired or invalid."
+                ) from e
+            if status_code == 403:
+                raise MicrosoftGraphError(
+                    "Permission denied: cannot access personal OneDrive. "
+                    "Requires Files.Read or Files.ReadWrite permission."
+                ) from e
+            raise MicrosoftGraphError(f"Failed to get personal OneDrive: HTTP {status_code}") from e
+
+    async def create_file(
+        self,
+        drive_id: str,
+        file_name: str,
+        content: str,
+        parent_folder_id: str = "root",
+        conflict_behavior: str = "rename",
+    ) -> MicrosoftGraphItem:
+        """Create a text file in a drive (SharePoint document library or OneDrive).
+
+        Uses Microsoft Graph's simple upload endpoint for files < 4MB.
+        Files are created as text/plain content.
+
+        Args:
+            drive_id: The ID of the drive (document library) where the file will be created.
+            file_name: The name of the file to create (e.g., 'report.txt').
+            content: The text content to store in the file.
+            parent_folder_id: ID of the parent folder. Defaults to "root" (drive root folder).
+            conflict_behavior: How to handle name conflicts. Options:
+                - "rename" (default): Auto-renames to 'filename (1).txt', etc.
+                - "fail": Returns 409 Conflict error
+                - "replace": Overwrites existing file
+
+        Returns
+        -------
+            MicrosoftGraphItem representing the created file.
+
+        Raises
+        ------
+            MicrosoftGraphError: If file creation fails.
+        """
+        if not drive_id or not drive_id.strip():
+            raise MicrosoftGraphError("drive_id cannot be empty")
+        if not file_name or not file_name.strip():
+            raise MicrosoftGraphError("file_name cannot be empty")
+
+        # URL encode the filename for path-based addressing
+        encoded_name = quote(file_name, safe="")
+
+        # Simple upload endpoint for files < 4MB
+        # Reference: https://learn.microsoft.com/en-us/graph/api/driveitem-put-content
+        upload_url = (
+            f"{GRAPH_API_BASE}/drives/{drive_id}/items/{parent_folder_id}:/{encoded_name}:/content"
+        )
+
+        try:
+            response = await self._client.put(
+                upload_url,
+                content=content.encode("utf-8"),
+                headers={"Content-Type": "text/plain"},
+                params={"@microsoft.graph.conflictBehavior": conflict_behavior},
+            )
+            response.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            raise self._handle_create_file_error(e, drive_id, file_name, parent_folder_id) from e
+
+        return MicrosoftGraphItem.from_api_response(response.json())
+
+    def _handle_create_file_error(
+        self,
+        error: httpx.HTTPStatusError,
+        drive_id: str,
+        file_name: str,
+        parent_folder_id: str,
+    ) -> MicrosoftGraphError:
+        """Handle HTTP errors for file creation and return appropriate MicrosoftGraphError."""
+        status_code = error.response.status_code
+        error_msg = f"Failed to create file: HTTP {status_code}"
+
+        if status_code == 400:
+            try:
+                error_data = error.response.json()
+                api_message = error_data.get("error", {}).get("message", "Invalid request")
+                error_msg = f"Bad request creating file: {api_message}"
+            except Exception:
+                error_msg = "Bad request: invalid parameters for file creation."
+        elif status_code == 401:
+            error_msg = "Authentication failed. Access token may be expired or invalid."
+        elif status_code == 403:
+            error_msg = (
+                f"Permission denied: you don't have permission to create files in drive "
+                f"'{drive_id}'. Requires Files.ReadWrite.All permission."
+            )
+        elif status_code == 404:
+            error_msg = (
+                f"Parent folder '{parent_folder_id}' not found in drive '{drive_id}'."
+                if parent_folder_id != "root"
+                else f"Drive '{drive_id}' not found."
+            )
+        elif status_code == 409:
+            error_msg = f"File '{file_name}' already exists and conflict behavior is set to 'fail'."
+        elif status_code == 429:
+            error_msg = "Rate limit exceeded. Please try again later."
+
+        return MicrosoftGraphError(error_msg)
+
+    async def share_item(
+        self,
+        file_id: str,
+        document_library_id: str,
+        recipient_emails: list[str],
+        role: Literal["read", "write"],
+        send_invitation: bool,
+    ) -> None:
+        """
+        Share sharepoint / ondrive item using Microsoft Graph API.
+        Under the hood all resources in sharepoint/onedrive
+            in MS Graph API are treated as 'driveItem'.
+
+        Args:
+            file_id: The ID of the file or folder to share.
+            document_library_id: The ID of the document library containing the item.
+            recipient_emails: A list of email addresses to invite.
+            role: The role to assign.
+            send_invitation: Flag determining if recipients should be notified
+
+        Returns
+        -------
+            None
+
+        Raises
+        ------
+            MicrosoftGraphError: If sharing fails
+        """
+        graph_url = f"{GRAPH_API_BASE}/drives/{document_library_id}/items/{file_id}/invite"
+
+        payload = {
+            "recipients": [{"email": email} for email in recipient_emails],
+            "requireSignIn": True,
+            "sendInvitation": send_invitation,
+            "roles": [role],
+        }
+
+        response = await self._client.post(url=graph_url, json=payload)
+
+        if response.status_code not in (200, 201):
+            raise MicrosoftGraphError(
+                f"Microsoft Graph API error {response.status_code}: {response.text}"
+            )
+
     async def __aenter__(self) -> "MicrosoftGraphClient":
         """Async context manager entry."""
         return self