datarobot-genai 0.1.59__tar.gz → 0.1.64__tar.gz

{datarobot_genai-0.1.59 → datarobot_genai-0.1.64}/PKG-INFO

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

{datarobot_genai-0.1.59 → datarobot_genai-0.1.64}/pyproject.toml

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

datarobot_genai-0.1.64/src/datarobot_genai/core/mcp/common.py

@@ -0,0 +1,161 @@
+# Copyright 2025 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import re
+from typing import Any
+from typing import Literal
+
+from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from pydantic import field_validator
+
+from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
+
+
+class MCPConfig(DataRobotAppFrameworkBaseSettings):
+    """Configuration for MCP server connection.
+
+    Derived values are exposed as properties rather than stored, avoiding
+    Pydantic field validation/serialization concerns for internal helpers.
+    """
+
+    external_mcp_url: str | None = None
+    external_mcp_headers: str | None = None
+    external_mcp_transport: Literal["sse", "streamable-http"] = "streamable-http"
+    mcp_deployment_id: str | None = None
+    datarobot_endpoint: str | None = None
+    datarobot_api_token: str | None = None
+    authorization_context: dict[str, Any] | None = None
+
+    _auth_context_handler: AuthContextHeaderHandler | None = None
+    _server_config: dict[str, Any] | None = None
+
+    @field_validator("external_mcp_headers", mode="before")
+    @classmethod
+    def validate_external_mcp_headers(cls, value: str | None) -> str | None:
+        if value is None:
+            return None
+
+        if not isinstance(value, str):
+            msg = "external_mcp_headers must be a JSON string"
+            raise TypeError(msg)
+
+        candidate = value.strip()
+
+        try:
+            json.loads(candidate)
+        except json.JSONDecodeError as exc:
+            msg = "external_mcp_headers must be valid JSON"
+            raise ValueError(msg) from exc
+
+        return candidate
+
+    @field_validator("mcp_deployment_id", mode="before")
+    @classmethod
+    def validate_mcp_deployment_id(cls, value: str | None) -> str | None:
+        if value is None:
+            return None
+
+        if not isinstance(value, str):
+            msg = "mcp_deployment_id must be a string"
+            raise TypeError(msg)
+
+        candidate = value.strip()
+
+        if not re.fullmatch(r"[0-9a-fA-F]{24}", candidate):
+            msg = "mcp_deployment_id must be a valid 24-character hex ID"
+            raise ValueError(msg)
+
+        return candidate
+
+    def _authorization_bearer_header(self) -> dict[str, str]:
+        """Return Authorization header with Bearer token or empty dict."""
+        if not self.datarobot_api_token:
+            return {}
+        auth = (
+            self.datarobot_api_token
+            if self.datarobot_api_token.startswith("Bearer ")
+            else f"Bearer {self.datarobot_api_token}"
+        )
+        return {"Authorization": auth}
+
+    @property
+    def auth_context_handler(self) -> AuthContextHeaderHandler:
+        if self._auth_context_handler is None:
+            self._auth_context_handler = AuthContextHeaderHandler()
+        return self._auth_context_handler
+
+    @property
+    def server_config(self) -> dict[str, Any] | None:
+        if self._server_config is None:
+            self._server_config = self._build_server_config()
+        return self._server_config
+
+    def _authorization_context_header(self) -> dict[str, str]:
+        """Return X-DataRobot-Authorization-Context header or empty dict."""
+        try:
+            return self.auth_context_handler.get_header(self.authorization_context)
+        except (LookupError, RuntimeError):
+            # Authorization context not available (e.g., in tests)
+            return {}
+
+    def _build_server_config(self) -> dict[str, Any] | None:
+        """
+        Get MCP server configuration.
+
+        Returns
+        -------
+            Server configuration dict with url, transport, and optional headers,
+            or None if not configured.
+        """
+        if self.external_mcp_url:
+            # External MCP URL - no authentication needed
+            if self.external_mcp_headers:
+                headers = json.loads(self.external_mcp_headers)
+            else:
+                headers = {}
+
+            config = {
+                "url": self.external_mcp_url.rstrip("/"),
+                "transport": self.external_mcp_transport,
+                "headers": headers,
+            }
+            return config
+        elif self.mcp_deployment_id:
+            # DataRobot deployment ID - requires authentication
+            if self.datarobot_endpoint is None:
+                raise ValueError(
+                    "When using a DataRobot hosted MCP deployment, datarobot_endpoint must be set."
+                )
+            if self.datarobot_api_token is None:
+                raise ValueError(
+                    "When using a DataRobot hosted MCP deployment, datarobot_api_token must be set."
+                )
+            base_url = self.datarobot_endpoint.rstrip("/")
+            if not base_url.endswith("/api/v2"):
+                base_url = base_url + "/api/v2"
+            url = f"{base_url}/deployments/{self.mcp_deployment_id}/directAccess/mcp"
+
+            headers = {
+                **self._authorization_bearer_header(),
+                **self._authorization_context_header(),
+            }
+
+            return {
+                "url": url,
+                "transport": "streamable-http",
+                "headers": headers,
+            }
+
+        return None

