datarobot_genai.dragent.memory_space_cache

memory_space_cache

DataRobot MemorySpace-backed KV cache for the agent card registry L2 layer.

Used only on enclave workloads (DR_WORKLOAD_EXTERNAL_URL_HOST + DR_WORKLOAD_EXTERNAL_URL_PREFIX with WORKLOAD_ID). Other runtimes use in-process L1 caching only.

Uses the Memory Service light ORM from datarobot.application_utils.persistence (DRMemorySpace / DRSession / DREvent / DRMemoryServiceClient). Each cache entry is one session located by a stable DRDeduplicationKey hash and carrying a single status event whose content is the opaque JSON payload. set_value patches that event in place instead of appending.

AgentCardCacheSession

Bases: DRSession

Memory Service session model for one agent-card registry L2 cache entry.

Source code in datarobot_genai/dragent/memory_space_cache.py

class AgentCardCacheSession(DRSession):
    """Memory Service session model for one agent-card registry L2 cache entry."""

    __description_prefix__ = "dragent_cache"
    __lifecycle_strategies__: ClassVar[list[dict[str, Any]]] = []

    cache_kind: Annotated[str, DRRangeKey]
    dedup_key: Annotated[str, DRDeduplicationKey]
    logical_key: str
    dragent_cache: bool = True

AgentCardCacheEvent

Bases: DREvent

Single-value status event carrying the opaque JSON cache payload.

Source code in datarobot_genai/dragent/memory_space_cache.py

class AgentCardCacheEvent(DREvent, session=AgentCardCacheSession):  # type: ignore[call-arg]
    """Single-value status event carrying the opaque JSON cache payload."""

    __event_type__ = CACHE_EVENT_TYPE

MemorySpaceKVCache

Store opaque JSON payloads in a DataRobot MemorySpace by logical key.

Source code in datarobot_genai/dragent/memory_space_cache.py

