datarobot-genai 0.2.5__tar.gz → 0.2.9__tar.gz

{datarobot_genai-0.2.5 → datarobot_genai-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datarobot-genai
-Version: 0.2.5
+Version: 0.2.9
 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 → datarobot_genai-0.2.9}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "datarobot-genai"
-version = "0.2.5"
+version = "0.2.9"
 description = "Generic helpers for GenAI"
 readme = "README.md"
 requires-python = ">=3.10, <3.13"
@@ -35,6 +35,8 @@ Homepage = "https://github.com/datarobot-oss/datarobot-genai"
 [project.entry-points.'nat.plugins']
 datarobot_llm_providers = "datarobot_genai.nat.datarobot_llm_providers"
 datarobot_llm_clients = "datarobot_genai.nat.datarobot_llm_clients"
+datarobot_mcp_client = "datarobot_genai.nat.datarobot_mcp_client"
+datarobot_auth_provider = "datarobot_genai.nat.datarobot_auth_provider"
 
 [project.optional-dependencies]
 crewai = [
@@ -63,6 +65,7 @@ nat = [
   "nvidia-nat==1.3.0; python_version >= '3.11'",
   "nvidia-nat-opentelemetry==1.3.0; python_version >= '3.11'",
   "nvidia-nat-langchain==1.3.0; python_version >= '3.11'",
+  "nvidia-nat-mcp==1.3.0; python_version >= '3.11'",
   "crewai>=1.1.0; python_version >= '3.11'",
   "llama-index-llms-litellm>=0.4.1,<0.7.0",  # Need this to support datarobot-llm plugin
   "opentelemetry-instrumentation-crewai>=0.40.5,<1.0.0",

{datarobot_genai-0.2.5 → datarobot_genai-0.2.9}/src/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-0.2.9/src/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-0.2.9/src/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