Skip to content

datarobot_genai.drmcputils.clients.datarobot

datarobot

Per-request, thread-safe DataRobot API client for tools.

Thread-safety: backed by :func:datarobot.client.client_configuration, which stores the client in a ContextVar. Concurrent asyncio tasks / threads each get their own client, so we never mutate the process-global dr.Client().

ThreadSafeDataRobotClient

Configure a per-request DataRobot SDK client from the caller's headers.

Source code in datarobot_genai/drmcputils/clients/datarobot.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
class ThreadSafeDataRobotClient:
    """Configure a per-request DataRobot SDK client from the caller's headers."""

    def __init__(self) -> None:
        self.endpoint = get_credentials().datarobot.datarobot_endpoint

    @contextmanager
    def request_user_client(self, *, headers_auth_only: bool = True) -> Iterator[RESTClientObject]:
        """Yield a scoped REST client; same semantics as :func:`request_user_dr_client`."""
        token = get_datarobot_access_token(headers_auth_only=headers_auth_only)
        with client_configuration(token=token, endpoint=self.endpoint):
            # Avoid use-case context from trafaret affecting tool calls.
            DRContext.use_case = None
            yield cast(RESTClientObject, dr.client.get_client())

request_user_client 

request_user_client(*, headers_auth_only: bool = True) -> Iterator[RESTClientObject]

Yield a scoped REST client; same semantics as :func:request_user_dr_client.

Source code in datarobot_genai/drmcputils/clients/datarobot.py 
115
116
117
118
119
120
121
122
@contextmanager
def request_user_client(self, *, headers_auth_only: bool = True) -> Iterator[RESTClientObject]:
    """Yield a scoped REST client; same semantics as :func:`request_user_dr_client`."""
    token = get_datarobot_access_token(headers_auth_only=headers_auth_only)
    with client_configuration(token=token, endpoint=self.endpoint):
        # Avoid use-case context from trafaret affecting tool calls.
        DRContext.use_case = None
        yield cast(RESTClientObject, dr.client.get_client())

get_datarobot_access_token 

get_datarobot_access_token(*, headers_auth_only: bool = True) -> str

Resolve the requesting user's DataRobot API token.

Resolution order:

  1. Token per auth_resolution_strategy (via :func:resolve_datarobot_token).
  2. If unset and headers_auth_only=False, the application API token from credentials (e.g. dynamic tool/prompt registration at server startup).
Raises

ToolError If no token is resolved and headers_auth_only=True.

Source code in datarobot_genai/drmcputils/clients/datarobot.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def get_datarobot_access_token(*, headers_auth_only: bool = True) -> str:
    """Resolve the requesting user's DataRobot API token.

    Resolution order:

    1. Token per ``auth_resolution_strategy`` (via :func:`resolve_datarobot_token`).
    2. If unset and ``headers_auth_only=False``, the application API token from
       credentials (e.g. dynamic tool/prompt registration at server startup).

    Raises
    ------
    ToolError
        If no token is resolved and ``headers_auth_only=True``.
    """
    token = resolve_datarobot_token()
    if token:
        return token
    if headers_auth_only:
        raise ToolError(
            "DataRobot API token not found in headers. "
            "Please provide it via 'Authorization' (Bearer), 'x-datarobot-api-token' headers.",
            kind=ToolErrorKind.AUTHENTICATION,
        )
    return get_credentials().datarobot.datarobot_api_token

request_user_dr_client 

request_user_dr_client(*, headers_auth_only: bool = True) -> Iterator[RESTClientObject]

Yield a request-user-scoped RESTClientObject for the block's duration.

Use inside a with so the configured client stays scoped to this task::

with request_user_dr_client() as client:
    client.post("entitlements/evaluate/", json=...)

Thread-safe: client_configuration() is ContextVar-scoped, so this does not mutate the global dr.Client() and will not mix tokens across concurrent requests.

Source code in datarobot_genai/drmcputils/clients/datarobot.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
@contextmanager
def request_user_dr_client(*, headers_auth_only: bool = True) -> Iterator[RESTClientObject]:
    """Yield a request-user-scoped ``RESTClientObject`` for the block's duration.

    Use inside a ``with`` so the configured client stays scoped to this task::

        with request_user_dr_client() as client:
            client.post("entitlements/evaluate/", json=...)

    Thread-safe: ``client_configuration()`` is ``ContextVar``-scoped, so this
    does not mutate the global ``dr.Client()`` and will not mix tokens across
    concurrent requests.
    """
    token = get_datarobot_access_token(headers_auth_only=headers_auth_only)
    endpoint = get_credentials().datarobot.datarobot_endpoint
    with client_configuration(token=token, endpoint=endpoint):
        # Avoid use-case context from trafaret affecting tool calls.
        DRContext.use_case = None
        yield cast(RESTClientObject, dr.client.get_client())

request_user_dr_sdk 

request_user_dr_sdk(*, headers_auth_only: bool = True) -> Iterator[Any]

Yield the datarobot module with a request-scoped SDK client configured.

Use for SDK calls (e.g. dr.Deployment.get) inside the with block::

with request_user_dr_sdk(headers_auth_only=True):
    deployment = dr.Deployment.get(deployment_id)

Thread-safe: same :func:client_configuration scoping as :func:request_user_dr_client.

Source code in datarobot_genai/drmcputils/clients/datarobot.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
@contextmanager
def request_user_dr_sdk(*, headers_auth_only: bool = True) -> Iterator[Any]:
    """Yield the ``datarobot`` module with a request-scoped SDK client configured.

    Use for SDK calls (e.g. ``dr.Deployment.get``) inside the ``with`` block::

        with request_user_dr_sdk(headers_auth_only=True):
            deployment = dr.Deployment.get(deployment_id)

    Thread-safe: same :func:`client_configuration` scoping as
    :func:`request_user_dr_client`.
    """
    token = get_datarobot_access_token(headers_auth_only=headers_auth_only)
    endpoint = get_credentials().datarobot.datarobot_endpoint
    with client_configuration(token=token, endpoint=endpoint):
        DRContext.use_case = None
        yield dr