class MemorySpaceKVCache:
    """Store opaque JSON payloads in a DataRobot MemorySpace by logical key."""

    def __init__(
        self,
        *,
        memory_space_id: str,
        key_prefix: str = "dragent:",
        client: DRMemoryServiceClient | None = None,
    ) -> None:
        self._memory_space_id = memory_space_id
        normalized = key_prefix if key_prefix.endswith(":") else f"{key_prefix}:"
        self._key_prefix = normalized
        self._client = client
        self._space: DRMemorySpace | None = None
        self._sessions: dict[str, AgentCardCacheSession] = {}

    def _logical_key(self, key: str) -> str:
        return f"{self._key_prefix}{CACHE_KIND}:{key}"

    def _client_or_global(self) -> DRMemoryServiceClient:
        return self._client or _require_memory_client()

    async def _resolve_space(self) -> DRMemorySpace:
        if self._space is None:
            client = self._client_or_global()
            self._space = await _call_with_transport_retry(
                lambda: DRMemorySpace.get(client, self._memory_space_id),
                op="resolve_space",
            )
        return self._space

    def _cache_session(self, logical_key: str, session: AgentCardCacheSession) -> None:
        self._sessions[logical_key] = session

    def _invalidate_session(self, logical_key: str) -> None:
        self._sessions.pop(logical_key, None)

    async def _resolve_session(self, logical_key: str) -> AgentCardCacheSession | None:
        if session := self._sessions.get(logical_key):
            return session

        space = await self._resolve_space()
        dedup_key = _cache_deduplication_key(logical_key)
        try:
            resolved = cast(
                AgentCardCacheSession,
                await _call_with_transport_retry(
                    lambda: AgentCardCacheSession.get(space, dedup_key=dedup_key),
                    op="resolve_session",
                ),
            )
        except DRMemoryNotFoundError:
            return None
        except DRMemoryUnavailableError:
            raise
        except Exception:
            logger.debug("MemorySpace session lookup failed for %s", logical_key)
            return None

        self._cache_session(logical_key, resolved)
        return resolved

    async def _read_payload(self, session: AgentCardCacheSession) -> str | None:
        events = await _call_with_transport_retry(
            lambda: AgentCardCacheEvent.last(session, n=1, type=CACHE_EVENT_TYPE),
            op="read_payload",
        )
        if not events:
            return None
        return events[0].content

    async def _write_payload(self, session: AgentCardCacheSession, payload: str) -> None:
        events = await _call_with_transport_retry(
            lambda: AgentCardCacheEvent.last(session, n=1, type=CACHE_EVENT_TYPE),
            op="write_payload_read",
        )
        if events:
            await _call_with_transport_retry(
                lambda: events[0].patch(content=payload),
                op="write_payload_update",
            )
            return
        await _call_with_transport_retry(
            lambda: AgentCardCacheEvent.post(
                session,
                content=payload,
                emitter_type="agent",
            ),
            op="write_payload_post",
        )

    async def get_value(self, key: str) -> str | None:
        """Return the stored JSON payload for *key*, or ``None`` when missing."""
        logical_key = self._logical_key(key)
        try:
            session = await self._resolve_session(logical_key)
            if session is None:
                return None
            return await self._read_payload(session)
        except DRMemoryUnavailableError:
            logger.exception("MemorySpace cache read failed for %s", logical_key)
            return None
        except Exception:
            logger.exception("MemorySpace cache read failed for %s", logical_key)
            return None

    async def set_value(self, key: str, payload: str) -> None:
        """Upsert a JSON payload for *key*."""
        logical_key = self._logical_key(key)
        dedup_key = _cache_deduplication_key(logical_key)

        try:
            space = await self._resolve_space()
            session = await self._resolve_session(logical_key)
            if session is None:
                session = cast(
                    AgentCardCacheSession,
                    await _call_with_transport_retry(
                        lambda: AgentCardCacheSession.post(
                            space,
                            cache_kind=CACHE_KIND,
                            dedup_key=dedup_key,
                            logical_key=logical_key,
                            participants=[DRAGENT_CACHE_PARTICIPANT_ID],
                        ),
                        op="create_cache_session",
                    ),
                )
                self._cache_session(logical_key, session)
                logger.info(
                    "MemorySpace L2 cache: created session %s for %s",
                    getattr(session, "id", None),
                    logical_key,
                )
            await self._write_payload(session, payload)
        except Exception:
            logger.exception("MemorySpace cache write failed for %s", logical_key)

    async def delete_value(self, key: str) -> None:
        """Remove a cached payload for *key* when present."""
        logical_key = self._logical_key(key)
        try:
            session = await self._resolve_session(logical_key)
            if session is not None:
                await _call_with_transport_retry(session.delete, op="delete_session")
            self._invalidate_session(logical_key)
        except Exception:
            logger.exception("MemorySpace cache delete failed for %s", logical_key)

get_value async

get_value(key: str) -> str | None

Return the stored JSON payload for key, or None when missing.

Source code in datarobot_genai/dragent/memory_space_cache.py

async def get_value(self, key: str) -> str | None:
    """Return the stored JSON payload for *key*, or ``None`` when missing."""
    logical_key = self._logical_key(key)
    try:
        session = await self._resolve_session(logical_key)
        if session is None:
            return None
        return await self._read_payload(session)
    except DRMemoryUnavailableError:
        logger.exception("MemorySpace cache read failed for %s", logical_key)
        return None
    except Exception:
        logger.exception("MemorySpace cache read failed for %s", logical_key)
        return None

set_value async

set_value(key: str, payload: str) -> None

Upsert a JSON payload for key.

Source code in datarobot_genai/dragent/memory_space_cache.py

