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

datarobot_genai/nat/datarobot_auth_provider.py

@@ -11,13 +11,22 @@
 # 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):
     """
@@ -49,5 +58,53 @@ class DataRobotAPIKeyAuthProviderConfig(APIKeyAuthProviderConfig, name="datarobo
 @register_auth_provider(config_type=DataRobotAPIKeyAuthProviderConfig)
 async def datarobot_api_key_client(
     config: DataRobotAPIKeyAuthProviderConfig, builder: Builder
-) -> APIKeyAuthProviderConfig:
+) -> 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

@@ -13,11 +13,15 @@
 # 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
@@ -47,7 +51,7 @@ class DataRobotMCPServerConfig(MCPServerConfig):
     )
     # Authentication configuration
     auth_provider: str | AuthenticationRef | None = Field(
-        default="datarobot_api_key" if config else None,
+        default="datarobot_mcp_auth" if config else None,
         description="Reference to authentication provider",
     )
     command: str | None = Field(
@@ -62,6 +66,55 @@ class DataRobotMCPClientConfig(MCPClientConfig, name="datarobot_mcp_client"):  #
     )
 
 
+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
@@ -108,7 +161,7 @@ async def datarobot_mcp_client_function_group(
     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 = MCPStreamableHTTPClient(
+        client = DataRobotMCPStreamableHTTPClient(
             str(config.server.url),
             auth_provider=auth_provider,
             user_id=base_user_id,

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

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

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

@@ -101,13 +101,13 @@ datarobot_genai/llama_index/base.py,sha256=ovcQQtC-djD_hcLrWdn93jg23AmD6NBEj7xtw
 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=jDeIS9f-8vGbeLy5gQkSjeuHINx5Fh_4BvXYERsgIIk,10516
-datarobot_genai/nat/datarobot_auth_provider.py,sha256=5SBs0xBEZAlBvK-_zOR2cfE_rnn2CoEXeb-Du-rqotc,2117
+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/nat/datarobot_mcp_client.py,sha256=XK3-Tp4hdQmTdI-Zwrl3Xd81qH5RhFTw_UaLTqwnYn4,7531
-datarobot_genai-0.2.6.dist-info/METADATA,sha256=2F0MD9IDjdYCmuhCO76dJHljxSk0cJP0rX17TJQoPaQ,6172
-datarobot_genai-0.2.6.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-datarobot_genai-0.2.6.dist-info/entry_points.txt,sha256=jEW3WxDZ8XIK9-ISmTyt5DbmBb047rFlzQuhY09rGrM,284
-datarobot_genai-0.2.6.dist-info/licenses/AUTHORS,sha256=isJGUXdjq1U7XZ_B_9AH8Qf0u4eX0XyQifJZ_Sxm4sA,80
-datarobot_genai-0.2.6.dist-info/licenses/LICENSE,sha256=U2_VkLIktQoa60Nf6Tbt7E4RMlfhFSjWjcJJfVC-YCE,11341
-datarobot_genai-0.2.6.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,,