cradle.sdk.client

   1import copy
   2import difflib
   3import logging
   4import time
   5import warnings
   6from collections.abc import Iterator
   7from importlib.metadata import version as pkg_version
   8from pathlib import Path
   9from typing import Any, BinaryIO, NoReturn, TypeVar
  10from urllib.parse import urlparse
  11
  12import httpx
  13import jwt
  14import pyarrow as pa
  15from pydantic import BaseModel, TypeAdapter
  16from typing_extensions import Generator, override
  17
  18from cradle.sdk.auth.device import DeviceAuth
  19from cradle.sdk.exceptions import ClientError as ClientError  # for re-export
  20from cradle.sdk.types.common import (
  21    ContextProject,
  22    ContextRound,
  23    ContextWorkspace,
  24    ListOptions,
  25)
  26from cradle.sdk.types.data import (
  27    AddTableRequest,
  28    ArtifactResponse,
  29    BaseTableCreate,
  30    BaseTableUpdate,
  31    DataLoadCreate,
  32    DataLoadResponse,
  33    DataVersionResponse,
  34    FileUploadResponse,
  35    ListArtifactResponse,
  36    ListDataLoadResponse,
  37    ListDataVersionResponse,
  38    ListTableResponse,
  39    QueryDataRequest,
  40    TableArchive,
  41    TableRename,
  42    TableResponse,
  43    ViewTableCreate,
  44    ViewTableUpdate,
  45)
  46from cradle.sdk.types.task import (
  47    ListTaskResponse,
  48    TaskCreate,
  49    TaskResponse,
  50    TaskState,
  51)
  52from cradle.sdk.types.workspace import (
  53    CustomPredictorContainerCreate,
  54    CustomPredictorContainerCreateResponse,
  55    CustomPredictorContainerUpdate,
  56    CustomPredictorCreateResponse,
  57    CustomPredictorResponse,
  58    ListCustomPredictorResponse,
  59    ListProjectResponse,
  60    ListRoundResponse,
  61    ProjectCreate,
  62    ProjectResponse,
  63    RoundCreate,
  64    RoundResponse,
  65    WorkspaceResponse,
  66)
  67
  68logger = logging.getLogger(__name__)
  69
  70T = TypeVar("T", bound=BaseModel)
  71
  72
  73class HttpClient:
  74    def __init__(
  75        self,
  76        prefix: str,
  77        base_url: str,
  78        auth: httpx.Auth | None = None,
  79        timeout: float = 60.0,
  80        user_agent: str | None = None,
  81        workspace: str | None = None,
  82        use_keyring: bool = True,
  83    ):
  84        base_url = base_url.rstrip("/")
  85        user_agent = user_agent or f"cradle-sdk-python/{pkg_version('cradle-sdk')}"
  86        headers = {
  87            "Accept": "application/json",
  88            "User-Agent": user_agent,
  89        }
  90        client = httpx.Client(auth=auth, headers=headers, timeout=timeout, follow_redirects=True)
  91        if auth is None:
  92            client.auth = DeviceAuth.from_strategy(
  93                client=client, workspace=workspace, base_url=base_url, use_keyring=use_keyring
  94            )
  95
  96        self.http_client = client
  97        self.prefix = prefix.strip("/")
  98        self.base_url = base_url.rstrip("/")
  99        self.workspace = workspace
 100
 101    def with_prefix(self, prefix: str) -> "HttpClient":
 102        c = copy.copy(self)
 103        c.prefix = f"{self.prefix}/{prefix.strip('/')}".strip("/")
 104        return c
 105
 106    def url(self, path: str) -> str:
 107        """Generate a URL for a given path.
 108
 109        This follows slightly weird rules, specific to how our API is designed.
 110
 111        The overall goal here is to make the URL generation loose with what it accepts, to never
 112        be confused by whether or not there are unexpected slashes in any part of the url. It's a very
 113        "do what I mean" approach.
 114
 115        - "path" is never interpreted as an absolute path, it is always relative to the prefix
 116        - if the path starts with ":", it is a verb endpoint and should be appended to the url + prefix, not assumed to be a subpath
 117        - everything else is assumed to be a subpath and should be appended to the url + prefix
 118        """
 119
 120        combined_path = f"{self.prefix}{path}"
 121        if not path.startswith(":") and not path.startswith("/"):
 122            combined_path = f"{self.prefix}/{path}"
 123
 124        # This special case is needed in case we have a ":verb" endpoint at the root of the API
 125        # like "https://cradle.bio/:list" - otherwise the logic below would generate
 126        # "https://cradle.bio:list" which is wrong.
 127        parsed_url = urlparse(self.base_url)
 128        base_url_has_path = bool(parsed_url.path and parsed_url.path not in {"", "/"})
 129        if combined_path.startswith(":") and base_url_has_path:
 130            return f"{self.base_url}{combined_path}"
 131
 132        return f"{self.base_url}/{combined_path.lstrip('/')}"
 133
 134    def request(
 135        self,
 136        method: str,
 137        path: str,
 138        json: dict[str, Any] | BaseModel | None = None,
 139        params: dict[str, Any] | None = None,
 140        **kwargs,
 141    ) -> httpx.Response:
 142        if params is not None:
 143            # HTTPX will send "...&foo=&...", which FastAPI would interpret as an empty string rather than None.
 144            kwargs["params"] = {k: v for k, v in params.items() if v is not None}
 145        if isinstance(json, BaseModel):
 146            json = json.model_dump(mode="json", by_alias=True)
 147        if json is not None:
 148            kwargs["json"] = json
 149
 150        return self.http_client.request(method, self.url(path), **kwargs)
 151
 152    def get(self, path: str, response_type: type[T] | Any | None, params: dict[str, Any] | None = None, **kwargs) -> T:
 153        return self._handle_response(self.request("GET", path, params=params, **kwargs), response_type)
 154
 155    def post(self, path: str, response_type: type[T] | Any | None, params: dict[str, Any] | None = None, **kwargs) -> T:
 156        return self._handle_response(self.request("POST", path, params=params, **kwargs), response_type)
 157
 158    def _handle_response(
 159        self, response: httpx.Response, response_type: type[T] | Iterator[str] | Iterator[bytes] | None
 160    ) -> Any:
 161        if not response.is_success:
 162            self.handle_error_response(response)
 163        if response_type is None:
 164            return None
 165
 166        ct = response.headers.get("content-type")
 167        if ct is None:
 168            raise ValueError("content type header missing")
 169
 170        ct = ct.split(";")[0].strip()
 171        if ct == "application/json":
 172            return TypeAdapter(response_type).validate_python(response.json())
 173        if ct == "application/x-ndjson" and response_type == Iterator[str]:
 174            return response.iter_lines()
 175        if ct == "application/octet-stream" and response_type == Iterator[bytes]:
 176            return response.iter_bytes()
 177        raise ValueError(f"Unsupported content type {ct}")
 178
 179    def handle_error_response(self, response: httpx.Response) -> NoReturn:
 180        auth = self.http_client.auth
 181        if (
 182            response.status_code == 403
 183            and self.workspace is not None
 184            and isinstance(auth, DeviceAuth)
 185            and auth.authorized_workspace_id != self.workspace
 186        ):
 187            auth.suggest_logout(self.workspace)
 188        try:
 189            body = response.json()
 190            error_msg = body.get("detail", response.text)
 191            if (error_id := body.get("error_id")) is not None:
 192                error_msg = f"{error_msg} (error ID: {error_id})"
 193            errors = body.get("errors", [])
 194        except Exception:  # noqa: BLE001
 195            error_msg = response.text
 196            errors = []
 197        raise ClientError(response.status_code, error_msg, errors=errors)
 198
 199
 200class DisabledAuth(httpx.Auth):
 201    """Auth implementation to use when authentication should be disabled"""
 202
 203    def __init__(self, workspace: str | None = None):
 204        self._workspace_name = workspace
 205
 206    @override
 207    def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]:
 208        claims = {
 209            "urn:cradle:email": "john-doe@cradle.bio",
 210            "urn:cradle:user_id": "john-doe",
 211            "urn:cradle:full_name": "John Doe",
 212            "urn:cradle:workspace_id": self._workspace_name,
 213        }
 214
 215        token = jwt.encode(claims, key="this_key_is_not_verified_by_the_backend")
 216
 217        request.headers["Authorization"] = f"Bearer ca_{token}"
 218        yield request
 219
 220
 221API_URL = "https://api.cradle.bio"
 222
 223
 224class Client:
 225    """The workspace API provides functionality for retrieving information about your workspace.
 226
 227    A workspace is the home for all data, projects and tasks for an organization. By
 228    design no information can be shared across workspaces, they are the top-level
 229    concept in the Cradle API.
 230
 231    Workspace creation and management is only available to Cradle administrators.
 232    """
 233
 234    def __init__(
 235        self,
 236        workspace: str,
 237        auth: httpx.Auth | None = None,
 238        base_url: str = API_URL,
 239        timeout: int = 60,
 240        user_agent: str | None = None,
 241        use_keyring: bool = True,
 242    ):
 243        self._base_client = HttpClient(
 244            prefix="/",
 245            base_url=base_url,
 246            auth=auth,
 247            timeout=timeout,
 248            user_agent=user_agent,
 249            workspace=workspace,
 250            use_keyring=use_keyring,
 251        )
 252        self._client = self._base_client.with_prefix(f"/v2/workspace/{workspace}/")
 253        self._workspace = workspace
 254
 255    def workspace(self) -> WorkspaceResponse:
 256        return self._base_client.get("/v2/workspace:get", WorkspaceResponse, dict(name=self._workspace))
 257
 258    @property
 259    def project(self) -> "ProjectClient":
 260        return ProjectClient(self._client)
 261
 262    @property
 263    def round(self) -> "RoundClient":
 264        return RoundClient(self._client)
 265
 266    @property
 267    def data(self) -> "DataClient":
 268        return DataClient(self._client)
 269
 270    @property
 271    def task(self) -> "TaskClient":
 272        return TaskClient(self._client)
 273
 274    @property
 275    def custom_predictor(self) -> "CustomPredictorClient":
 276        return CustomPredictorClient(self._client)
 277
 278
 279class ProjectClient:
 280    """The projects API provides functionality for managing and interacting with
 281    projects in a workspace.
 282
 283    Projects organize tasks in a workspace and represent the work involved in
 284    optimizing a specific protein or achieving a specific goal. Data is currently
 285    isolated to a specific project, but in the future we will support cross project
 286    data sharing within a workspace.
 287    """
 288
 289    def __init__(self, client: HttpClient):
 290        self._client = client.with_prefix("/project")
 291
 292    def get(self, project_id: int) -> ProjectResponse:
 293        """Get a project by `ID`."""
 294        return self._client.get(":get", ProjectResponse, dict(id=project_id))
 295
 296    def list(self) -> Generator[ProjectResponse, None, None]:
 297        """List all projects in the workspace."""
 298        opts = ListOptions()
 299        while True:
 300            resp = self._client.get(":list", ListProjectResponse, params=opts.model_dump())
 301            yield from resp.items
 302            if resp.cursor is None:
 303                return
 304            opts.cursor = resp.cursor
 305
 306    def get_by_name(self, name: str) -> ProjectResponse:
 307        """Get a project by name. Raises ValueError if not found."""
 308        for project in self.list():
 309            if project.name == name:
 310                return project
 311        raise ValueError(f"Project with name '{name}' not found")
 312
 313    def create(self, project: ProjectCreate) -> ProjectResponse:
 314        """Create a new project in the workspace."""
 315        return self._client.post(":create", ProjectResponse, json=project)
 316
 317    def update(self, project_id: int, project: ProjectCreate) -> ProjectResponse:
 318        """Update a project's information."""
 319        return self._client.post(":update", ProjectResponse, dict(id=project_id), json=project)
 320
 321    def archive(self, project_id: int) -> ProjectResponse:
 322        """Archive a project.
 323
 324        The data from the project will still be available.
 325        It will not be possible to create new tasks, data loads, etc. in the project.
 326        """
 327        return self._client.post(":archive", ProjectResponse, dict(id=project_id))
 328
 329    def unarchive(self, project_id: int) -> ProjectResponse:
 330        """Unarchive a previously archived project.
 331
 332        Returns an HTTP 422 `Unprocessable Entity` error if the project is not archived.
 333        """
 334        return self._client.post(":unarchive", ProjectResponse, dict(id=project_id))
 335
 336
 337class RoundClient:
 338    """The rounds API provides functionality for managing and interacting with rounds within a project.
 339
 340    Rounds are a way to organize the experiments within a project. Commonly, a round represents all
 341    in-vitro experiments related to a specific sub-goal or timeframe of a project.
 342    While rounds are typically ordered sequentially, multiple rounds can exist in parallel (for example, when
 343    testing multiple design methods in parallel).
 344
 345    Tasks, data loads, and reports can be (but don't have to be!) assigned to a round.
 346    """
 347
 348    def __init__(self, client: HttpClient):
 349        self._client = client.with_prefix("/round")
 350
 351    def get(self, round_id: int) -> RoundResponse:
 352        """Get a round by ID."""
 353        return self._client.get(":get", RoundResponse, dict(id=round_id))
 354
 355    def list(self, project_id: int | None = None) -> Generator[RoundResponse, None, None]:
 356        """List all rounds in the project."""
 357        params = {}
 358        if project_id is not None:
 359            params["project_id"] = project_id
 360
 361        opts = ListOptions()
 362        while True:
 363            resp = self._client.get(
 364                ":list",
 365                ListRoundResponse,
 366                params={**opts.model_dump(), **params},
 367            )
 368            yield from resp.items
 369            if resp.cursor is None:
 370                return
 371            opts.cursor = resp.cursor
 372
 373    def get_by_name(self, project_id: int, name: str) -> RoundResponse:
 374        """Get a round by name within a specific project. Raises ValueError if not found."""
 375        for round_ in self.list(project_id=project_id):
 376            if round_.name == name:
 377                return round_
 378        raise ValueError(f"Round with name '{name}' not found in project {project_id}")
 379
 380    def create(self, round_: RoundCreate) -> RoundResponse:
 381        """Create a new round in a project."""
 382        return self._client.post(":create", RoundResponse, json=round_)
 383
 384    def update(self, round_id: int, round_: RoundCreate) -> RoundResponse:
 385        """Update a round's name and description."""
 386        return self._client.post(":update", RoundResponse, dict(id=round_id), json=round_)
 387
 388    def archive(self, round_id: int) -> RoundResponse:
 389        """Archive a round.
 390
 391        The data from the round will still be available.
 392        It will not be possible to create new tasks, data loads, etc. in the round.
 393        """
 394        return self._client.post(":archive", RoundResponse, dict(id=round_id))
 395
 396    def unarchive(self, round_id: int) -> RoundResponse:
 397        """Unarchive a previously archived round.
 398
 399        Returns an HTTP 422 `Unprocessable Entity` error if the round is not archived.
 400        """
 401        return self._client.post(":unarchive", RoundResponse, dict(id=round_id))
 402
 403
 404class DataClient:
 405    """The data API provides functionality for managing and interacting with data in a workspace.
 406
 407    Data comes in the form of *tables* and *artifacts*.
 408    *Tables* are used to store row-based data, such as measurements, sequences, and other structured data.
 409    *Artifacts* are used to store the results of tasks, such as machine learning models, predictions, and other outputs.
 410
 411    While tables can be created and uploaded by users, artifacts are only produced by tasks and cannot be uploaded directly.
 412
 413    Artifacts and table rows (via their data load source) have a context specified. This context refers to the
 414    origin of the data and is used in other parts of the API to determine data visibility, e.g. to
 415    isolate data between projects.
 416
 417    See the [Artifacts](#tag/Artifacts) API endpoints for more details on artifact data.
 418
 419    See the [Tables](#tag/Tables) API endpoints for more details on tabular data.
 420    """
 421
 422    def __init__(self, client: HttpClient):
 423        self._client = client.with_prefix("/data")
 424
 425    @property
 426    def artifact(self) -> "ArtifactClient":
 427        return ArtifactClient(self._client)
 428
 429    @property
 430    def table(self) -> "TableClient":
 431        return TableClient(self._client)
 432
 433    @property
 434    def load(self) -> "DataLoadClient":
 435        return DataLoadClient(self._client)
 436
 437    def list_versions(
 438        self,
 439        project_id: int | None = None,
 440        table_id: int | None = None,
 441    ) -> Generator[DataVersionResponse, None, None]:
 442        warnings.warn(
 443            "DataClient.list_versions() is deprecated. Use snapshots instead to work with older state of data.",
 444            category=DeprecationWarning,
 445            stacklevel=2,
 446        )
 447
 448        params = {}
 449        if project_id is not None:
 450            params["project_id"] = project_id
 451        if table_id is not None:
 452            params["table_id"] = table_id
 453
 454        opts = ListOptions()
 455        while True:
 456            resp = self._client.get("/version:list", ListDataVersionResponse, params={**params, **opts.model_dump()})
 457            yield from resp.items
 458            if resp.cursor is None:
 459                return
 460            opts.cursor = resp.cursor
 461
 462
 463class ArtifactClient:
 464    """Artifacts are produced by tasks and can be referenced as inputs to other tasks. They can generally
 465    not be uploaded and downloaded directly. They have a well-defined type that determines for which
 466    purpose they can be used.
 467    """
 468
 469    def __init__(self, client: HttpClient):
 470        self._client = client.with_prefix("/artifact")
 471
 472    def get(self, artifact_id: int) -> ArtifactResponse:
 473        """Get information about an artifact."""
 474        return self._client.get(":get", ArtifactResponse, dict(id=artifact_id))
 475
 476    def list(
 477        self,
 478        project_id: int | None = None,
 479        round_id: int | None = None,
 480    ) -> Generator[ArtifactResponse, None, None]:
 481        """List artifacts in the workspace."""
 482        opts = ListOptions()
 483
 484        params = {}
 485        if project_id is not None:
 486            params["project_id"] = project_id
 487        if round_id is not None:
 488            params["round_id"] = round_id
 489
 490        while True:
 491            resp = self._client.get(
 492                ":list",
 493                ListArtifactResponse,
 494                params={**opts.model_dump(), **params},
 495            )
 496            yield from resp.items
 497            if resp.cursor is None:
 498                return
 499            opts.cursor = resp.cursor
 500
 501
 502class _IteratorReader:
 503    def __init__(self, iterator: Iterator[bytes]):
 504        self.iterator = iterator
 505        self.buffer = b""
 506        self._closed = False
 507
 508    @property
 509    def closed(self) -> bool:
 510        return self._closed
 511
 512    def close(self) -> None:
 513        self._closed = True
 514        self._buffer = b""
 515
 516    def read(self, size=-1):
 517        try:
 518            while size < 0 or len(self.buffer) < size:
 519                self.buffer += next(self.iterator)
 520        except StopIteration:
 521            pass
 522        if size < 0:
 523            out, self.buffer = self.buffer, b""
 524        else:
 525            out, self.buffer = self.buffer[:size], self.buffer[size:]
 526        return out
 527
 528    def __enter__(self):
 529        return self
 530
 531    def __exit__(self, exc_type, exc_val, exc_tb):
 532        self.close()
 533
 534
 535class TableClient:
 536    """Tables store row-based data with a well-defined column schema. Such data may come from spreadsheets
 537    or other tabular data stores.
 538
 539    Each table is created with a well-defined column schema, which can include primitive columns such
 540    as numbers and strings, as well as nested struct columns and arrays.
 541    A table's schema can be extended but existing columns cannot generally be modified or removed.
 542
 543    Table data as well as it's column schema itself are version controlled. Each change to the contained
 544    rows is tracked as a data version.
 545    Tables can be queried at any historic version unless some destructive change has been made, such as
 546    hard deletion of historic data for compliance reasons.
 547
 548    Because tables are versioned, each row is immutable. This means that it is not possible to change individual
 549    values in a row. Instead, the data load that added the row can be undone and the rows can be loaded
 550    again with the changed data.
 551
 552    See the [Data Loads](#tag/Data-Loads) API endpoints on how to add and remove data to tables.
 553    """
 554
 555    def __init__(self, client: HttpClient):
 556        self._client = client.with_prefix("/table")
 557
 558    def query(
 559        self,
 560        query: str,
 561        *,
 562        project_id: int | None,
 563        version_id: int | None = None,
 564    ) -> Generator[pa.RecordBatch, None, None]:
 565        """Query data with the given SQL query and retrieve the result as arrow record batches.
 566
 567        Args:
 568            query: The SQL query the execute.
 569            project_id: project_id to filter the source data by. Can be set to none to query data
 570                access all projects the authenticated user has access to.
 571            version_id: Data version ID at which to execute the query.
 572
 573        Example:
 574            ```
 575            result = client.query("SELECT * FROM source_table")
 576            table = pyarrow.Table.from_batches(result)
 577
 578            # Write as CSV. Does only work with primitive columns.
 579            pyarrow.csv.write_csv(t, "result.csv")
 580
 581            # Write as Parquet. Does work with structs and array columns as well.
 582            pyarrow.parquet.write_table(t, "result.parquet")
 583            ```
 584        """
 585        request = QueryDataRequest(
 586            query=query,
 587            project_id=project_id,
 588            version_id=version_id,
 589        )
 590        with self._client.http_client.stream(
 591            "POST", self._client.url(":query"), json=request.model_dump(by_alias=True)
 592        ) as response:
 593            if not response.is_success:
 594                response.read()  # Required in streaming request to the error handling can access content.
 595                self._client.handle_error_response(response)
 596
 597            ct = response.headers.get("content-type")
 598            if ct is None:
 599                raise ValueError("content type header missing")
 600            ct = ct.split(";")[0].strip()
 601
 602            if ct != "application/vnd.apache.arrow.stream":
 603                raise ValueError(f"unexpected content type {ct}")
 604            yield from pa.ipc.open_stream(_IteratorReader(response.iter_bytes()))
 605
 606    def get_by_id(self, table_id: int) -> TableResponse:
 607        """Retrieve the table by its ID.
 608        The ID is an opaque identifier for the table returned at table creation time and is distinct from the table reference.
 609        """
 610        return self._client.get(":getById", TableResponse, dict(id=table_id))
 611
 612    def get(self, reference: str) -> TableResponse:
 613        """Get a table by its reference."""
 614        return self._client.get(":get", TableResponse, dict(reference=reference))
 615
 616    def list(self) -> Generator[TableResponse, None, None]:
 617        """List all current base tables and views in the workspace."""
 618        opts = ListOptions()
 619        while True:
 620            resp = self._client.get(":list", ListTableResponse, params=opts.model_dump())
 621            yield from resp.items
 622            if resp.cursor is None:
 623                return
 624            opts.cursor = resp.cursor
 625
 626    def list_versions(self, table_id: int) -> Generator[TableResponse, None, None]:
 627        """List all versions of a specific table.
 628
 629        Deprecated: Table versions will no longer be tracked in future platform versions.
 630        """
 631        warnings.warn(
 632            "TableClient.list_versions() is deprecated. Use snapshots instead to work with older state of data.",
 633            category=DeprecationWarning,
 634            stacklevel=2,
 635        )
 636
 637        opts = ListOptions()
 638        while True:
 639            resp = self._client.get(
 640                ":listVersions",
 641                ListTableResponse,
 642                params=dict(**opts.model_dump(), id=table_id),
 643            )
 644            yield from resp.items
 645            if resp.cursor is None:
 646                return
 647            opts.cursor = resp.cursor
 648
 649    def list_archived(self, reference: str) -> Generator[TableResponse, None, None]:
 650        """List all archived tables that match the given reference."""
 651        opts = ListOptions()
 652        while True:
 653            resp = self._client.get(
 654                ":listArchived",
 655                ListTableResponse,
 656                params=dict(**opts.model_dump(), reference=reference),
 657            )
 658            yield from resp.items
 659            if resp.cursor is None:
 660                return
 661            opts.cursor = resp.cursor
 662
 663    def create(self, table: ViewTableCreate | BaseTableCreate) -> TableResponse:
 664        """Create a new table or view."""
 665        return self._client.post(":create", TableResponse, json=table)
 666
 667    def update(self, reference: str, table: ViewTableUpdate | BaseTableUpdate) -> TableResponse:
 668        """Update a table or view.
 669
 670        Tables can only be updated in forward-compatible ways.
 671        New columns can be added but not be removed.
 672        Existing columns can be turned from non-nullable to nullable but not the other way around.
 673        Existing columns cannot be renamed or changed in type.
 674
 675        Views can be updated arbitrarily.
 676        """
 677        # Server differentiates between assembly_format being None or unset,
 678        # so make sure to propagate it correctly.
 679        body = table.model_dump(mode="json", by_alias=True)
 680        if "assembly_format" not in table.model_fields_set:
 681            body.pop("assembly_format")
 682
 683        return self._client.post(":update", TableResponse, params=dict(reference=reference), json=body)
 684
 685    def archive(self, reference: str, archive: TableArchive | None = None) -> TableResponse:
 686        """Archive a table or view.
 687
 688        Archiving makes the table or view inaccessible for regular operations
 689        but preserves it for historical purposes.
 690        """
 691        if not archive:
 692            archive = TableArchive()
 693        return self._client.post(":archive", TableResponse, params=dict(reference=reference), json=archive)
 694
 695    def unarchive(self, table_id: int) -> TableResponse:
 696        """Unarchive a previously archived table or view.
 697
 698        This makes the table or view available for regular operations again.
 699        If an active table with the same reference already exists, the unarchive operation
 700        will fail with an HTTP 409 `Conflict error`.
 701        """
 702        return self._client.post(":unarchive", TableResponse, params=dict(id=table_id))
 703
 704    def rename(self, reference: str, rename: TableRename) -> TableResponse:
 705        """Rename the table."""
 706        return self._client.post(":rename", TableResponse, params=dict(reference=reference), json=rename)
 707
 708
 709class DataLoadClient:
 710    """The data API provides functionality for managing and interacting with data in a workspace.
 711
 712    Data comes in the form of *tables* and *artifacts*.
 713    *Tables* are used to store row-based data, such as measurements, sequences, and other structured data.
 714    *Artifacts* are used to store the results of tasks, such as machine learning models, predictions, and other outputs.
 715
 716    While tables can be created and uploaded by users, artifacts are only produced by tasks and cannot be uploaded directly.
 717
 718    Artifacts and table rows (via their data load source) have a context specified. This context refers to the
 719    origin of the data and is used in other parts of the API to determine data visibility, e.g. to
 720    isolate data between projects.
 721
 722    See the [Artifacts](#tag/Artifacts) API endpoints for more details on artifact data.
 723
 724    See the [Tables](#tag/Tables) API endpoints for more details on tabular data.
 725    """
 726
 727    def __init__(self, client: HttpClient):
 728        self._client = client.with_prefix("/load")
 729
 730    def list(
 731        self, project_id: int | None = None, round_id: int | None = None
 732    ) -> Generator[DataLoadResponse, None, None]:
 733        """List all data loads in the workspace."""
 734        params = {}
 735        if project_id is not None:
 736            params["project_id"] = project_id
 737        if round_id is not None:
 738            params["round_id"] = round_id
 739
 740        opts = ListOptions()
 741        while True:
 742            resp = self._client.get(":list", ListDataLoadResponse, params={**params, **opts.model_dump()})
 743            yield from resp.items
 744            if resp.cursor is None:
 745                return
 746            opts.cursor = resp.cursor
 747
 748    def get(self, load_id: int) -> DataLoadResponse:
 749        """Get a data load by ID."""
 750        return self._client.get(":get", DataLoadResponse, dict(id=load_id))
 751
 752    def create(self, load: DataLoadCreate) -> DataLoadResponse:
 753        """Create a new data load.
 754        After creation, the data load will be empty. Files can subsequently be added using the `uploadFile` endpoint.
 755        """
 756        return self._client.post(":create", DataLoadResponse, json=load)
 757
 758    def upload_file(
 759        self,
 760        load_id: int,
 761        file: Path,
 762        *,
 763        description: str | None = None,
 764        filepath: str | Path | None = None,
 765        table_reference: str | None = None,
 766        source_file_id: int | None = None,
 767    ) -> FileUploadResponse:
 768        """Upload a file to the cloud storage bucket and add upload information to load."""
 769        data = {}
 770        if description is not None:
 771            data["description"] = description
 772        if filepath is not None:
 773            data["filepath"] = str(filepath)
 774        if table_reference is not None:
 775            data["table_reference"] = table_reference
 776        if source_file_id is not None:
 777            data["source_file_id"] = source_file_id
 778
 779        with file.open("rb") as f:
 780            return self._client.post(
 781                ":uploadFile",
 782                FileUploadResponse,
 783                params=dict(id=load_id),
 784                files=dict(file=(file.name, f, "application/octet-stream")),
 785                data=data,
 786            )
 787
 788    def download_file(self, file_id: str, buffer: BinaryIO):
 789        """Download a previously uploaded file from a load."""
 790        bytes_stream: Iterator[bytes] = self._client.get(":downloadFile", Iterator[bytes], dict(file_id=file_id))
 791        buffer.writelines(bytes_stream)
 792
 793    def finalize(self, load_id: int) -> DataLoadResponse:
 794        """Finalize the data load and start ingesting rows from the uploaded files.
 795
 796        The ingestion process will run in the background. The progress can be tracked by polling the data load
 797        and waiting for the status to switch from `LOADING` to `COMPLETED` or `FAILED`.
 798        """
 799        return self._client.post(":finalize", DataLoadResponse, params=dict(id=load_id))
 800
 801    def delete(self, load_id: int) -> DataLoadResponse:
 802        """Delete the data load if it has not been completed yet.
 803
 804        The deletion will happen asynchronously afterwards and is irreversible.
 805        """
 806        return self._client.post(":delete", DataLoadResponse, params=dict(id=load_id))
 807
 808    def undo(self, load_id: int) -> DataLoadResponse:
 809        """Append an entry to the changelog undoing the load.
 810
 811        The rows from this data load will no longer show at subsequent data versions.
 812        Previously undone loads can be redone to restore their data.
 813        """
 814        return self._client.post(":undo", DataLoadResponse, params=dict(id=load_id))
 815
 816    def redo(self, load_id: int) -> DataLoadResponse:
 817        """Append an entry to the changelog redoing the load. The rows from this data load will be visible again at subsequent data versions."""
 818        return self._client.post(":redo", DataLoadResponse, params=dict(id=load_id))
 819
 820    def add_table(self, load_id: int, request: AddTableRequest) -> DataLoadResponse:
 821        """Register a new table in the data load configuration."""
 822        return self._client.post(":addTable", DataLoadResponse, params=dict(id=load_id), json=request)
 823
 824
 825class TaskClient:
 826    """Create endpoints for all supported task types.
 827
 828    When creating a task, a context must be specified. The context determines data visibility: data can either be visible
 829    in the entire project, or only in a specific round within a project.
 830    Artifact and table data also has a context and a task can only access data with a compatible context. A data context
 831    is compatible when it a workspace context or, if it is a project or round context, belongs to the same project.
 832
 833    For example, suppose project `A` has rounds `abc` and `def` and project `B` has round `xyz`.
 834    A task with context `project=A` or `round=abc` can only see data with context `project=A` or `round=abc` or `round=def`.
 835    A task with context `project=B` can only see data with context `project=B` or `round=xyz`.
 836
 837    Any data a task creates will be saved with the same context as the task.
 838
 839    Tasks always use the latest version of the data in the workspace, unless a specific version is provided.
 840    """
 841
 842    def __init__(self, client: HttpClient):
 843        self._client = client.with_prefix("/task")
 844
 845    def list(self, project_id: int | None = None, round_id: int | None = None) -> Generator[TaskResponse, None, None]:
 846        """List all tasks in a workspace, project or round.
 847
 848        To list all tasks in a workspace, leave `project_id` and `round_id` unset.
 849        To list all tasks in a project, specify `project_id` and leave `round_id` unset.
 850        To list all tasks in a round, specify `round_id` and leave `project_id` unset.
 851
 852        Only one of `project_id` or `round_id` must be specified at once.
 853        """
 854        params = {}
 855        if project_id is not None:
 856            params["project_id"] = project_id
 857        if round_id is not None:
 858            params["round_id"] = round_id
 859
 860        opts = ListOptions()
 861        while True:
 862            resp = self._client.get(":list", ListTaskResponse, params={**params, **opts.model_dump()})
 863            yield from resp.items
 864            if resp.cursor is None:
 865                return
 866            opts.cursor = resp.cursor
 867
 868    def get(self, task_id: int) -> TaskResponse:
 869        """Obtain the current state and the result of a task.
 870
 871        This endpoint can be called periodically to monitor the status of a task and to detect
 872        when it finished executing.
 873
 874        If the `state` field of the response is `COMPLETED`, then the `result` field contains the
 875        task result (often containing IDs of artifacts and tables created by the task).
 876
 877        If the `state` field of the response is `FAILED`, then the `error` field contains
 878        a description explaining why the task has failed.
 879
 880        The states `COMPLETED`, `FAILED` and `CANCELLED` are terminal states. When the task is in
 881        any other state, it will eventually transition into one of these three.
 882        """
 883        return self._client.get(":get", TaskResponse, dict(id=task_id))
 884
 885    def get_by_name(
 886        self,
 887        name: str,
 888        project_id: int | None = None,
 889        round_id: int | None = None,
 890    ) -> TaskResponse:
 891        """Get a task by unique name within their workspace, project, or round context.
 892
 893        If a task with name "TaskA" exists in project with ID 123, the `project_id` must be specified.
 894        It is not sufficient to leave it unset as tasks with the same name may exist at the workspace or
 895        round context or in a different project.
 896        """
 897        tasks = self.list(project_id=project_id, round_id=round_id)
 898        # When listing for workspace or project we still need to filter it down since they'll
 899        # return tasks from their sub-contexts.
 900        if round_id is not None:
 901            scope_desc = f"round {round_id}"
 902            tasks = [task for task in tasks if task.context == ContextRound(round_id=round_id)]
 903        elif project_id is not None:
 904            scope_desc = f"project {project_id}"
 905            tasks = [task for task in tasks if task.context == ContextProject(project_id=project_id)]
 906        else:
 907            scope_desc = "workspace"
 908            tasks = [task for task in tasks if task.context == ContextWorkspace()]
 909
 910        for task in tasks:
 911            if task.name == name:
 912                return task
 913        raise ValueError(f"Task with name '{name}' not found in {scope_desc}")
 914
 915    def get_by_idempotency_key(
 916        self,
 917        idempotency_key: str,
 918        round_id: int | None = None,
 919        project_id: int | None = None,
 920    ) -> TaskResponse:
 921        """Get a task by its idempotency key within a specific round or project.
 922
 923        If a task with idempotency key "TaskA" exists in project with ID 123, the `project_id` must be specified.
 924        It is not sufficient to leave it unset as tasks with the same idempotency key may exist at the workspace or
 925        round context or in a different project.
 926        """
 927        params = {"idempotency_key": idempotency_key}
 928        if round_id is not None:
 929            params["round_id"] = str(round_id)
 930        if project_id is not None:
 931            params["project_id"] = str(project_id)
 932        return self._client.get(":getByIdempotencyKey", TaskResponse, params=params)
 933
 934    def create(self, task: TaskCreate) -> TaskResponse:
 935        return self._client.post(f"/{task.parameters.task_type}:create", TaskResponse, json=task)
 936
 937    def create_or_get(self, task: TaskCreate, skip_parameter_check: bool = False) -> TaskResponse:
 938        """Create a task or get existing one with matching (name or idempotency key) and parameters.
 939
 940        The task must have a name or an idempotency key set. If the parameter check is not skipped the existing task must have
 941        the same parameters as `task`. If an explicit `data_version_id` is specified, it must also match.
 942
 943        Raises ValueError if parameters don't match or if name is None.
 944        """
 945
 946        error_prefix: str = ""
 947
 948        # If using idempotency_key, try to create and handle 409 Conflict by fetching existing task
 949        if task.idempotency_key is not None:
 950            error_prefix = f"Task with idempotency key '{task.idempotency_key}'"
 951            try:
 952                return self.create(task)
 953            except ClientError as e:
 954                if e.status_code != 409:
 955                    raise
 956
 957                if isinstance(task.context, ContextRound):
 958                    existing_task = self.get_by_idempotency_key(task.idempotency_key, round_id=task.context.round_id)
 959                elif isinstance(task.context, ContextProject):
 960                    existing_task = self.get_by_idempotency_key(
 961                        task.idempotency_key, project_id=task.context.project_id
 962                    )
 963                else:
 964                    existing_task = self.get_by_idempotency_key(task.idempotency_key)
 965        elif name := task.name:
 966            error_prefix = f"Task with name '{name}'"
 967            try:
 968                if isinstance(task.context, ContextRound):
 969                    existing_task = self.get_by_name(name, round_id=task.context.round_id)
 970                elif isinstance(task.context, ContextProject):
 971                    existing_task = self.get_by_name(name, project_id=task.context.project_id)
 972                else:
 973                    existing_task = self.get_by_name(name)
 974            except ValueError:
 975                existing_task = None
 976        else:
 977            raise ValueError("Task must have `name` or `idempotency_key` set")
 978
 979        if existing_task is None:
 980            return self.create(task)
 981
 982        if skip_parameter_check:
 983            return existing_task
 984
 985        if existing_task.parameters != task.parameters:
 986            existing_params = existing_task.parameters.model_dump_json(indent=2, by_alias=True)
 987            new_params = task.parameters.model_dump_json(indent=2, by_alias=True)
 988            diff = "\n".join(
 989                difflib.unified_diff(
 990                    existing_params.splitlines(),
 991                    new_params.splitlines(),
 992                    fromfile="existing_parameters",
 993                    tofile="new_parameters",
 994                    lineterm="",
 995                )
 996            )
 997            raise ValueError(f"{error_prefix} exists but has different parameters:\n{diff}")
 998
 999        if task.data_version_id is not None and existing_task.data_version_id != task.data_version_id:
1000            raise ValueError(
1001                f"{error_prefix} exists but has different data_version_id (existing: {existing_task.data_version_id}, new: {task.data_version_id})"
1002            )
1003
1004        return existing_task
1005
1006    def cancel(self, task_id: int) -> TaskResponse:
1007        """Cancel the task if possible.
1008
1009        Only tasks that have not yet terminated may be canceled. There is no guarantee that cancellation succeeds
1010        before the task successfully completes.
1011        """
1012        return self._client.post(":cancel", TaskResponse, dict(id=task_id))
1013
1014    def recover(self, task_id: int) -> TaskResponse:
1015        """Recover a task that has failed or been cancelled by putting it into state EXECUTING.
1016
1017        Note that this is best-effort and is not guaranteed to fix any issues in the underlying
1018        backend execution, and may also cause unexpected side-effects such as duplicate data loads.
1019        Main use-case is tasks that failed due to intermittent issues in the backend which have been resolved.
1020        """
1021        return self._client.post(":recover", TaskResponse, dict(id=task_id))
1022
1023    def archive(self, task_id: int) -> TaskResponse:
1024        """Archive the task.
1025
1026        This will make the task disappear from regular task listings. It will not remove data that has been
1027        produced by the task. If data removal is required, use the `data_load_id` for the task's results
1028        and call the `/data/load:undo` endpoint separately.
1029        """
1030        return self._client.post(":archive", TaskResponse, dict(id=task_id))
1031
1032    def unarchive(self, task_id: int) -> TaskResponse:
1033        """Unarchive the task."""
1034        return self._client.post(":unarchive", TaskResponse, dict(id=task_id))
1035
1036    def wait(self, task_id: int, timeout: float = 60.0) -> TaskResponse:
1037        start = time.monotonic()
1038
1039        while time.monotonic() - start < timeout:
1040            try:
1041                task = self.get(task_id)
1042                if task.state == TaskState.FAILED:
1043                    raise ValueError(f"Task {task_id} failed: {task.errors}")
1044                if task.state == TaskState.CANCELLED:
1045                    raise ValueError(f"Task {task_id} was cancelled")
1046                if task.state == TaskState.COMPLETED:
1047                    break
1048            except httpx.TransportError:
1049                logger.warning("Transport error while waiting for task, retrying...")
1050            except ClientError as e:
1051                if 500 <= e.status_code < 600:
1052                    logger.warning("Internal server error while waiting for task, retrying...")
1053                else:
1054                    raise
1055
1056            time.sleep(3)
1057        print()
1058        return self.get(task_id)
1059
1060
1061class CustomPredictorClient:
1062    """Custom predictors can be used instead of or in addition to predictors trained on the Cradle platform.
1063
1064    A custom predictor is a function that, given a sequence, provides one or more scores for that sequence.
1065    These APIs allow to upload a Docker image that contains a custom predictor to the Cradle platform.
1066
1067    If you would like to use custom predictors, please get in touch with your customer success representative.
1068    """
1069
1070    def __init__(self, client: HttpClient):
1071        self._client = client.with_prefix("/customPredictor")
1072
1073    def get_current(self, name: str) -> CustomPredictorResponse:
1074        """Return the custom predictor with the most recent version"""
1075        return self._client.get(":get", CustomPredictorResponse, params={"name": name})
1076
1077    def get_version(self, version_id: int) -> CustomPredictorResponse:
1078        """Return the custom predictor with the given version ID"""
1079        return self._client.get(":getVersion", CustomPredictorResponse, params={"version_id": version_id})
1080
1081    def list(self) -> Generator[CustomPredictorResponse, None, None]:
1082        """List the active and being-provisioned versions of every predictor in the workspace."""
1083        opts = ListOptions()
1084        while True:
1085            resp = self._client.get(":list", ListCustomPredictorResponse, params=opts.model_dump())
1086            yield from resp.items
1087            if resp.cursor is None:
1088                return
1089            opts.cursor = resp.cursor
1090
1091    def create(
1092        self, cp: CustomPredictorContainerCreate, container_image: Path, upload_timeout: float = 300
1093    ) -> CustomPredictorContainerCreateResponse:
1094        """Create a new custom predictor.
1095
1096        The newly created predictor will initially be in `PROVISIONING` state.
1097        This method provides the container image at the given path on the
1098        filesystem, which should be a `tar` file as created by `docker save`.
1099
1100        The returned custom predictor can be polled regularly (using the
1101        `get_version()` method) until it transitions to either `ACTIVE` or
1102        `FAILED` state. Once it is in `ACTIVE` state it can be used. When it
1103        is in `FAILED` state, it means the uploaded container image could not
1104        be used. In this case, the `error` field will have more information and
1105        a new upload can be attempted by using the `update()` method.
1106        """
1107        with container_image.open("rb") as f:
1108            client = self._client
1109            predictor = client.post(":create", CustomPredictorCreateResponse, json=cp)
1110
1111            # Upload to a signed URL - use a fresh client so the API client's auth isn't attached.
1112            response = httpx.put(
1113                predictor.root.upload_url,
1114                headers=predictor.root.upload_headers,
1115                content=f,
1116                timeout=upload_timeout,
1117                follow_redirects=True,
1118            )
1119        if not response.is_success:
1120            raise ClientError(response.status_code, response.text, [])
1121
1122        return predictor.root
1123
1124    def update(
1125        self, name: str, cp: CustomPredictorContainerUpdate, container_image: Path, upload_timeout: float = 300
1126    ) -> CustomPredictorContainerCreateResponse:
1127        """Create a new version for a custom predictor.
1128
1129        See the `create()` method for how to provide a new container image for
1130        the updated custom predictor.
1131        """
1132        with container_image.open("rb") as f:
1133            client = self._client
1134            predictor = client.post(":update", CustomPredictorCreateResponse, params={"name": name}, json=cp)
1135
1136            # Upload to a signed URL - use a fresh client so the API client's auth isn't attached.
1137            response = httpx.put(
1138                predictor.root.upload_url,
1139                headers=predictor.root.upload_headers,
1140                content=f,
1141                timeout=upload_timeout,
1142                follow_redirects=True,
1143            )
1144
1145        if not response.is_success:
1146            raise ClientError(response.status_code, response.text, [])
1147
1148        return predictor.root
1149
1150    def delete(self, name: str) -> None:
1151        """Delete the custom predictor with the given name
1152
1153        This transitions all versions of this custom predictor (the initial one
1154        created with the `:create` endpoint and all subsequent versions created with
1155        the `:update` endpoint) to the `INACTIVE` state, unless it is already in
1156        one of the `INACTIVE` or `FAILED` states.
1157        """
1158        self._client.post(":delete", response_type=None, params={"name": name})