async def set_value(self, key: str, payload: str) -> None:
    """Upsert a JSON payload for *key*."""
    logical_key = self._logical_key(key)
    dedup_key = _cache_deduplication_key(logical_key)

    try:
        space = await self._resolve_space()
        session = await self._resolve_session(logical_key)
        if session is None:
            session = cast(
                AgentCardCacheSession,
                await _call_with_transport_retry(
                    lambda: AgentCardCacheSession.post(
                        space,
                        cache_kind=CACHE_KIND,
                        dedup_key=dedup_key,
                        logical_key=logical_key,
                        participants=[DRAGENT_CACHE_PARTICIPANT_ID],
                    ),
                    op="create_cache_session",
                ),
            )
            self._cache_session(logical_key, session)
            logger.info(
                "MemorySpace L2 cache: created session %s for %s",
                getattr(session, "id", None),
                logical_key,
            )
        await self._write_payload(session, payload)
    except Exception:
        logger.exception("MemorySpace cache write failed for %s", logical_key)

delete_value async

delete_value(key: str) -> None

Remove a cached payload for key when present.

Source code in datarobot_genai/dragent/memory_space_cache.py

async def delete_value(self, key: str) -> None:
    """Remove a cached payload for *key* when present."""
    logical_key = self._logical_key(key)
    try:
        session = await self._resolve_session(logical_key)
        if session is not None:
            await _call_with_transport_retry(session.delete, op="delete_session")
        self._invalidate_session(logical_key)
    except Exception:
        logger.exception("MemorySpace cache delete failed for %s", logical_key)

is_enclave_l2_workload

is_enclave_l2_workload() -> bool

Return True when this process can use the registry L2 MemorySpace cache.

Source code in datarobot_genai/dragent/memory_space_cache.py

def is_enclave_l2_workload() -> bool:
    """Return True when this process can use the registry L2 MemorySpace cache."""
    return resolve_external_workload_api_endpoint() is not None and is_workload_mode()

registry_cache_deduplication_key

registry_cache_deduplication_key(workload_id: str) -> str

Return the stable deduplication key for an enclave workload's registry L2 space.

Source code in datarobot_genai/dragent/memory_space_cache.py

def registry_cache_deduplication_key(workload_id: str) -> str:
    """Return the stable deduplication key for an enclave workload's registry L2 space."""
    return f"{_REGISTRY_CACHE_SPACE_DEDUP_PREFIX}:workload:{workload_id}"

get_memory_service_client

get_memory_service_client() -> DRMemoryServiceClient | None

Return the configured Memory Service client, if any.

Source code in datarobot_genai/dragent/memory_space_cache.py

def get_memory_service_client() -> DRMemoryServiceClient | None:
    """Return the configured Memory Service client, if any."""
    return _MemoryClientState.client

try_resolve_memory_space_id

try_resolve_memory_space_id() -> str | None

Return the registry L2 MemorySpace ID on enclave workloads, else None.

Creates or adopts the space on first call. Uses WORKLOAD_ID as a deduplication_key so replicas share one space. No-op when not on an enclave workload, when credentials are unavailable, or after the first successful provision in this process.

This is the agent card registry L2 cache, not agent memory (AGENT_MEMORY_SPACE_ID). Agent memory is provisioned on the control hub (Pulumi / task deploy-dev) and the Mem0 client talks to that same host. Other runtimes use in-process L1 caching only.

Import-time bootstrap uses this synchronous helper. Lifespan warmup and :func:~datarobot_genai.dragent.agent_card_registry.get_default_registry must call :func:try_resolve_memory_space_id_async to provision — asyncio.run cannot nest inside an already running event loop.

An id already stored by a prior provision (import bootstrap or the async helper) is returned synchronously even from a running loop, so YAML parse and AgentCardRegistry.__init__ can still attach L2 write-behind.

Source code in datarobot_genai/dragent/memory_space_cache.py

