datarobot-genai 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

datarobot_genai/nat/agent.py

@@ -123,6 +123,8 @@ class NatAgent(BaseAgent[None]):
         model: str | None = None,
         verbose: bool | str | None = True,
         timeout: int | None = 90,
+        authorization_context: dict[str, Any] | None = None,
+        forwarded_headers: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(
@@ -131,6 +133,8 @@ class NatAgent(BaseAgent[None]):
             model=model,
             verbose=verbose,
             timeout=timeout,
+            authorization_context=authorization_context,
+            forwarded_headers=forwarded_headers,
             **kwargs,
         )
         self.workflow_path = workflow_path

datarobot_genai/nat/datarobot_auth_provider.py

@@ -0,0 +1,110 @@
+# 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.
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from datarobot.core.config import DataRobotAppFrameworkBaseSettings
+from nat.authentication.api_key.api_key_auth_provider import APIKeyAuthProvider
+from nat.authentication.api_key.api_key_auth_provider_config import APIKeyAuthProviderConfig
+from nat.authentication.interfaces import AuthProviderBase
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_auth_provider
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.data_models.authentication import AuthResult
+from nat.data_models.authentication import HeaderCred
+from pydantic import Field
+
+from datarobot_genai.core.mcp.common import MCPConfig
+
+
+class Config(DataRobotAppFrameworkBaseSettings):
+    """
+    Finds variables in the priority order of: env
+    variables (including Runtime Parameters), .env, file_secrets, then
+    Pulumi output variables.
+    """
+
+    datarobot_api_token: str | None = None
+
+
+config = Config()
+
+
+class DataRobotAPIKeyAuthProviderConfig(APIKeyAuthProviderConfig, name="datarobot_api_key"):  # type: ignore[call-arg]
+    raw_key: str = Field(
+        description=(
+            "Raw API token or credential to be injected into the request parameter. "
+            "Used for 'bearer','x-api-key','custom', and other schemes. "
+        ),
+        default=config.datarobot_api_token,
+    )
+    default_user_id: str | None = Field(default="default-user", description="Default user ID")
+    allow_default_user_id_for_tool_calls: bool = Field(
+        default=True, description="Allow default user ID for tool calls"
+    )
+
+
+@register_auth_provider(config_type=DataRobotAPIKeyAuthProviderConfig)
+async def datarobot_api_key_client(
+    config: DataRobotAPIKeyAuthProviderConfig, builder: Builder
+) -> AsyncGenerator[APIKeyAuthProvider]:
+    yield APIKeyAuthProvider(config=config)
+
+
+mcp_config = MCPConfig().server_config
+
+
+class DataRobotMCPAuthProviderConfig(AuthProviderBaseConfig, name="datarobot_mcp_auth"):  # type: ignore[call-arg]
+    headers: dict[str, str] | None = Field(
+        description=("Headers to be used for authentication. "),
+        default=mcp_config["headers"] if mcp_config else None,
+    )
+    default_user_id: str | None = Field(default="default-user", description="Default user ID")
+    allow_default_user_id_for_tool_calls: bool = Field(
+        default=True, description="Allow default user ID for tool calls"
+    )
+
+
+class DataRobotMCPAuthProvider(AuthProviderBase[DataRobotMCPAuthProviderConfig]):
+    def __init__(
+        self, config: DataRobotMCPAuthProviderConfig, config_name: str | None = None
+    ) -> None:
+        assert isinstance(config, DataRobotMCPAuthProviderConfig), (
+            "Config is not DataRobotMCPAuthProviderConfig"
+        )
+        super().__init__(config)
+
+    async def authenticate(self, user_id: str | None = None, **kwargs: Any) -> AuthResult | None:
+        """
+        Authenticate the user using the API key credentials.
+
+        Args:
+            user_id (str): The user ID to authenticate.
+
+        Returns
+        -------
+            AuthenticatedContext: The authenticated context containing headers
+        """
+        return AuthResult(
+            credentials=[
+                HeaderCred(name=name, value=value) for name, value in self.config.headers.items()
+            ]
+        )
+
+
+@register_auth_provider(config_type=DataRobotMCPAuthProviderConfig)
+async def datarobot_mcp_auth_provider(
+    config: DataRobotMCPAuthProviderConfig, builder: Builder
+) -> AsyncGenerator[DataRobotMCPAuthProvider]:
+    yield DataRobotMCPAuthProvider(config=config)

