class DataRobotClientWithAsyncAPI:
"""A async wrapper of DR RESTFul APIs.
This client was designed for internal use only (Some methods are created as private methods).
It is the interim solution before async support is added in
https://github.com/datarobot/public_api_client.
"""
def __init__(self, dr_host: str) -> None:
self._dr_host = dr_host
self._session = get_async_https_session()
self._retry_client = get_async_https_retry_client(self._session)
async def clean_up(self) -> None:
await self._retry_client.close()
async def __aenter__(self) -> "DataRobotClientWithAsyncAPI":
return self
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> None:
await self.clean_up()
@staticmethod
def get_api_v2_endpoint(host_url: str, v2_path: str) -> str:
parsed_url = urlparse(host_url)
full_path = f"api/v2/{v2_path}/"
full_path = full_path.replace("//", "/")
new_parsed_url = parsed_url._replace(path=full_path)
return urlunparse(new_parsed_url)
async def _unpaginate(
self,
initial_url: str,
headers: dict[str, str],
) -> AsyncIterator[Any]:
url: str | None = initial_url
while url is not None:
async with self._retry_client.get(url, headers=headers) as resp:
resp.raise_for_status()
resp_data = await resp.json()
for item in resp_data.get("data", []):
yield item
url = resp_data.get("next")
async def _get_feature_entitlement_evaluate_result(
self,
feature_flag_name: str,
dr_bearer_token: str,
) -> dict[str, Any]:
url = self.get_api_v2_endpoint(self._dr_host, "/entitlements/evaluate/")
async with self._retry_client.post(
url,
json={"entitlements": [{"name": feature_flag_name}]},
headers={"Authorization": f"Bearer {dr_bearer_token}"},
) as resp:
resp.raise_for_status()
return await resp.json()
async def _is_feature_flag_enabled(self, feature_flag_name: str, dr_bearer_token: str) -> bool:
response = await self._get_feature_entitlement_evaluate_result(
feature_flag_name, dr_bearer_token
)
feature_flag_info = response["entitlements"][0]
return bool(feature_flag_info["value"])
async def _list_mcp_deployment_ids(self, dr_bearer_token: str) -> list[str]:
url = self.get_api_v2_endpoint(self._dr_host, "/deployments/")
headers = {"Authorization": f"Bearer {dr_bearer_token}"}
ids: list[str] = []
async for deployment in self._unpaginate(url, headers):
model = deployment.get("model") or {}
if model.get("targetType") == "MCP":
ids.append(deployment["id"])
return ids
async def _list_mcp_tool_custom_model_deployment_ids(self, dr_bearer_token: str) -> list[str]:
url = self.get_api_v2_endpoint(self._dr_host, "/deployments")
url = f"{url}?tagValues=tool&tagKeys=tool"
headers = {"Authorization": f"Bearer {dr_bearer_token}"}
ids: list[str] = []
async for deployment in self._unpaginate(url, headers):
for tag in deployment.get("tags", []):
if tag.get("name") == "tool" and tag.get("value") == "tool":
ids.append(deployment["id"])
return ids
async def _get_datarobot_deployment(
self, deployment_id: str, dr_bearer_token: str
) -> dr.Deployment:
url = self.get_api_v2_endpoint(self._dr_host, f"/deployments/{deployment_id}/")
headers = {"Authorization": f"Bearer {dr_bearer_token}"}
async with self._retry_client.get(
url,
headers=headers,
) as resp:
resp.raise_for_status()
deployment_json_response = await resp.json()
return dr.Deployment.from_server_data(deployment_json_response)
async def _get_deployment_directaccess_info(
self, deployment_id: str, dr_bearer_token: str
) -> dict[str, Any]:
"""Fetch and normalise the /directAccess/info/ metadata for a deployment.
Applies the same camelCase → snake_case normalisation that the synchronous
DR SDK client applies when the drmcp server fetches this endpoint.
"""
url = self.get_api_v2_endpoint(
self._dr_host, f"/deployments/{deployment_id}/directAccess/info/"
)
headers = {"Authorization": f"Bearer {dr_bearer_token}"}
async with self._retry_client.get(url, headers=headers) as resp:
resp.raise_for_status()
response_json = await resp.json()
data = from_api(response_json)
if isinstance(data, list):
return data[0] if data else {}
if isinstance(data, dict):
return data
return {}
async def _get_deployment_supports_chat_api(
self, deployment_id: str, dr_bearer_token: str
) -> bool:
"""Return whether the deployment advertises chat completions support.
Reads ``supports_chat_api`` from the deployment's ``/capabilities/`` endpoint.
Defaults to ``False`` on any failure or when the flag is absent.
"""
url = self.get_api_v2_endpoint(self._dr_host, f"/deployments/{deployment_id}/capabilities/")
headers = {"Authorization": f"Bearer {dr_bearer_token}"}
try:
async with self._retry_client.get(url, headers=headers) as resp:
resp.raise_for_status()
payload = await resp.json()
capabilities = (payload or {}).get("data") or []
for capability in capabilities:
if not isinstance(capability, dict):
continue
if capability.get("name") == "supports_chat_api":
return bool(capability.get("supported", False))
return False
except Exception as exc:
logger.warning(
"Could not fetch capabilities for deployment %s; "
"assuming chat API not supported: %s",
deployment_id,
exc,
)
return False