{datarobot_genai-0.1.59 → datarobot_genai-0.1.64}/src/datarobot_genai/core/utils/auth.py

@@ -16,9 +16,14 @@ import warnings
 from typing import Any
 
 import jwt
+from datarobot.auth.datarobot.oauth import AsyncOAuth as DatarobotAsyncOAuthClient
+from datarobot.auth.identity import Identity
+from datarobot.auth.oauth import AsyncOAuthComponent
 from datarobot.auth.session import AuthCtx
 from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from datarobot.models.genai.agent.auth import ToolAuth
 from datarobot.models.genai.agent.auth import get_authorization_context
+from pydantic import BaseModel
 
 logger = logging.getLogger(__name__)
 
@@ -27,6 +32,13 @@ class AuthContextConfig(DataRobotAppFrameworkBaseSettings):
     session_secret_key: str = ""
 
 
+class DRAppCtx(BaseModel):
+    """DataRobot application context from authorization metadata."""
+
+    email: str | None = None
+    api_key: str | None = None
+
+
 class AuthContextHeaderHandler:
     """Manages encoding and decoding of authorization context into JWT tokens.
 
@@ -146,6 +158,7 @@ class AuthContextHeaderHandler:
 
         auth_ctx_dict = self.decode(token)
         if not auth_ctx_dict:
+            logger.debug("Failed to decode auth context from token")
             return None
 
         try:
@@ -153,3 +166,54 @@ class AuthContextHeaderHandler:
         except Exception as e:
             logger.error(f"Failed to create AuthCtx from decoded token: {e}", exc_info=True)
             return None
+
+
+class AsyncOAuthTokenProvider:
+    """Manages OAuth access tokens using generic OAuth client."""
+
+    def __init__(self, auth_ctx: AuthCtx) -> None:
+        self.auth_ctx = auth_ctx
+        self.oauth_client = self._create_oauth_client()
+
+    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.")
+
+        if provider_type is None:
+            if len(identities) > 1:
+                raise ValueError(
+                    "Multiple identities found. Please specify 'provider_type' parameter."
+                )
+            return identities[0]
+
+        identity = next((id for id in identities if id.provider_type == provider_type), None)
+
+        if identity is None:
+            raise ValueError(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."""
+        if auth_type != ToolAuth.OBO:
+            raise ValueError(
+                f"Unsupported auth type: {auth_type}. Only {ToolAuth.OBO} is supported."
+            )
+
+        identity = self._get_identity(provider_type)
+        token_data = await self.oauth_client.refresh_access_token(
+            identity_id=identity.provider_identity_id
+        )
+        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.1.59 → datarobot_genai-0.1.64}/src/datarobot_genai/crewai/base.py

@@ -92,8 +92,6 @@ class CrewAIAgent(BaseAgent[BaseTool], abc.ABC):
 
         # Use MCP context manager to handle connection lifecycle
         with mcp_tools_context(
-            api_base=self.api_base,
-            api_key=self.api_key,
             authorization_context=self._authorization_context,
         ) as mcp_tools:
             # Set MCP tools for all agents if MCP is not configured this is effectively a no-op

{datarobot_genai-0.1.59 → datarobot_genai-0.1.64}/src/datarobot_genai/crewai/mcp.py

@@ -29,15 +29,10 @@ from datarobot_genai.core.mcp.common import MCPConfig
 
 @contextmanager
 def mcp_tools_context(
-    api_base: str | None = None,
-    api_key: str | None = None,
     authorization_context: dict[str, Any] | None = None,
 ) -> Generator[list[Any], None, None]:
     """Context manager for MCP tools that handles connection lifecycle."""
