Skip to content

datarobot_genai.core.utils.auth

auth

DRAppCtx

Bases: BaseModel

DataRobot application context from authorization metadata.

Source code in datarobot_genai/core/utils/auth.py 
41
42
43
44
45
class DRAppCtx(BaseModel):
    """DataRobot application context from authorization metadata."""

    email: str | None = None
    api_key: str | None = None

AuthContextHeaderHandler

Manages encoding and decoding of authorization context into JWT tokens.

This class provides a consistent interface for encoding auth context into JWT tokens and exchanging them via HTTP headers across multiple applications.

Source code in datarobot_genai/core/utils/auth.py 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
class AuthContextHeaderHandler:
    """Manages encoding and decoding of authorization context into JWT tokens.

    This class provides a consistent interface for encoding auth context into JWT tokens
    and exchanging them via HTTP headers across multiple applications.
    """

    HEADER_NAME = "X-DataRobot-Authorization-Context"
    DEFAULT_ALGORITHM = "HS256"

    def __init__(
        self,
        secret_key: str | None = None,
        algorithm: str = DEFAULT_ALGORITHM,
        validate_signature: bool = True,
    ) -> None:
        """Initialize the handler.

        Parameters
        ----------
        secret_key : Optional[str]
            Secret key for JWT encoding/decoding. If None, tokens will be unsigned (insecure).
        algorithm : str
            JWT algorithm. Default is "HS256".
        validate_signature : bool
            Whether to validate JWT signatures. Default is True.

        Raises
        ------
        ValueError
            If algorithm is 'none' (insecure).
        """
        if algorithm is None:
            raise ValueError("Algorithm None is not allowed. Use a secure algorithm like HS256.")

        # Get secret key from parameter, config, or environment variable
        # Handle the case where AuthContextConfig() initialization fails due to
        # a bug in the datarobot package when SESSION_SECRET_KEY is not set
        if secret_key:
            self.secret_key = secret_key
        else:
            try:
                config = AuthContextConfig()
                self.secret_key = config.session_secret_key or ""
            except (TypeError, AttributeError, Exception):
                # Fallback to reading environment variable directly if config initialization fails
                # This can happen when SESSION_SECRET_KEY is not set and the datarobot package's
                # getenv function encounters a bug with None values
                # it tries to check if "apiToken" in payload: when payload is None
                self.secret_key = ""

        self.algorithm = algorithm
        self.validate_signature = validate_signature

    @property
    def header(self) -> str:
        """Get the header name for authorization context."""
        return self.HEADER_NAME

    def get_header(self, authorization_context: dict[str, Any] | None = None) -> dict[str, str]:
        """Get the authorization context header with encoded JWT token."""
        token = self.encode(authorization_context)
        if not token:
            return {}

        return {self.header: token}

    def encode(self, authorization_context: dict[str, Any] | None = None) -> str | None:
        """Encode the current authorization context into a JWT token."""
        auth_context = authorization_context or get_authorization_context()
        if not auth_context:
            return None

        if not self.secret_key:
            warnings.warn(
                "No secret key provided. Please make sure SESSION_SECRET_KEY is set. "
                "JWT tokens will be signed with an insecure key. This should only be "
                "used for testing."
            )

        # PyJWT rejects empty string keys; use a dummy key when none is configured
        key = self.secret_key or "insecure-placeholder-key"
        return jwt.encode(auth_context, key, algorithm=self.algorithm)

    def decode(self, token: str) -> dict[str, Any] | None:
        """Decode a JWT token into the authorization context."""
        if not token:
            return None

        if not self.secret_key and self.validate_signature:
            logger.error(
                "No secret key provided. Cannot validate signature. "
                "Provide a secret key or set validate_signature to False."
            )
            return None

        try:
            decoded = jwt.decode(
                jwt=token,
                key=self.secret_key,
                algorithms=[self.algorithm],
                options={"verify_signature": self.validate_signature},
            )
        except jwt.ExpiredSignatureError:
            logger.info("JWT token has expired.")
            return None
        except jwt.InvalidTokenError:
            logger.warning("JWT token is invalid or malformed.")
            return None

        if not isinstance(decoded, dict):
            logger.warning("Decoded JWT token is not a dictionary.")
            return None

        return decoded

    def get_context(self, headers: dict[str, str]) -> AuthCtx | None:
        """Extract and validate authorization context from headers.

        Parameters
        ----------
        headers : Dict[str, str]
            HTTP headers containing the authorization context.

        Returns
        -------
        Optional[AuthCtx]
            Validated authorization context or None if validation fails.
        """
        token = headers.get(self.header) or headers.get(self.header.lower())
        if not token:
            logger.debug("No authorization context header found")
            return None

        auth_ctx_dict = self.decode(token)
        if not auth_ctx_dict:
            logger.debug("Failed to decode auth context from token")
            return None

        try:
            return AuthCtx(**auth_ctx_dict)
        except Exception as e:
            logger.error(f"Failed to create AuthCtx from decoded token: {e}", exc_info=True)
            return None

