Skip to content

datarobot_genai.core.chat.auth

auth

Authorization context helpers for chat flows.

get_authorization_context_from_headers 

get_authorization_context_from_headers(headers: dict[str, str], secret_key: str | None = None) -> dict[str, Any] | None

Extract authorization context from headers using AuthContextHeaderHandler.

Parameters

headers : dict[str, str] HTTP headers from which to extract the authorization context. secret_key : str | None Secret key for JWT decoding. If None, retrieves from environment variable.

Returns

dict[str, Any] | None The extracted authorization context, or None if not found.

Source code in datarobot_genai/core/chat/auth.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def get_authorization_context_from_headers(
    headers: dict[str, str],
    secret_key: str | None = None,
) -> dict[str, Any] | None:
    """Extract authorization context from headers using AuthContextHeaderHandler.

    Parameters
    ----------
    headers : dict[str, str]
        HTTP headers from which to extract the authorization context.
    secret_key : str | None
        Secret key for JWT decoding. If None, retrieves from environment variable.

    Returns
    -------
    dict[str, Any] | None
        The extracted authorization context, or None if not found.
    """
    handler = AuthContextHeaderHandler(secret_key=secret_key)
    if context := handler.get_context(headers):
        return context.model_dump()
    return None

get_authorization_context_from_params 

get_authorization_context_from_params(completion_create_params: CompletionCreateParams | CompletionCreateParamsNonStreaming | CompletionCreateParamsStreaming) -> dict[str, Any] | None

Extract authorization context from completion create parameters.

Parameters

completion_create_params : CompletionCreateParams The parameters used to create the completion.

Returns

dict[str, Any] | None The extracted authorization context, or None if not found.

Source code in datarobot_genai/core/chat/auth.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def get_authorization_context_from_params(
    completion_create_params: CompletionCreateParams
    | CompletionCreateParamsNonStreaming
    | CompletionCreateParamsStreaming,
) -> dict[str, Any] | None:
    """Extract authorization context from completion create parameters.

    Parameters
    ----------
    completion_create_params : CompletionCreateParams
        The parameters used to create the completion.

    Returns
    -------
    dict[str, Any] | None
        The extracted authorization context, or None if not found.
    """
    return completion_create_params.get("authorization_context", None)

resolve_authorization_context 

resolve_authorization_context(completion_create_params: CompletionCreateParams | CompletionCreateParamsNonStreaming | CompletionCreateParamsStreaming, **kwargs: Any) -> dict[str, Any]

Resolve the authorization context for the agent.

Authorization context is required for propagating information needed by downstream agents and tools to retrieve access tokens to connect to external services. This method extracts the authorization context from either the incoming HTTP headers or the completion create parameters.

Parameters

completion_create_params : CompletionCreateParams | CompletionCreateParamsNonStreaming | CompletionCreateParamsStreaming Parameters supplied to the completion API. May include a fallback authorization_context mapping under the same key. **kwargs : Any Additional keyword arguments. Expected to include a headers key containing incoming HTTP headers as dict[str, str].

Returns

dict[str, Any] The initialized authorization context.

Source code in datarobot_genai/core/chat/auth.py 
 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
def resolve_authorization_context(
    completion_create_params: CompletionCreateParams
    | CompletionCreateParamsNonStreaming
    | CompletionCreateParamsStreaming,
    **kwargs: Any,
) -> dict[str, Any]:
    """Resolve the authorization context for the agent.

    Authorization context is required for propagating information needed by downstream
    agents and tools to retrieve access tokens to connect to external services. This method
    extracts the authorization context from either the incoming HTTP headers or the completion
    create parameters.

    Parameters
    ----------
    completion_create_params : CompletionCreateParams | CompletionCreateParamsNonStreaming |
        CompletionCreateParamsStreaming
        Parameters supplied to the completion API. May include a fallback
        ``authorization_context`` mapping under the same key.
    **kwargs : Any
        Additional keyword arguments. Expected to include a ``headers`` key
        containing incoming HTTP headers as ``dict[str, str]``.

    Returns
    -------
    dict[str, Any]
        The initialized authorization context.
    """
    incoming_headers = kwargs.get("headers", {})

    # Recommended way of propagating authorization context is via headers
    # with JWT endoding/decoding for additional security. The completion params
    # is used as a fallback for backward compatibility only and may be removed in
    # the future.
    authorization_context: dict[str, Any] = (
        get_authorization_context_from_headers(incoming_headers)
        or get_authorization_context_from_params(completion_create_params)
        or {}
    )

    return authorization_context