datarobot_genai/nat/datarobot_mcp_client.py

@@ -0,0 +1,234 @@
+# 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 logging
+from datetime import timedelta
+from typing import Literal
+
+import httpx
+from nat.authentication.interfaces import AuthProviderBase
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_function_group
+from nat.data_models.component_ref import AuthenticationRef
+from nat.plugins.mcp.client_base import AuthAdapter
+from nat.plugins.mcp.client_base import MCPSSEClient
+from nat.plugins.mcp.client_base import MCPStdioClient
+from nat.plugins.mcp.client_base import MCPStreamableHTTPClient
+from nat.plugins.mcp.client_config import MCPServerConfig
+from nat.plugins.mcp.client_impl import MCPClientConfig
+from nat.plugins.mcp.client_impl import MCPFunctionGroup
+from nat.plugins.mcp.client_impl import mcp_apply_tool_alias_and_description
+from nat.plugins.mcp.client_impl import mcp_session_tool_function
+from pydantic import Field
+from pydantic import HttpUrl
+
+from datarobot_genai.core.mcp.common import MCPConfig
+
+logger = logging.getLogger(__name__)
+
+config = MCPConfig().server_config
+
+
+class DataRobotMCPServerConfig(MCPServerConfig):
+    transport: Literal["streamable-http", "sse", "stdio"] = Field(
+        default=config["transport"] if config else "stdio",
+        description="Transport type to connect to the MCP server (sse or streamable-http)",
+    )
+    url: HttpUrl | None = Field(
+        default=config["url"] if config else None,
+        description="URL of the MCP server (for sse or streamable-http transport)",
+    )
+    # Authentication configuration
+    auth_provider: str | AuthenticationRef | None = Field(
+        default="datarobot_mcp_auth" if config else None,
+        description="Reference to authentication provider",
+    )
+    command: str | None = Field(
+        default=None if config else "docker",
+        description="Command to run for stdio transport (e.g. 'python' or 'docker')",
+    )
+
+
+class DataRobotMCPClientConfig(MCPClientConfig, name="datarobot_mcp_client"):  # type: ignore[call-arg]
+    server: DataRobotMCPServerConfig = Field(
+        default=DataRobotMCPServerConfig(), description="DataRobot MCP Server configuration"
+    )
+
+
+class DataRobotAuthAdapter(AuthAdapter):
+    async def _get_auth_headers(
+        self, request: httpx.Request | None = None, response: httpx.Response | None = None
+    ) -> dict[str, str]:
+        """Get authentication headers from the NAT auth provider."""
+        try:
+            # Use the user_id passed to this AuthAdapter instance
+            auth_result = await self.auth_provider.authenticate(
+                user_id=self.user_id, response=response
+            )
+            as_kwargs = auth_result.as_requests_kwargs()
+            return as_kwargs["headers"]
+        except Exception as e:
+            logger.warning("Failed to get auth token: %s", e)
+            return {}
+
+
+class DataRobotMCPStreamableHTTPClient(MCPStreamableHTTPClient):
+    def __init__(
+        self,
+        url: str,
+        auth_provider: AuthProviderBase | None = None,
+        user_id: str | None = None,
+        tool_call_timeout: timedelta = timedelta(seconds=60),
+        auth_flow_timeout: timedelta = timedelta(seconds=300),
+        reconnect_enabled: bool = True,
+        reconnect_max_attempts: int = 2,
+        reconnect_initial_backoff: float = 0.5,
+        reconnect_max_backoff: float = 50.0,
+    ):
+        super().__init__(
+            url=url,
+            auth_provider=auth_provider,
+            user_id=user_id,
+            tool_call_timeout=tool_call_timeout,
+            auth_flow_timeout=auth_flow_timeout,
+            reconnect_enabled=reconnect_enabled,
+            reconnect_max_attempts=reconnect_max_attempts,
+            reconnect_initial_backoff=reconnect_initial_backoff,
+            reconnect_max_backoff=reconnect_max_backoff,
+        )
+        effective_user_id = user_id or (
+            auth_provider.config.default_user_id if auth_provider else None
+        )
+        self._httpx_auth = (
+            DataRobotAuthAdapter(auth_provider, effective_user_id) if auth_provider else None
+        )
+
+
+@register_function_group(config_type=DataRobotMCPClientConfig)
+async def datarobot_mcp_client_function_group(
+    config: DataRobotMCPClientConfig, _builder: Builder
+) -> MCPFunctionGroup:
+    """
+    Connect to an MCP server and expose tools as a function group.
+
+    Args:
+        config: The configuration for the MCP client
+        _builder: The builder
+    Returns:
+        The function group
+    """
+    # Resolve auth provider if specified
+    auth_provider = None
+    if config.server.auth_provider:
+        auth_provider = await _builder.get_auth_provider(config.server.auth_provider)
+
+    # Build the appropriate client
+    if config.server.transport == "stdio":
+        if not config.server.command:
+            raise ValueError("command is required for stdio transport")
+        client = MCPStdioClient(
+            config.server.command,
+            config.server.args,
+            config.server.env,
+            tool_call_timeout=config.tool_call_timeout,
+            auth_flow_timeout=config.auth_flow_timeout,
+            reconnect_enabled=config.reconnect_enabled,
+            reconnect_max_attempts=config.reconnect_max_attempts,
+            reconnect_initial_backoff=config.reconnect_initial_backoff,
+            reconnect_max_backoff=config.reconnect_max_backoff,
+        )
+    elif config.server.transport == "sse":
+        client = MCPSSEClient(
+            str(config.server.url),
+            tool_call_timeout=config.tool_call_timeout,
+            auth_flow_timeout=config.auth_flow_timeout,
+            reconnect_enabled=config.reconnect_enabled,
+            reconnect_max_attempts=config.reconnect_max_attempts,
+            reconnect_initial_backoff=config.reconnect_initial_backoff,
+            reconnect_max_backoff=config.reconnect_max_backoff,
+        )
+    elif config.server.transport == "streamable-http":
+        # Use default_user_id for the base client
+        base_user_id = auth_provider.config.default_user_id if auth_provider else None
+        client = DataRobotMCPStreamableHTTPClient(
+            str(config.server.url),
+            auth_provider=auth_provider,
+            user_id=base_user_id,
+            tool_call_timeout=config.tool_call_timeout,
+            auth_flow_timeout=config.auth_flow_timeout,
+            reconnect_enabled=config.reconnect_enabled,
+            reconnect_max_attempts=config.reconnect_max_attempts,
+            reconnect_initial_backoff=config.reconnect_initial_backoff,
+            reconnect_max_backoff=config.reconnect_max_backoff,
+        )
+    else:
+        raise ValueError(f"Unsupported transport: {config.server.transport}")
+
+    logger.info("Configured to use MCP server at %s", client.server_name)
+
+    # Create the MCP function group
+    group = MCPFunctionGroup(config=config)
+
+    # Store shared components for session client creation
+    group._shared_auth_provider = auth_provider
+    group._client_config = config
+
+    async with client:
+        # Expose the live MCP client on the function group instance so other components
+        # (e.g., HTTP endpoints) can reuse the already-established session instead of creating a
+        # new client per request.
+        group.mcp_client = client
+        group.mcp_client_server_name = client.server_name
+        group.mcp_client_transport = client.transport
+
+        all_tools = await client.get_tools()
+        tool_overrides = mcp_apply_tool_alias_and_description(all_tools, config.tool_overrides)
+
+        # Add each tool as a function to the group
+        for tool_name, tool in all_tools.items():
+            # Get override if it exists
+            override = tool_overrides.get(tool_name)
+
+            # Use override values or defaults
+            function_name = override.alias if override and override.alias else tool_name
+            description = (
+                override.description if override and override.description else tool.description
+            )
+
+            # Create the tool function according to configuration
+            tool_fn = mcp_session_tool_function(tool, group)
+
+            # Normalize optional typing for linter/type-checker compatibility
+            single_fn = tool_fn.single_fn
+            if single_fn is None:
+                # Should not happen because FunctionInfo always sets a single_fn
+                logger.warning("Skipping tool %s because single_fn is None", function_name)
+                continue
+
+            input_schema = tool_fn.input_schema
+            # Convert NoneType sentinel to None for FunctionGroup.add_function signature
+            if input_schema is type(None):  # noqa: E721
+                input_schema = None
+
+            # Add to group
+            logger.info("Adding tool %s to group", function_name)
+            group.add_function(
+                name=function_name,
+                description=description,
+                fn=single_fn,
+                input_schema=input_schema,
+                converters=tool_fn.converters,
+            )
+
+        yield group