header property 

header: str

Get the header name for authorization context.

__init__ 

__init__(secret_key: str | None = None, algorithm: str = DEFAULT_ALGORITHM, validate_signature: bool = True) -> None

Initialize the handler.

Parameters

secret_key : Optional[str] Secret key for JWT encoding/decoding. If None, tokens will be unsigned (insecure). algorithm : str JWT algorithm. Default is "HS256". validate_signature : bool Whether to validate JWT signatures. Default is True.

Raises

ValueError If algorithm is 'none' (insecure).

Source code in datarobot_genai/core/utils/auth.py 
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def __init__(
    self,
    secret_key: str | None = None,
    algorithm: str = DEFAULT_ALGORITHM,
    validate_signature: bool = True,
) -> None:
    """Initialize the handler.

    Parameters
    ----------
    secret_key : Optional[str]
        Secret key for JWT encoding/decoding. If None, tokens will be unsigned (insecure).
    algorithm : str
        JWT algorithm. Default is "HS256".
    validate_signature : bool
        Whether to validate JWT signatures. Default is True.

    Raises
    ------
    ValueError
        If algorithm is 'none' (insecure).
    """
    if algorithm is None:
        raise ValueError("Algorithm None is not allowed. Use a secure algorithm like HS256.")

    # Get secret key from parameter, config, or environment variable
    # Handle the case where AuthContextConfig() initialization fails due to
    # a bug in the datarobot package when SESSION_SECRET_KEY is not set
    if secret_key:
        self.secret_key = secret_key
    else:
        try:
            config = AuthContextConfig()
            self.secret_key = config.session_secret_key or ""
        except (TypeError, AttributeError, Exception):
            # Fallback to reading environment variable directly if config initialization fails
            # This can happen when SESSION_SECRET_KEY is not set and the datarobot package's
            # getenv function encounters a bug with None values
            # it tries to check if "apiToken" in payload: when payload is None
            self.secret_key = ""

    self.algorithm = algorithm
    self.validate_signature = validate_signature

get_header 

get_header(authorization_context: dict[str, Any] | None = None) -> dict[str, str]

Get the authorization context header with encoded JWT token.

Source code in datarobot_genai/core/utils/auth.py 
107
108
109
110
111
112
113
def get_header(self, authorization_context: dict[str, Any] | None = None) -> dict[str, str]:
    """Get the authorization context header with encoded JWT token."""
    token = self.encode(authorization_context)
    if not token:
        return {}

    return {self.header: token}

encode 

encode(authorization_context: dict[str, Any] | None = None) -> str | None

Encode the current authorization context into a JWT token.

