Skip to content

datarobot_genai.core.chat.client

client

Client for interacting with Agent Tools deployments for chat and scoring.

ToolClient

Client for interacting with Agent Tools Deployments.

This class provides methods to call the custom model tool using various hooks: score, score_unstructured, and chat. When the authorization_context is set, the client automatically propagates it to the agent tool. The authorization_context is required for retrieving access tokens to connect to external services.

Source code in datarobot_genai/core/chat/client.py 
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 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
class ToolClient:
    """Client for interacting with Agent Tools Deployments.

    This class provides methods to call the custom model tool using various hooks:
    `score`, `score_unstructured`, and `chat`. When the `authorization_context` is set,
    the client automatically propagates it to the agent tool. The `authorization_context`
    is required for retrieving access tokens to connect to external services.
    """

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        authorization_context: dict[str, Any] | None = None,
    ):
        """Initialize the ToolClient.

        Args:
            api_key (str | None): API key for authentication. Defaults to the
                environment variable `DATAROBOT_API_TOKEN`.
            base_url (str | None): Base URL for the DataRobot API. Defaults to the
                environment variable `DATAROBOT_ENDPOINT` or 'app.datarobot.com'.
            authorization_context (dict[str, Any] | None): Authorization context to use
                for tool calls. If None, will attempt to get from ContextVar (for backward
                compatibility).
        """
        self.api_key = api_key or os.getenv("DATAROBOT_API_TOKEN")
        base_url = base_url or os.getenv("DATAROBOT_ENDPOINT") or "https://app.datarobot.com"
        base_url = get_api_base(base_url, deployment_id=None)
        self.base_url = base_url
        self._authorization_context = authorization_context

    @property
    def datarobot_api_endpoint(self) -> str:
        return self.base_url + "api/v2"

    def get_deployment(self, deployment_id: str) -> dr.Deployment:
        """Retrieve a deployment by its ID.

        Args:
            deployment_id (str): The ID of the deployment.

        Returns
        -------
            dr.Deployment: The deployment object.
        """
        dr.Client(self.api_key, self.datarobot_api_endpoint)
        return dr.Deployment.get(deployment_id=deployment_id)

    def call(
        self,
        deployment_id: str,
        payload: dict[str, Any],
        authorization_context: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> UnstructuredPredictionResult:
        """Run the custom model tool using score_unstructured hook.

        Args:
            deployment_id (str): The ID of the deployment.
            payload (dict[str, Any]): The input payload.
            authorization_context (dict[str, Any] | None): Authorization context to use.
                If None, uses the context from initialization or falls back to ContextVar.
            **kwargs: Additional keyword arguments.

        Returns
        -------
            UnstructuredPredictionResult: The response content and headers.
        """
        # Use explicit context, fall back to instance context, then ContextVar
        auth_ctx = authorization_context or self._authorization_context
        if auth_ctx is None:
            try:
                auth_ctx = get_authorization_context()
            except LookupError:
                auth_ctx = {}

        data = {
            "payload": payload,
            "authorization_context": auth_ctx,
        }
        return predict_unstructured(
            deployment=self.get_deployment(deployment_id),
            data=json.dumps(data),
            content_type="application/json",
            **kwargs,
        )

    def score(
        self, deployment_id: str, data_frame: pd.DataFrame, **kwargs: Any
    ) -> PredictionResult:
        """Run the custom model tool using score hook.

        Args:
            deployment_id (str): The ID of the deployment.
            data_frame (pd.DataFrame): The input data frame.
            **kwargs: Additional keyword arguments.

        Returns
        -------
            PredictionResult: The response content and headers.
        """
        return predict(
            deployment=self.get_deployment(deployment_id),
            data_frame=data_frame,
            **kwargs,
        )

    def chat(
        self,
        completion_create_params: CompletionCreateParams,
        model: str,
        authorization_context: dict[str, Any] | None = None,
    ) -> ChatCompletion | Iterator[ChatCompletionChunk]:
        """Run the custom model tool with the chat hook.

        Args:
            completion_create_params (CompletionCreateParams): Parameters for the chat completion.
            model (str): The model to use.
            authorization_context (dict[str, Any] | None): Authorization context to use.
                If None, uses the context from initialization or falls back to ContextVar.

        Returns
        -------
            Union[ChatCompletion, Iterator[ChatCompletionChunk]]: The chat completion response.
        """
        # Use explicit context, fall back to instance context, then ContextVar
        auth_ctx = authorization_context or self._authorization_context
        if auth_ctx is None:
            try:
                auth_ctx = get_authorization_context()
            except LookupError:
                auth_ctx = {}

        extra_body = {
            "authorization_context": auth_ctx,
        }
        return openai.chat.completions.create(
            **completion_create_params,
            model=model,
            extra_body=extra_body,
        )

__init__ 

__init__(api_key: str | None = None, base_url: str | None = None, authorization_context: dict[str, Any] | None = None)

Initialize the ToolClient.

Parameters:

Name Type Description Default
api_key str | None

API key for authentication. Defaults to the environment variable DATAROBOT_API_TOKEN.

 None
base_url str | None

Base URL for the DataRobot API. Defaults to the environment variable DATAROBOT_ENDPOINT or 'app.datarobot.com'.

 None
authorization_context dict[str, Any] | None

Authorization context to use for tool calls. If None, will attempt to get from ContextVar (for backward compatibility).

 None
Source code in datarobot_genai/core/chat/client.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __init__(
    self,
    api_key: str | None = None,
    base_url: str | None = None,
    authorization_context: dict[str, Any] | None = None,
):
    """Initialize the ToolClient.

    Args:
        api_key (str | None): API key for authentication. Defaults to the
            environment variable `DATAROBOT_API_TOKEN`.
        base_url (str | None): Base URL for the DataRobot API. Defaults to the
            environment variable `DATAROBOT_ENDPOINT` or 'app.datarobot.com'.
        authorization_context (dict[str, Any] | None): Authorization context to use
            for tool calls. If None, will attempt to get from ContextVar (for backward
            compatibility).
    """
    self.api_key = api_key or os.getenv("DATAROBOT_API_TOKEN")
    base_url = base_url or os.getenv("DATAROBOT_ENDPOINT") or "https://app.datarobot.com"
    base_url = get_api_base(base_url, deployment_id=None)
    self.base_url = base_url
    self._authorization_context = authorization_context

get_deployment 

get_deployment(deployment_id: str) -> dr.Deployment

Retrieve a deployment by its ID.

Parameters:

Name Type Description Default
deployment_id str

The ID of the deployment.

 required
Returns
dr.Deployment: The deployment object.
Source code in datarobot_genai/core/chat/client.py 
73
74
75
76
77
78
79
80
81
82
83
84
def get_deployment(self, deployment_id: str) -> dr.Deployment:
    """Retrieve a deployment by its ID.

    Args:
        deployment_id (str): The ID of the deployment.

    Returns
    -------
        dr.Deployment: The deployment object.
    """
    dr.Client(self.api_key, self.datarobot_api_endpoint)
    return dr.Deployment.get(deployment_id=deployment_id)

call 

call(deployment_id: str, payload: dict[str, Any], authorization_context: dict[str, Any] | None = None, **kwargs: Any) -> UnstructuredPredictionResult

Run the custom model tool using score_unstructured hook.

Parameters:

Name Type Description Default
deployment_id str

The ID of the deployment.

 required
payload dict[str, Any]

The input payload.

 required
authorization_context dict[str, Any] | None

Authorization context to use. If None, uses the context from initialization or falls back to ContextVar.

 None
**kwargs Any

Additional keyword arguments.

 {}
Returns
UnstructuredPredictionResult: The response content and headers.
Source code in datarobot_genai/core/chat/client.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
123
def call(
    self,
    deployment_id: str,
    payload: dict[str, Any],
    authorization_context: dict[str, Any] | None = None,
    **kwargs: Any,
) -> UnstructuredPredictionResult:
    """Run the custom model tool using score_unstructured hook.

    Args:
        deployment_id (str): The ID of the deployment.
        payload (dict[str, Any]): The input payload.
        authorization_context (dict[str, Any] | None): Authorization context to use.
            If None, uses the context from initialization or falls back to ContextVar.
        **kwargs: Additional keyword arguments.

    Returns
    -------
        UnstructuredPredictionResult: The response content and headers.
    """
    # Use explicit context, fall back to instance context, then ContextVar
    auth_ctx = authorization_context or self._authorization_context
    if auth_ctx is None:
        try:
            auth_ctx = get_authorization_context()
        except LookupError:
            auth_ctx = {}

    data = {
        "payload": payload,
        "authorization_context": auth_ctx,
    }
    return predict_unstructured(
        deployment=self.get_deployment(deployment_id),
        data=json.dumps(data),
        content_type="application/json",
        **kwargs,
    )

score 

score(deployment_id: str, data_frame: DataFrame, **kwargs: Any) -> PredictionResult

Run the custom model tool using score hook.

Parameters:

Name Type Description Default
deployment_id str

The ID of the deployment.

 required
data_frame DataFrame

The input data frame.

 required
**kwargs Any

Additional keyword arguments.

 {}
Returns
PredictionResult: The response content and headers.
Source code in datarobot_genai/core/chat/client.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
def score(
    self, deployment_id: str, data_frame: pd.DataFrame, **kwargs: Any
) -> PredictionResult:
    """Run the custom model tool using score hook.

    Args:
        deployment_id (str): The ID of the deployment.
        data_frame (pd.DataFrame): The input data frame.
        **kwargs: Additional keyword arguments.

    Returns
    -------
        PredictionResult: The response content and headers.
    """
    return predict(
        deployment=self.get_deployment(deployment_id),
        data_frame=data_frame,
        **kwargs,
    )

chat 

chat(completion_create_params: CompletionCreateParams, model: str, authorization_context: dict[str, Any] | None = None) -> ChatCompletion | Iterator[ChatCompletionChunk]

Run the custom model tool with the chat hook.

Parameters:

Name Type Description Default
completion_create_params CompletionCreateParams

Parameters for the chat completion.

 required
model str

The model to use.

 required
authorization_context dict[str, Any] | None

Authorization context to use. If None, uses the context from initialization or falls back to ContextVar.

 None
Returns
Union[ChatCompletion, Iterator[ChatCompletionChunk]]: The chat completion response.
Source code in datarobot_genai/core/chat/client.py 
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
def chat(
    self,
    completion_create_params: CompletionCreateParams,
    model: str,
    authorization_context: dict[str, Any] | None = None,
) -> ChatCompletion | Iterator[ChatCompletionChunk]:
    """Run the custom model tool with the chat hook.

    Args:
        completion_create_params (CompletionCreateParams): Parameters for the chat completion.
        model (str): The model to use.
        authorization_context (dict[str, Any] | None): Authorization context to use.
            If None, uses the context from initialization or falls back to ContextVar.

    Returns
    -------
        Union[ChatCompletion, Iterator[ChatCompletionChunk]]: The chat completion response.
    """
    # Use explicit context, fall back to instance context, then ContextVar
    auth_ctx = authorization_context or self._authorization_context
    if auth_ctx is None:
        try:
            auth_ctx = get_authorization_context()
        except LookupError:
            auth_ctx = {}

    extra_body = {
        "authorization_context": auth_ctx,
    }
    return openai.chat.completions.create(
        **completion_create_params,
        model=model,
        extra_body=extra_body,
    )