-    config = MCPConfig(
-        api_base=api_base, api_key=api_key, authorization_context=authorization_context
-    )
-
+    config = MCPConfig(authorization_context=authorization_context)
     # If no MCP server configured, return empty tools list
     if not config.server_config:
         print("No MCP server configured, using empty tools list", flush=True)
@@ -47,10 +42,8 @@ def mcp_tools_context(
     print(f"Connecting to MCP server: {config.server_config['url']}", flush=True)
 
     # Use MCPServerAdapter as context manager with the server config
-    adapter_setting = config.server_config.copy()
-    adapter_setting["transport"] = "streamable-http"
     try:
-        with MCPServerAdapter(adapter_setting) as tools:
+        with MCPServerAdapter(config.server_config) as tools:
             print(
                 f"Successfully connected to MCP server, got {len(tools)} tools",
                 flush=True,

{datarobot_genai-0.1.59 → datarobot_genai-0.1.64}/src/datarobot_genai/drmcp/core/auth.py

@@ -18,7 +18,6 @@ import logging
 from typing import Any
 
 from datarobot.auth.session import AuthCtx
-from datarobot.models.genai.agent.auth import OAuthAccessTokenProvider
 from datarobot.models.genai.agent.auth import ToolAuth
 from fastmcp.server.dependencies import get_context
 from fastmcp.server.dependencies import get_http_headers
@@ -27,12 +26,15 @@ from fastmcp.server.middleware import Middleware
 from fastmcp.server.middleware import MiddlewareContext
 from fastmcp.tools.tool import ToolResult
 
+from datarobot_genai.core.utils.auth import AsyncOAuthTokenProvider
 from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
-from datarobot_genai.drmcp import get_config
 
 logger = logging.getLogger(__name__)
 
 
+AUTH_CTX_KEY = "authorization_context"
+
+
 class OAuthMiddleWare(Middleware):
     """Middleware that parses `x-datarobot-authorization-context` for tool calls.
 
@@ -45,16 +47,8 @@ class OAuthMiddleWare(Middleware):
         Handler for encoding/decoding JWT tokens containing auth context.
     """
 
-    def __init__(self, secret_key: str | None = None) -> None:
-        """Initialize the middleware with authentication handler.
-
-        Parameters
-        ----------
-        secret_key : Optional[str]
-            Secret key for JWT validation. If None, uses the value from config.
-        """
-        secret_key = secret_key or get_config().session_secret_key
-        self.auth_handler = AuthContextHeaderHandler(secret_key)
+    def __init__(self, auth_handler: AuthContextHeaderHandler | None = None) -> None:
+        self.auth_handler = auth_handler or AuthContextHeaderHandler()
 
     async def on_call_tool(
         self, context: MiddlewareContext, call_next: CallNext[Any, ToolResult]
@@ -74,9 +68,12 @@ class OAuthMiddleWare(Middleware):
             The result from the tool execution.
         """
         auth_context = self._extract_auth_context()
+        if not auth_context:
+            logger.debug("No valid authorization context extracted from request headers.")
 
         if context.fastmcp_context is not None:
-            context.fastmcp_context.auth_context = auth_context
+            context.fastmcp_context.set_state(AUTH_CTX_KEY, auth_context)
+            logger.debug("Authorization context attached to state.")
 
         return await call_next(context)
 
@@ -99,8 +96,8 @@ class OAuthMiddleWare(Middleware):
             return None
 
 
-async def get_auth_context() -> AuthCtx:
-    """Retrieve the AuthCtx from the current request context, if available.
+async def must_get_auth_context() -> AuthCtx:
+    """Retrieve the AuthCtx from the current request context or raise error.
 
     Raises
     ------
@@ -113,14 +110,15 @@ async def get_auth_context() -> AuthCtx:
         The authorization context associated with the current request.
     """
     context = get_context()
-    auth_ctx = getattr(context, "auth_context", None)
+
+    auth_ctx = context.get_state(AUTH_CTX_KEY)
     if not auth_ctx:
-        raise RuntimeError("No authorization context found.")
+        raise RuntimeError("Could not retrieve authorization context from FastMCP context state.")
 
     return auth_ctx
 
 
-async def get_access_token(provider: str | None = None) -> str:
+async def get_access_token(provider_type: str | None = None) -> str:
     """Retrieve access token from the DataRobot OAuth Provider Service.
 
     OAuth access tokens can be retrieved only for providers where the user completed
@@ -132,7 +130,7 @@ async def get_access_token(provider: str | None = None) -> str:
 
     Parameters
     ----------
-    provider : str, optional
+    provider_type : str, optional
         The name of the OAuth provider. It should match the name of the provider configured
         during provider setup. If no value is provided and only one OAuth provider exists, that
         provider will be used. If multiple providers exist and none is specified, an error will be
@@ -142,12 +140,18 @@ async def get_access_token(provider: str | None = None) -> str:
     -------
     The oauth access token.
     """
-    token_provider = OAuthAccessTokenProvider(await get_auth_context())
-    access_token = token_provider.get_token(ToolAuth.OBO, provider)
-    return access_token
+    auth_ctx = await must_get_auth_context()
+    logger.debug("Retrieved authorization context")
+
+    oauth_token_provider = AsyncOAuthTokenProvider(auth_ctx)
+    oauth_access_token = await oauth_token_provider.get_token(
+        auth_type=ToolAuth.OBO,
+        provider_type=provider_type,
+    )
+    return oauth_access_token
 
 
-def initialize_oauth_middleware(mcp: Any, secret_key: str | None = None) -> None:
+def initialize_oauth_middleware(mcp: Any) -> None:
     """Initialize and register OAuth middleware with the MCP server.
 
     Parameters
@@ -157,6 +161,5 @@ def initialize_oauth_middleware(mcp: Any, secret_key: str | None = None) -> None
     secret_key : Optional[str]
         Secret key for JWT validation. If None, uses the value from config.
     """
-    middleware = OAuthMiddleWare(secret_key=secret_key)
-    mcp.add_middleware(middleware)
+    mcp.add_middleware(OAuthMiddleWare())
     logger.info("OAuth middleware registered successfully")

{datarobot_genai-0.1.59 → datarobot_genai-0.1.64}/src/datarobot_genai/drmcp/core/clients.py

@@ -21,6 +21,9 @@ from datarobot.context import Context as DRContext
 from datarobot.rest import RESTClientObject
 from fastmcp.server.dependencies import get_http_headers
 
+from datarobot_genai.core.utils.auth import AuthContextHeaderHandler
+from datarobot_genai.core.utils.auth import DRAppCtx
+
 from .credentials import get_credentials
 
 logger = logging.getLogger(__name__)
@@ -66,19 +69,80 @@ def _extract_token_from_headers(headers: dict[str, str]) -> str | None:
     return None
 
 
+def _extract_token_from_auth_context(headers: dict[str, str]) -> str | None:
+    """
+    Extract API token from authorization context metadata as a fallback.
+
+    Args:
+        headers: Dictionary of headers (keys should be lowercase)
+
+    Returns
+    -------
+        The extracted API key from auth context metadata, or None if not found
+    """
+    try:
+        auth_handler = AuthContextHeaderHandler()
+
+        auth_ctx = auth_handler.get_context(headers)
+        if not auth_ctx or not auth_ctx.metadata:
+            return None
+
+        metadata = auth_ctx.metadata
+        if not isinstance(metadata, dict):
+            return None
+
+        dr_ctx: DRAppCtx = DRAppCtx(**metadata.get("dr_ctx", {}))
+        if dr_ctx.api_key:
+            logger.debug("Extracted token from auth context")
+            return dr_ctx.api_key
+
+        return None
+
+    except Exception as e:
+        logger.debug(f"Failed to get token from auth context: {e}")
+        return None
+
+
+def extract_token_from_headers(headers: dict[str, str]) -> str | None:
+    """
+    Extract a token from headers with multiple fallback strategies.
+
+    This function attempts to extract a token in the following order:
+    1. From standard authorization headers (Bearer token, x-datarobot-api-token, etc.)
+    2. From authorization context metadata (dr_ctx.api_key)
+
+    Args:
+        headers: Dictionary of headers (keys should be lowercase)
+
+    Returns
+    -------
+        The extracted token string, or None if not found
+    """
+    if token := _extract_token_from_headers(headers):
+        return token
+
+    if token := _extract_token_from_auth_context(headers):
+        return token
+
+    return None
+
+
 def get_sdk_client() -> Any:
     """
     Get a DataRobot SDK client, using the user's Bearer token from the request.
 
-    This function attempts to extract the Bearer token from the HTTP request headers.
-    If no token is found, it uses the application credentials.
+    This function attempts to extract the Bearer token from the HTTP request headers
+    with fallback strategies:
+    1. Standard authorization headers (Bearer token, x-datarobot-api-token, etc.)
+    2. Authorization context metadata (dr_ctx.api_key)
+    3. Application credentials as final fallback
     """
     token = None
 
     try:
         headers = get_http_headers()
         if headers:
-            token = _extract_token_from_headers(headers)
+            token = extract_token_from_headers(headers)
             if token:
                 logger.debug("Using API token found in HTTP headers")
     except Exception:

{datarobot_genai-0.1.59 → datarobot_genai-0.1.64}/src/datarobot_genai/drmcp/core/config.py

@@ -196,14 +196,6 @@ class MCPServerConfig(BaseSettings):
         ),
         description="Enable/disable predictive tools",
     )
-    session_secret_key: str | None = Field(
-        default=None,
-        validation_alias=AliasChoices(
-            RUNTIME_PARAM_ENV_VAR_NAME_PREFIX + "SESSION_SECRET_KEY",
-            "SESSION_SECRET_KEY",
-        ),
-        description="Session secret key for the MCP server",
-    )
 
     @field_validator(
         "otel_attributes",

datarobot_genai-0.1.64/src/datarobot_genai/drmcp/core/dynamic_prompts/controllers.py

@@ -0,0 +1,85 @@
+# Copyright 2025 DataRobot, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from fastmcp.prompts.prompt import Prompt
+
+from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_template
+from datarobot_genai.drmcp.core.dynamic_prompts.dr_lib import get_datarobot_prompt_template_version
+from datarobot_genai.drmcp.core.dynamic_prompts.register import (
+    register_prompt_from_datarobot_prompt_management,
+)
+from datarobot_genai.drmcp.core.exceptions import DynamicPromptRegistrationError
+from datarobot_genai.drmcp.core.mcp_instance import mcp
+
+logger = logging.getLogger(__name__)
+
+
+async def register_prompt_from_prompt_template_id_and_version(
+    prompt_template_id: str, prompt_template_version_id: str | None
+) -> Prompt:
+    """Register a Prompt for a specific prompt template ID and version.
+
+    Args:
+        prompt_template_id: The ID of the DataRobot prompt template.
+        prompt_template_version_id: Optional ID of the DataRobot prompt template version.
+            If not provided latest will be used
+
+    Raises
+    ------
+        DynamicPromptRegistrationError: If registration fails at any step.
+
+    Returns
+    -------
+        The registered Prompt instance.
+    """
+    prompt_template = get_datarobot_prompt_template(prompt_template_id)
+
+    if not prompt_template:
+        raise DynamicPromptRegistrationError("Registration failed. Could not find prompt template.")
+
+    if not prompt_template_version_id:
+        return await register_prompt_from_datarobot_prompt_management(
+            prompt_template=prompt_template
+        )
+
+    prompt_template_version = get_datarobot_prompt_template_version(
+        prompt_template_id, prompt_template_version_id
+    )
+
+    if not prompt_template_version:
+        raise DynamicPromptRegistrationError(
+            "Registration failed. Could not find prompt template version."
+        )
+
+    return await register_prompt_from_datarobot_prompt_management(
+        prompt_template=prompt_template, prompt_template_version=prompt_template_version
+    )
+
+
+async def delete_registered_prompt_template(prompt_template_id: str) -> bool:
+    """Delete the prompt registered for the prompt template id in the MCP instance."""
+    prompt_templates_mappings = await mcp.get_prompt_mapping()
+    if prompt_template_id not in prompt_templates_mappings:
+        logger.debug(f"No prompt registered for prompt template id {prompt_template_id}")
+        return False
+
+    prompt_template_version_id, prompt_name = prompt_templates_mappings[prompt_template_id]
+    await mcp.remove_prompt_mapping(prompt_template_id, prompt_template_version_id)
+    logger.info(
+        f"Deleted prompt name {prompt_name} for prompt template id {prompt_template_id}, "
+        f"version {prompt_template_version_id}"
+    )
+    return True