Source code in datarobot_genai/core/utils/auth.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def encode(self, authorization_context: dict[str, Any] | None = None) -> str | None:
    """Encode the current authorization context into a JWT token."""
    auth_context = authorization_context or get_authorization_context()
    if not auth_context:
        return None

    if not self.secret_key:
        warnings.warn(
            "No secret key provided. Please make sure SESSION_SECRET_KEY is set. "
            "JWT tokens will be signed with an insecure key. This should only be "
            "used for testing."
        )

    # PyJWT rejects empty string keys; use a dummy key when none is configured
    key = self.secret_key or "insecure-placeholder-key"
    return jwt.encode(auth_context, key, algorithm=self.algorithm)

decode 

decode(token: str) -> dict[str, Any] | None

Decode a JWT token into the authorization context.

Source code in datarobot_genai/core/utils/auth.py 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def decode(self, token: str) -> dict[str, Any] | None:
    """Decode a JWT token into the authorization context."""
    if not token:
        return None

    if not self.secret_key and self.validate_signature:
        logger.error(
            "No secret key provided. Cannot validate signature. "
            "Provide a secret key or set validate_signature to False."
        )
        return None

    try:
        decoded = jwt.decode(
            jwt=token,
            key=self.secret_key,
            algorithms=[self.algorithm],
            options={"verify_signature": self.validate_signature},
        )
    except jwt.ExpiredSignatureError:
        logger.info("JWT token has expired.")
        return None
    except jwt.InvalidTokenError:
        logger.warning("JWT token is invalid or malformed.")
        return None

    if not isinstance(decoded, dict):
        logger.warning("Decoded JWT token is not a dictionary.")
        return None

    return decoded

get_context 

get_context(headers: dict[str, str]) -> AuthCtx | None

Extract and validate authorization context from headers.

Parameters

headers : Dict[str, str] HTTP headers containing the authorization context.

Returns

Optional[AuthCtx] Validated authorization context or None if validation fails.

Source code in datarobot_genai/core/utils/auth.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
def get_context(self, headers: dict[str, str]) -> AuthCtx | None:
    """Extract and validate authorization context from headers.

    Parameters
    ----------
    headers : Dict[str, str]
        HTTP headers containing the authorization context.

    Returns
    -------
    Optional[AuthCtx]
        Validated authorization context or None if validation fails.
    """
    token = headers.get(self.header) or headers.get(self.header.lower())
    if not token:
        logger.debug("No authorization context header found")
        return None

    auth_ctx_dict = self.decode(token)
    if not auth_ctx_dict:
        logger.debug("Failed to decode auth context from token")
        return None

    try:
        return AuthCtx(**auth_ctx_dict)
    except Exception as e:
        logger.error(f"Failed to create AuthCtx from decoded token: {e}", exc_info=True)
        return None

TokenRetriever

Bases: Protocol

Protocol for OAuth token retrievers.

Source code in datarobot_genai/core/utils/auth.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
class TokenRetriever(Protocol):
    """Protocol for OAuth token retrievers."""

    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
        """Filter identities to only those valid for this retriever implementation."""
        ...

    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
        """Refresh the access token for the given identity ID.

        Parameters
        ----------
        identity_id : str
            The provider identity ID to refresh the token for.

        Returns
        -------
        OAuthToken
            The refreshed OAuth token.
        """
        ...

filter_identities 

filter_identities(identities: Sequence[Identity]) -> list[Identity]

Filter identities to only those valid for this retriever implementation.

Source code in datarobot_genai/core/utils/auth.py 
200
201
202
def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
    """Filter identities to only those valid for this retriever implementation."""
    ...

refresh_access_token async 

refresh_access_token(identity: Identity) -> OAuthToken

Refresh the access token for the given identity ID.

Parameters

identity_id : str The provider identity ID to refresh the token for.

Returns

OAuthToken The refreshed OAuth token.

Source code in datarobot_genai/core/utils/auth.py 
204
205
206
207
208
209
210
211
212
213
214
215
216
217
async def refresh_access_token(self, identity: Identity) -> OAuthToken:
    """Refresh the access token for the given identity ID.

    Parameters
    ----------
    identity_id : str
        The provider identity ID to refresh the token for.

    Returns
    -------
    OAuthToken
        The refreshed OAuth token.
    """
    ...

