Skip to content

datarobot_genai.dragent.plugins.datarobot_auth_provider

datarobot_auth_provider

DataRobotMCPAuthProvider

Bases: AuthProviderBase[DataRobotMCPAuthProviderConfig]

Source code in datarobot_genai/dragent/plugins/datarobot_auth_provider.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
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
        """
        forwarded_headers = extract_datarobot_headers_from_context()
        authentication_context = extract_authorization_from_context()
        mcp_config = MCPConfig(
            forwarded_headers=forwarded_headers, authorization_context=authentication_context
        ).server_config

        # in dragent we get forwarded_headers and authentication_context from Context
        # in drum we write self.config.headers with a custom loader
        auth_headers = {}
        if mcp_config:
            auth_headers.update(mcp_config["headers"])
        if self.config.headers:
            auth_headers.update(self.config.headers)

        return AuthResult(
            credentials=[HeaderCred(name=name, value=value) for name, value in auth_headers.items()]
        )

authenticate async 

authenticate(user_id: str | None = None, **kwargs: Any) -> AuthResult | None

Authenticate the user using the API key credentials.

Parameters:

Name Type Description Default
user_id str

The user ID to authenticate.

 None
Returns
AuthenticatedContext: The authenticated context containing headers
Source code in datarobot_genai/dragent/plugins/datarobot_auth_provider.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
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
    """
    forwarded_headers = extract_datarobot_headers_from_context()
    authentication_context = extract_authorization_from_context()
    mcp_config = MCPConfig(
        forwarded_headers=forwarded_headers, authorization_context=authentication_context
    ).server_config

    # in dragent we get forwarded_headers and authentication_context from Context
    # in drum we write self.config.headers with a custom loader
    auth_headers = {}
    if mcp_config:
        auth_headers.update(mcp_config["headers"])
    if self.config.headers:
        auth_headers.update(self.config.headers)

    return AuthResult(
        credentials=[HeaderCred(name=name, value=value) for name, value in auth_headers.items()]
    )