{datarobot_genai-0.2.5.dist-info → datarobot_genai-0.2.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.5
+Version: 0.2.7
 Summary: Generic helpers for GenAI
 Project-URL: Homepage, https://github.com/datarobot-oss/datarobot-genai
 Author: DataRobot, Inc.
@@ -63,6 +63,7 @@ Requires-Dist: anyio==4.11.0; extra == 'nat'
 Requires-Dist: crewai>=1.1.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: llama-index-llms-litellm<0.7.0,>=0.4.1; extra == 'nat'
 Requires-Dist: nvidia-nat-langchain==1.3.0; (python_version >= '3.11') and extra == 'nat'
+Requires-Dist: nvidia-nat-mcp==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat-opentelemetry==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: nvidia-nat==1.3.0; (python_version >= '3.11') and extra == 'nat'
 Requires-Dist: opentelemetry-instrumentation-crewai<1.0.0,>=0.40.5; extra == 'nat'

{datarobot_genai-0.2.5.dist-info → datarobot_genai-0.2.7.dist-info}/RECORD

@@ -100,12 +100,14 @@ datarobot_genai/llama_index/agent.py,sha256=V6ZsD9GcBDJS-RJo1tJtIHhyW69_78gM6_fO
 datarobot_genai/llama_index/base.py,sha256=ovcQQtC-djD_hcLrWdn93jg23AmD6NBEj7xtw4a6K6c,14481
 datarobot_genai/llama_index/mcp.py,sha256=leXqF1C4zhuYEKFwNEfZHY4dsUuGZk3W7KArY-zxVL8,2645
 datarobot_genai/nat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datarobot_genai/nat/agent.py,sha256=siBLDWAff2-JwZ8Q3iNpM_e4_IoSwG9IvY0hyEjNenw,10292
+datarobot_genai/nat/agent.py,sha256=jDeIS9f-8vGbeLy5gQkSjeuHINx5Fh_4BvXYERsgIIk,10516
+datarobot_genai/nat/datarobot_auth_provider.py,sha256=Z4NSsrHxK8hUeiqtK_lryHsUuZC74ziNo_FHbsZgtiM,4230
 datarobot_genai/nat/datarobot_llm_clients.py,sha256=STzAZ4OF8U-Y_cUTywxmKBGVotwsnbGP6vTojnu6q0g,9921
 datarobot_genai/nat/datarobot_llm_providers.py,sha256=aDoQcTeGI-odqydPXEX9OGGNFbzAtpqzTvHHEkmJuEQ,4963
-datarobot_genai-0.2.5.dist-info/METADATA,sha256=ekA7W-PGh36MkW_Whpb4o7SP2ALvXk9-jzibzcT_Z-8,6088
-datarobot_genai-0.2.5.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.5.dist-info/entry_points.txt,sha256=CZhmZcSyt_RBltgLN_b9xasJD6J5SaDc_z7K0wuOY9Y,150
-datarobot_genai-0.2.5.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.5.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.5.dist-info/RECORD,,
+datarobot_genai/nat/datarobot_mcp_client.py,sha256=35FzilxNp4VqwBYI0NsOc91-xZm1C-AzWqrOdDy962A,9612
+datarobot_genai-0.2.7.dist-info/METADATA,sha256=dWVPDIxuaXhv5tFM-_yhFvatzWmHzLvZtcC5pM_NsLw,6172
+datarobot_genai-0.2.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+datarobot_genai-0.2.7.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
+datarobot_genai-0.2.7.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
+datarobot_genai-0.2.7.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
+datarobot_genai-0.2.7.dist-info/RECORD,,

{datarobot_genai-0.2.5.dist-info → datarobot_genai-0.2.7.dist-info}/entry_points.txt

@@ -1,3 +1,5 @@
 [nat.plugins]
+datarobot_auth_provider = datarobot_genai.nat.datarobot_auth_provider
 datarobot_llm_clients = datarobot_genai.nat.datarobot_llm_clients
 datarobot_llm_providers = datarobot_genai.nat.datarobot_llm_providers
+datarobot_mcp_client = datarobot_genai.nat.datarobot_mcp_client