DatarobotTokenRetriever

Retrieves OAuth tokens using the DataRobot platform.

Source code in datarobot_genai/core/utils/auth.py 
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
class DatarobotTokenRetriever:
    """Retrieves OAuth tokens using the DataRobot platform."""

    def __init__(self) -> None:
        self._client = DatarobotOAuthClient()

    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
        """Filter oauth2 identities to only those with provider_identity_id.

        The `provider_identity_id` is required in order to identify the provider
        and retrieve the access token from DataRobot OAuth Providers service.
        """
        return [i for i in identities if i.type == "oauth2" and i.provider_identity_id]

    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
        """Refresh the access token using DataRobot's OAuth client."""
        return await self._client.refresh_access_token(identity_id=identity.provider_identity_id)

filter_identities 

filter_identities(identities: Sequence[Identity]) -> list[Identity]

Filter oauth2 identities to only those with provider_identity_id.

The provider_identity_id is required in order to identify the provider and retrieve the access token from DataRobot OAuth Providers service.

Source code in datarobot_genai/core/utils/auth.py 
226
227
228
229
230
231
232
def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
    """Filter oauth2 identities to only those with provider_identity_id.

    The `provider_identity_id` is required in order to identify the provider
    and retrieve the access token from DataRobot OAuth Providers service.
    """
    return [i for i in identities if i.type == "oauth2" and i.provider_identity_id]

refresh_access_token async 

refresh_access_token(identity: Identity) -> OAuthToken

Refresh the access token using DataRobot's OAuth client.

Source code in datarobot_genai/core/utils/auth.py 
234
235
236
async def refresh_access_token(self, identity: Identity) -> OAuthToken:
    """Refresh the access token using DataRobot's OAuth client."""
    return await self._client.refresh_access_token(identity_id=identity.provider_identity_id)

AuthlibTokenRetriever

Retrieves OAuth tokens from a generic Authlib-based endpoint.

Source code in datarobot_genai/core/utils/auth.py 
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
class AuthlibTokenRetriever:
    """Retrieves OAuth tokens from a generic Authlib-based endpoint."""

    def __init__(self, application_endpoint: str) -> None:
        if not application_endpoint:
            raise ValueError("AuthlibTokenRetriever requires 'application_endpoint'.")
        self.application_endpoint = application_endpoint.rstrip("/")

    def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
        """Filter identities to only OAuth2 identities."""
        return [i for i in identities if i.type == "oauth2" and i.provider_identity_id is None]

    async def refresh_access_token(self, identity: Identity) -> OAuthToken:
        """Retrieve an OAuth token via an HTTP POST request.

        Parameters
        ----------
        identity : Identity
            The identity to retrieve the token for.

        Returns
        -------
        OAuthToken
            The retrieved OAuth token.
        """
        api_token = os.environ.get("DATAROBOT_API_TOKEN")
        if not api_token:
            raise ValueError("DATAROBOT_API_TOKEN environment variable is required but not set.")

        token_url = f"{self.application_endpoint}/oauth/token/"
        headers = {"Authorization": f"Bearer {api_token}"}
        payload = {"identity_id": identity.id}
        timeout = aiohttp.ClientTimeout(total=30)

        try:
            async with aiohttp.ClientSession(timeout=timeout) as session:
                async with session.post(token_url, headers=headers, json=payload) as response:
                    response.raise_for_status()
                    data = await response.json()
                    logger.debug(f"Retrieved access token from {token_url}")
                    return OAuthToken(**data)
        except aiohttp.ClientError as e:
            logger.error(f"Error retrieving token from {token_url}: {e}")
            raise

filter_identities 

filter_identities(identities: Sequence[Identity]) -> list[Identity]

Filter identities to only OAuth2 identities.