def try_resolve_memory_space_id() -> str | None:
    """Return the registry L2 MemorySpace ID on enclave workloads, else ``None``.

    Creates or adopts the space on first call. Uses ``WORKLOAD_ID`` as a
    ``deduplication_key`` so replicas share one space. No-op when not on an enclave
    workload, when credentials are unavailable, or after the first successful
    provision in this process.

    This is the agent card registry L2 cache, not agent memory
    (``AGENT_MEMORY_SPACE_ID``). Agent memory is provisioned on the control hub
    (Pulumi / ``task deploy-dev``) and the Mem0 client talks to that same host.
    Other runtimes use in-process L1 caching only.

    Import-time bootstrap uses this synchronous helper. Lifespan warmup and
    :func:`~datarobot_genai.dragent.agent_card_registry.get_default_registry`
    must call :func:`try_resolve_memory_space_id_async` to *provision* — ``asyncio.run``
    cannot nest inside an already running event loop.

    An id already stored by a prior provision (import bootstrap or the async
    helper) is returned synchronously even from a running loop, so YAML parse
    and ``AgentCardRegistry.__init__`` can still attach L2 write-behind.
    """
    if _ProvisionedRegistryCacheSpaceState.space_id is not None:
        return _ProvisionedRegistryCacheSpaceState.space_id
    try:
        return _run_async(_try_resolve_memory_space_id)
    except RuntimeError as exc:
        if "running event loop" not in str(exc):
            raise
        logger.debug(
            "try_resolve_memory_space_id called from a running event loop; "
            "use try_resolve_memory_space_id_async"
        )
        return None

try_resolve_memory_space_id_async async

try_resolve_memory_space_id_async() -> str | None

Async counterpart of :func:try_resolve_memory_space_id.

Use this from a running event loop (lifespan warmup, get_default_registry).

Source code in datarobot_genai/dragent/memory_space_cache.py

async def try_resolve_memory_space_id_async() -> str | None:
    """Async counterpart of :func:`try_resolve_memory_space_id`.

    Use this from a running event loop (lifespan warmup, ``get_default_registry``).
    """
    return await _try_resolve_memory_space_id()

try_configure_datarobot_memory_client

try_configure_datarobot_memory_client(*, api_token: str | None = None) -> bool

Configure the enclave memory client when possible; return False on failure.

Source code in datarobot_genai/dragent/memory_space_cache.py

def try_configure_datarobot_memory_client(
    *,
    api_token: str | None = None,
) -> bool:
    """Configure the enclave memory client when possible; return ``False`` on failure."""
    try:
        configure_datarobot_memory_client(api_token=api_token)
    except Exception as exc:
        logger.debug("MemorySpace client unavailable: %s", exc)
        return False
    return True

configure_datarobot_memory_client

configure_datarobot_memory_client(*, api_token: str | None = None) -> None

Configure the process-global Memory Service ORM client for enclave workloads.

Uses the enclave API gateway (DR_WORKLOAD_EXTERNAL_URL_HOST + DR_WORKLOAD_EXTERNAL_URL_PREFIX → {host}/api/v2). The ORM client talks directly to {endpoint}/memory and does not require dr.Client()'s GET /version/ probe.

Source code in datarobot_genai/dragent/memory_space_cache.py

def configure_datarobot_memory_client(
    *,
    api_token: str | None = None,
) -> None:
    """Configure the process-global Memory Service ORM client for enclave workloads.

    Uses the enclave API gateway (``DR_WORKLOAD_EXTERNAL_URL_HOST`` +
    ``DR_WORKLOAD_EXTERNAL_URL_PREFIX`` → ``{host}/api/v2``). The ORM client talks
    directly to ``{endpoint}/memory`` and does not require ``dr.Client()``'s
    ``GET /version/`` probe.
    """
    enclave_endpoint = resolve_external_workload_api_endpoint()
    if enclave_endpoint is None:
        raise ValueError(
            "MemorySpace cache backends require an enclave API gateway "
            "(DR_WORKLOAD_EXTERNAL_URL_HOST and DR_WORKLOAD_EXTERNAL_URL_PREFIX)."
        )
    token = api_token or os.getenv("DATAROBOT_API_TOKEN")
    if not token:
        raise ValueError("DATAROBOT_API_TOKEN is required when using memory_space cache backends.")
    logger.info(
        "Configuring DataRobot memory client for enclave gateway %s.",
        enclave_endpoint,
    )
    _MemoryClientState.endpoint = enclave_endpoint.rstrip("/")
    _MemoryClientState.api_token = token
    _MemoryClientState.client = DRMemoryServiceClient(
        endpoint=_MemoryClientState.endpoint,
        api_token=token,
    )