datarobot_genai.drmcputils.files.store

store

Block storage for drtools, backed by the DataRobot Files API (v3.10+).

Defines a minimal :class:BlobStore Protocol and a Files-API-backed implementation. Higher-level domains (e.g. panels) depend on the Protocol, not the concrete backend, so storage stays swappable (a local/in-memory backend can satisfy the same contract for tests) and the Files API is the production default.

The DataRobot Files SDK (datarobot.models.Files) is synchronous; its calls are dispatched to a worker thread via :func:asyncio.to_thread, which copies the current context so the per-request client configured by :func:request_user_dr_sdk remains in effect inside the thread.

BlobRef dataclass

A lightweight, serializable handle to a stored blob.

Every field is durably round-tripped by the Files API, so a ref returned by :meth:BlobStore.put and the same blob's ref returned by :meth:BlobStore.list compare equal:

• files_id — the durable handle (the DataRobot Files container id).
• name — the stored container name.
• tags — the stored tags (a tuple so the ref stays hashable).

Note: byte size and MIME content-type are intentionally not fields. The Files API does not persist them on the container (and search_catalog never returns them), so they could not be populated symmetrically by list. Callers that need a content-type should record it themselves (e.g. a panel manifest); see content_type on :meth:BlobStore.put.

Source code in datarobot_genai/drmcputils/files/store.py

@dataclass(frozen=True, slots=True)
class BlobRef:
    """A lightweight, serializable handle to a stored blob.

    Every field is durably round-tripped by the Files API, so a ref returned by
    :meth:`BlobStore.put` and the same blob's ref returned by
    :meth:`BlobStore.list` compare equal:

    - ``files_id`` — the durable handle (the DataRobot Files container id).
    - ``name`` — the stored container name.
    - ``tags`` — the stored tags (a ``tuple`` so the ref stays hashable).

    Note: byte size and MIME content-type are intentionally *not* fields. The
    Files API does not persist them on the container (and ``search_catalog``
    never returns them), so they could not be populated symmetrically by
    ``list``. Callers that need a content-type should record it themselves
    (e.g. a panel manifest); see ``content_type`` on :meth:`BlobStore.put`.
    """

    files_id: str
    name: str
    tags: tuple[str, ...] = ()

BlobStore

Bases: Protocol

Storage seam for opaque byte payloads. Implementations must be async-safe.

Source code in datarobot_genai/drmcputils/files/store.py

@runtime_checkable
class BlobStore(Protocol):
    """Storage seam for opaque byte payloads. Implementations must be async-safe."""

    async def put(
        self,
        data: bytes,
        *,
        name: str,
        content_type: str | None = None,
        tags: list[str] | None = None,
        timeout: int = DEFAULT_PUT_TIMEOUT_SECONDS,
    ) -> BlobRef:
        """Store ``data`` and return a handle to it.

        ``content_type`` is advisory: the Files backend does not persist it, so
        it is not reflected on the returned :class:`BlobRef`. Callers that need
        it should record it alongside their own metadata.

        ``timeout`` bounds the upload (seconds); raise it for very large blobs.
        """
        ...

    async def get(self, ref: BlobRef | str) -> bytes:
        """Fetch the bytes for ``ref`` (a :class:`BlobRef` or raw files id)."""
        ...

    async def delete(self, ref: BlobRef | str) -> None:
        """Delete the blob referenced by ``ref``."""
        ...

    async def list(
        self,
        *,
        search: str | None = None,
        tags: list[str] | None = None,
        limit: int = DEFAULT_LIST_LIMIT,
        offset: int = 0,
    ) -> list[BlobRef]:
        """List stored blobs the requesting user can access."""
        ...

put async

put(data: bytes, *, name: str, content_type: str | None = None, tags: list[str] | None = None, timeout: int = DEFAULT_PUT_TIMEOUT_SECONDS) -> BlobRef

Store data and return a handle to it.

content_type is advisory: the Files backend does not persist it, so it is not reflected on the returned :class:BlobRef. Callers that need it should record it alongside their own metadata.

timeout bounds the upload (seconds); raise it for very large blobs.

Source code in datarobot_genai/drmcputils/files/store.py

async def put(
    self,
    data: bytes,
    *,
    name: str,
    content_type: str | None = None,
    tags: list[str] | None = None,
    timeout: int = DEFAULT_PUT_TIMEOUT_SECONDS,
) -> BlobRef:
    """Store ``data`` and return a handle to it.

    ``content_type`` is advisory: the Files backend does not persist it, so
    it is not reflected on the returned :class:`BlobRef`. Callers that need
    it should record it alongside their own metadata.

    ``timeout`` bounds the upload (seconds); raise it for very large blobs.
    """
    ...

get async