Source code in datarobot_genai/core/utils/auth.py 
247
248
249
def filter_identities(self, identities: Sequence[Identity]) -> list[Identity]:
    """Filter identities to only OAuth2 identities."""
    return [i for i in identities if i.type == "oauth2" and i.provider_identity_id is None]

refresh_access_token async 

refresh_access_token(identity: Identity) -> OAuthToken

Retrieve an OAuth token via an HTTP POST request.

Parameters

identity : Identity The identity to retrieve the token for.

Returns

OAuthToken The retrieved OAuth token.

Source code in datarobot_genai/core/utils/auth.py 
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
async def refresh_access_token(self, identity: Identity) -> OAuthToken:
    """Retrieve an OAuth token via an HTTP POST request.

    Parameters
    ----------
    identity : Identity
        The identity to retrieve the token for.

    Returns
    -------
    OAuthToken
        The retrieved OAuth token.
    """
    api_token = os.environ.get("DATAROBOT_API_TOKEN")
    if not api_token:
        raise ValueError("DATAROBOT_API_TOKEN environment variable is required but not set.")

    token_url = f"{self.application_endpoint}/oauth/token/"
    headers = {"Authorization": f"Bearer {api_token}"}
    payload = {"identity_id": identity.id}
    timeout = aiohttp.ClientTimeout(total=30)

    try:
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.post(token_url, headers=headers, json=payload) as response:
                response.raise_for_status()
                data = await response.json()
                logger.debug(f"Retrieved access token from {token_url}")
                return OAuthToken(**data)
    except aiohttp.ClientError as e:
        logger.error(f"Error retrieving token from {token_url}: {e}")
        raise

OAuthConfig

Bases: BaseModel

Configuration extracted from AuthCtx metadata for OAuth operations.

Source code in datarobot_genai/core/utils/auth.py 
285
286
287
288
289
290
291
292
293
294
295
296
297
class OAuthConfig(BaseModel):
    """Configuration extracted from AuthCtx metadata for OAuth operations."""

    implementation: str = "datarobot"
    application_endpoint: str | None = None

    @classmethod
    def from_auth_ctx(cls, auth_ctx: AuthCtx) -> "OAuthConfig":
        metadata = auth_ctx.metadata or {}
        return cls(
            implementation=metadata.get("oauth_implementation", "datarobot"),
            application_endpoint=metadata.get("application_endpoint"),
        )

AsyncOAuthTokenProvider

Provides OAuth tokens for authorized users.

This class manages OAuth token retrieval for users with multiple identity providers. It uses either DataRobot or Authlib as the OAuth token storage and refresh backend based on the auth context metadata.

Source code in datarobot_genai/core/utils/auth.py 
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
class AsyncOAuthTokenProvider:
    """Provides OAuth tokens for authorized users.

    This class manages OAuth token retrieval for users with multiple identity providers.
    It uses either DataRobot or Authlib as the OAuth token storage and refresh backend
    based on the auth context metadata.
    """

    def __init__(self, auth_ctx: AuthCtx) -> None:
        """Initialize the provider with an authorization context.

        Parameters
        ----------
        auth_ctx : AuthCtx
            The authorization context containing user identities and metadata.
        """
        self.auth_ctx = auth_ctx
        config = OAuthConfig.from_auth_ctx(auth_ctx)
        self._retriever = create_token_retriever(config)

    def _get_identity(self, provider_type: str | None) -> Identity:
        """Get identity from auth context, filtered by provider_type if specified."""
        oauth_identities = self._retriever.filter_identities(self.auth_ctx.identities)
        if not oauth_identities:
            raise OAuthProviderNotFound("No OAuth provider found.")

        if provider_type is None:
            if len(oauth_identities) > 1:
                raise OAuthValidationErr(
                    "Multiple OAuth providers found. Specify 'provider_type' parameter."
                )
            return oauth_identities[0]

        identity = next((i for i in oauth_identities if i.provider_type == provider_type), None)
        if identity is None:
            raise OAuthValidationErr(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 an OAuth access token for the specified auth type and provider.

        Parameters
        ----------
        auth_type : ToolAuth
            Authentication type (only OBO is supported).
        provider_type : str, optional
            The specific provider to use (e.g., 'google'). Required if multiple
            identities are available.

        Returns
        -------
        str
            The retrieved OAuth access token.

        Raises
        ------
        ValueError
            If the auth type is unsupported or if a suitable identity cannot be found.
        """
        if auth_type != ToolAuth.OBO:
            raise ValueError(
                f"Unsupported auth type: {auth_type}. Only OBO (on-behalf-of) is supported."
            )

        identity = self._get_identity(provider_type)
        token_data = await self._retriever.refresh_access_token(identity)
        return token_data.access_token

__init__ 

__init__(auth_ctx: AuthCtx) -> None

Initialize the provider with an authorization context.

Parameters

auth_ctx : AuthCtx The authorization context containing user identities and metadata.

Source code in datarobot_genai/core/utils/auth.py 
335
336
337
338
339
340
341
342
343
344
345
def __init__(self, auth_ctx: AuthCtx) -> None:
    """Initialize the provider with an authorization context.

    Parameters
    ----------
    auth_ctx : AuthCtx
        The authorization context containing user identities and metadata.
    """
    self.auth_ctx = auth_ctx
    config = OAuthConfig.from_auth_ctx(auth_ctx)
    self._retriever = create_token_retriever(config)

get_token async 

get_token(auth_type: ToolAuth, provider_type: str | None = None) -> str

Get an OAuth access token for the specified auth type and provider.

Parameters

auth_type : ToolAuth Authentication type (only OBO is supported). provider_type : str, optional The specific provider to use (e.g., 'google'). Required if multiple identities are available.

Returns

str The retrieved OAuth access token.

Raises

ValueError If the auth type is unsupported or if a suitable identity cannot be found.

Source code in datarobot_genai/core/utils/auth.py 
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
async def get_token(self, auth_type: ToolAuth, provider_type: str | None = None) -> str:
    """Get an OAuth access token for the specified auth type and provider.

    Parameters
    ----------
    auth_type : ToolAuth
        Authentication type (only OBO is supported).
    provider_type : str, optional
        The specific provider to use (e.g., 'google'). Required if multiple
        identities are available.

    Returns
    -------
    str
        The retrieved OAuth access token.

    Raises
    ------
    ValueError
        If the auth type is unsupported or if a suitable identity cannot be found.
    """
    if auth_type != ToolAuth.OBO:
        raise ValueError(
            f"Unsupported auth type: {auth_type}. Only OBO (on-behalf-of) is supported."
        )

    identity = self._get_identity(provider_type)
    token_data = await self._retriever.refresh_access_token(identity)
    return token_data.access_token

create_token_retriever 

create_token_retriever(config: OAuthConfig) -> TokenRetriever

Create a token retriever based on the OAuth configuration.

Parameters

config : OAuthConfig The OAuth configuration specifying implementation type and endpoints.

Returns

TokenRetriever The configured token retriever instance.

Source code in datarobot_genai/core/utils/auth.py 
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
def create_token_retriever(config: OAuthConfig) -> TokenRetriever:
    """Create a token retriever based on the OAuth configuration.

    Parameters
    ----------
    config : OAuthConfig
        The OAuth configuration specifying implementation type and endpoints.

    Returns
    -------
    TokenRetriever
        The configured token retriever instance.
    """
    if config.implementation == "datarobot":
        return DatarobotTokenRetriever()

    if config.implementation == "authlib":
        if not config.application_endpoint:
            raise ValueError("Required 'application_endpoint' not found in metadata.")
        return AuthlibTokenRetriever(config.application_endpoint)

    raise ValueError(
        f"Unsupported OAuth implementation: '{config.implementation}'. "
        f"Supported values: datarobot, authlib."
    )