get(ref: BlobRef | str) -> bytes

Fetch the bytes for ref (a :class:BlobRef or raw files id).

Source code in datarobot_genai/drmcputils/files/store.py

async def get(self, ref: BlobRef | str) -> bytes:
    """Fetch the bytes for ``ref`` (a :class:`BlobRef` or raw files id)."""
    ...

delete async

delete(ref: BlobRef | str) -> None

Delete the blob referenced by ref.

Source code in datarobot_genai/drmcputils/files/store.py

async def delete(self, ref: BlobRef | str) -> None:
    """Delete the blob referenced by ``ref``."""
    ...

list async

list(*, search: str | None = None, tags: list[str] | None = None, limit: int = DEFAULT_LIST_LIMIT, offset: int = 0) -> list[BlobRef]

List stored blobs the requesting user can access.

Source code in datarobot_genai/drmcputils/files/store.py

async def list(
    self,
    *,
    search: str | None = None,
    tags: list[str] | None = None,
    limit: int = DEFAULT_LIST_LIMIT,
    offset: int = 0,
) -> list[BlobRef]:
    """List stored blobs the requesting user can access."""
    ...

DataRobotFilesBlobStore

:class:BlobStore backed by the DataRobot Files API (datarobot.models.Files).

Each blob is stored as its own single-file Files container; the container id is the durable handle (:attr:BlobRef.files_id). Scoped to the requesting user's DataRobot token; blocking SDK calls run in a worker thread.

Source code in datarobot_genai/drmcputils/files/store.py

class DataRobotFilesBlobStore:
    """:class:`BlobStore` backed by the DataRobot Files API (``datarobot.models.Files``).

    Each blob is stored as its own single-file Files container; the container id
    is the durable handle (:attr:`BlobRef.files_id`). Scoped to the requesting
    user's DataRobot token; blocking SDK calls run in a worker thread.
    """

    def __init__(self, *, headers_auth_only: bool = True) -> None:
        # When False, fall back to the application API token outside HTTP contexts.
        self._headers_auth_only = headers_auth_only

    async def put(
        self,
        data: bytes,
        *,
        name: str,
        content_type: str | None = None,
        tags: list[str] | None = None,
        timeout: int = DEFAULT_PUT_TIMEOUT_SECONDS,
    ) -> BlobRef:
        # content_type is advisory only — the Files API has no field for it.
        def _upload() -> dr.models.Files:
            buf = _NamedBytesIO(data, name)
            # use_archive_contents=False: store the payload as-is, never auto-extract.
            return dr.models.Files.create_from_file(
                filelike=buf,
                tags=tags,
                use_archive_contents=False,
                read_timeout=timeout,
                max_wait=timeout,
            )

        with request_user_dr_sdk(headers_auth_only=self._headers_auth_only):
            try:
                files = await asyncio.to_thread(_upload)
            except ClientError as exc:
                raise_tool_error_for_client_error(exc)
        logger.debug("Stored blob %s (name=%s, %d bytes)", files.id, name, len(data))
        return _blob_ref(files, default_name=name)

    async def get(self, ref: BlobRef | str) -> bytes:
        files_id = _files_id(ref)

        def _download() -> bytes:
            files = dr.models.Files.get(files_id)
            buffer = io.BytesIO()
            files.download(filelike=buffer)
            return buffer.getvalue()

        with request_user_dr_sdk(headers_auth_only=self._headers_auth_only):
            try:
                return await asyncio.to_thread(_download)
            except ClientError as exc:
                raise_tool_error_for_client_error(exc)

    async def delete(self, ref: BlobRef | str) -> None:
        files_id = _files_id(ref)
        with request_user_dr_sdk(headers_auth_only=self._headers_auth_only):
            try:
                await asyncio.to_thread(dr.models.Files.delete, files_id)
            except ClientError as exc:
                raise_tool_error_for_client_error(exc)

    async def list(
        self,
        *,
        search: str | None = None,
        tags: list[str] | None = None,
        limit: int = DEFAULT_LIST_LIMIT,
        offset: int = 0,
    ) -> list[BlobRef]:
        def _search() -> list[object]:
            return dr.models.Files.search_catalog(
                search=search,
                tags=tags,
                limit=limit,
                offset=offset,
            )

        with request_user_dr_sdk(headers_auth_only=self._headers_auth_only):
            try:
                results = await asyncio.to_thread(_search)
            except ClientError as exc:
                raise_tool_error_for_client_error(exc)
        if limit and len(results) >= limit:
            # The Files API caps each page at ``limit``; more may exist. Callers
            # that need them must page via ``offset``.
            logger.debug(
                "list() hit the page limit of %d; additional blobs may exist (use offset to page)",
                limit,
            )
        return [_blob_ref(item) for item in results]