datarobot_genai.dragent.plugins.datarobot_mem0_memory

datarobot_mem0_memory

NAT memory provider backed by DataRobot's Mem0 client.

This provider wires datarobot-genai[dragent] into NAT's MemoryEditor interface so auto_memory_agent can store and retrieve long-term memory.

Backend selection is driven by config:

• agent_memory_space_id → the DataRobot Memory Service's mem0-compatible API, reached at {DATAROBOT_ENDPOINT}/memory/{agent_memory_space_id}/ and authenticated with the DataRobot API token. See PBMP-7431 ("Agentic Memory Service"), section "Connect to DataRobot memory" / "API Layout".
• api_key (or MEM0_API_KEY) → Mem0's hosted SaaS at https://api.mem0.ai. Used when no agent_memory_space_id is configured.

Both routes share the same MemoryEditor adapter because the DR endpoint is API-compatible with mem0 — only the host and auth token differ.

When neither route is configured (no agent_memory_space_id + token and no api_key / MEM0_API_KEY), the provider yields an :class:UnconfiguredMemoryEditor so workflows can declare dr_mem0_memory unconditionally and enable memory later via runtime parameters or env vars.

Config

Bases: DataRobotAppFrameworkBaseSettings

Finds variables in the priority order of: env variables (including Runtime Parameters), .env, file_secrets, then Pulumi output variables.

Source code in datarobot_genai/dragent/plugins/datarobot_mem0_memory.py

class Config(DataRobotAppFrameworkBaseSettings):
    """
    Finds variables in the priority order of: env
    variables (including Runtime Parameters), .env, file_secrets, then
    Pulumi output variables.
    """

    mem0_api_key: str | None = None
    agent_memory_space_id: str | None = None
    agent_memory_ttl_days: int | None = None

DRMem0MemoryClientConfig

Bases: MemoryBaseConfig, RetryMixin

A NAT memory backend backed by datarobot-genai's Mem0 client.

Backend selection:

• If agent_memory_space_id is set, requests go to the DataRobot Memory Service's mem0-compatible endpoint at {datarobot_endpoint}/memory/{agent_memory_space_id}/ authenticated with datarobot_api_token (or DATAROBOT_API_TOKEN).
• Otherwise, api_key (or MEM0_API_KEY) is used against Mem0's hosted SaaS (host defaults to https://api.mem0.ai).

Source code in datarobot_genai/dragent/plugins/datarobot_mem0_memory.py

class DRMem0MemoryClientConfig(  # type: ignore[call-arg]
    MemoryBaseConfig,  # type: ignore[misc]
    RetryMixin,  # type: ignore[misc]
    name="dr_mem0_memory",
):
    """A NAT memory backend backed by ``datarobot-genai``'s Mem0 client.

    Backend selection:

    * If ``agent_memory_space_id`` is set, requests go to the DataRobot Memory
      Service's mem0-compatible endpoint at
      ``{datarobot_endpoint}/memory/{agent_memory_space_id}/`` authenticated with
      ``datarobot_api_token`` (or ``DATAROBOT_API_TOKEN``).
    * Otherwise, ``api_key`` (or ``MEM0_API_KEY``) is used against Mem0's
      hosted SaaS (``host`` defaults to ``https://api.mem0.ai``).
    """

    api_key: str | None = Field(
        default_factory=_get_default_mem0_api_key_for_memory_backend,
        description="Mem0 API key used when targeting Mem0's hosted SaaS.",
    )
    host: str | None = Field(
        default=None,
        description=(
            "Mem0 base URL for the SaaS backend. Ignored when ``agent_memory_space_id`` is set."
        ),
    )
    org_id: str | None = None
    project_id: str | None = None
    agent_memory_space_id: str | None = Field(
        default_factory=_get_default_agent_memory_space_id,
        description=(
            "DataRobot MemorySpace ID. When set, the editor uses the DataRobot "
            "Memory Service's mem0-compatible endpoint instead of Mem0 SaaS. "
            "The endpoint is built as ``{datarobot_endpoint}/memory/{id}/``."
        ),
    )
    datarobot_endpoint: str | None = Field(
        default=None,
        description=(
            "DataRobot API base URL used to build the mem0 endpoint when "
            "``agent_memory_space_id`` is set (e.g. ``https://app.datarobot.com/api/v2``). "
            "Defaults to the ``DATAROBOT_ENDPOINT`` env var."
        ),
    )
    datarobot_api_token: str | None = Field(
        default=None,
        description=(
            "DataRobot API token used when ``agent_memory_space_id`` is set. "
            "Defaults to the ``DATAROBOT_API_TOKEN`` env var."
        ),
    )
    default_ttl_days: int | None = Field(
        default_factory=_get_default_ttl_days,
        ge=0,
        description=(
            "Default TTL in days for stored memories. When set to a "
            "positive value, the editor passes "
            "``expiration_date = today + default_ttl_days`` (UTC, "
            "``YYYY-MM-DD``) through to Mem0's ``add`` API so memories "
            "auto-expire. Callers may override per-call by passing "
            "``expiration_date`` in ``add_params``. ``None`` or ``0`` "
            "leaves the field unset (no expiration). Defaults from the "
            "``AGENT_MEMORY_TTL_DAYS`` env var."
        ),
    )

UnconfiguredMemoryEditor

Bases: MemoryEditor

No-op memory backend returned when dr_mem0_memory has no credentials.

Source code in datarobot_genai/dragent/plugins/datarobot_mem0_memory.py

class UnconfiguredMemoryEditor(MemoryEditor):  # type: ignore[misc]
    """No-op memory backend returned when ``dr_mem0_memory`` has no credentials."""

    async def add_items(self, items: list[MemoryItem], **kwargs: Any) -> None:
        return

    async def search(self, query: str, top_k: int = 5, **kwargs: Any) -> list[MemoryItem]:
        return []

    async def remove_items(self, **kwargs: Any) -> None:
        return

DRMem0Editor

Bases: MemoryEditor

Adapt Mem0Client to NAT's MemoryEditor interface.

Source code in datarobot_genai/dragent/plugins/datarobot_mem0_memory.py

class DRMem0Editor(MemoryEditor):  # type: ignore[misc]
    """Adapt ``Mem0Client`` to NAT's ``MemoryEditor`` interface."""

    def __init__(
        self,
        client: Any,
        ttl_days: int | None = None,
        *,
        store_name: str = "mem0",
        store_id: str | None = None,
    ) -> None:
        self._client = client
        self._mem0 = client._memory
        self._ttl_days = ttl_days
        self._store_name = store_name
        self._store_id = store_id

    def _memory_span_attributes(
        self,
        *,
        user_id: str | None = None,
        extra: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        attrs: dict[str, Any] = dict(extra or {})
        if user_id:
            attrs["gen_ai.memory.scope"] = "user"
            attrs["memory.user_id"] = user_id
        return attrs

    async def add_items(self, items: list[MemoryItem], **kwargs: Any) -> None:
        configured_user_id = kwargs.get("user_id")
        resolved_user_ids = {
            item.user_id or configured_user_id or self._mem0.user_id for item in items
        }
        user_id = next(iter(resolved_user_ids)) if len(resolved_user_ids) == 1 else None

        with trace_memory_operation(
            "update_memory",
            store_name=self._store_name,
            store_id=self._store_id,
            attributes=self._memory_span_attributes(
                user_id=user_id,
                extra={"memory.item_count": len(items)},
            ),
        ):
            await self._add_items_impl(items, **kwargs)

    async def _add_items_impl(self, items: list[MemoryItem], **kwargs: Any) -> None:
        add_kwargs = dict(kwargs)
        output_format = add_kwargs.pop("output_format", "v1.1")
        configured_run_id = add_kwargs.pop("run_id", None)
        configured_tags = add_kwargs.pop("tags", None)
        configured_metadata = dict(add_kwargs.pop("metadata", None) or {})
        configured_user_id = add_kwargs.pop("user_id", None)

        # Inject the configured TTL as Mem0's ``expiration_date`` only when
        # the caller hasn't supplied one. Per-call overrides (e.g. via
        # ``add_params: {expiration_date: ...}`` in workflow.yaml) win so
        # special-case memories can opt out of or extend the default.
        if "expiration_date" not in add_kwargs and self._ttl_days:
            add_kwargs["expiration_date"] = _ttl_to_expiration_date(self._ttl_days)

        coroutines = []
        for item in items:
            metadata = configured_metadata | dict(item.metadata or {})
            run_id = metadata.pop("run_id", configured_run_id)
            user_id = item.user_id or configured_user_id or self._mem0.user_id
            item_kwargs = add_kwargs | {
                "user_id": user_id,
                "tags": item.tags or configured_tags or [],
                "metadata": metadata,
                "output_format": output_format,
            }
            if run_id:
                item_kwargs["run_id"] = run_id
            coroutines.append(self._mem0.add(item.conversation, **item_kwargs))

        if coroutines:
            await asyncio.gather(*coroutines)

    async def search(self, query: str, top_k: int = 5, **kwargs: Any) -> list[MemoryItem]:
        user_id = kwargs.get("user_id") or self._mem0.user_id
        with trace_memory_operation(
            "search_memory",
            store_name=self._store_name,
            store_id=self._store_id,
            attributes=self._memory_span_attributes(
                user_id=user_id,
                extra={
                    "gen_ai.memory.query.text": truncate_memory_text(query),
                    "memory.top_k": top_k,
                },
            ),
        ) as span:
            memories = await self._search_impl(query, top_k=top_k, **kwargs)
            span.set_attribute("gen_ai.memory.search.result.count", len(memories))
            return memories

    async def _search_impl(self, query: str, top_k: int = 5, **kwargs: Any) -> list[MemoryItem]:
        user_id = kwargs.pop("user_id", None) or self._mem0.user_id
        conditions: list[dict[str, Any]] = [{"user_id": user_id}]
        for key in ("run_id", "agent_id", "app_id"):
            value = kwargs.pop(key, None)
            if value:
                conditions.append({key: value})

        metadata_filter = kwargs.pop("metadata", None)
        if metadata_filter:
            conditions.append({"metadata": metadata_filter})

        filters = kwargs.pop("filters", None) or {"AND": conditions}
        output_format = kwargs.pop("output_format", "v1.1")
        result = await self._mem0.search(
            query,
            filters=filters,
            top_k=top_k,
            output_format=output_format,
            **kwargs,
        )

        memories: list[MemoryItem] = []
        for raw_result in result.get("results", []) or []:
            if not isinstance(raw_result, dict):
                continue
            item_meta = raw_result.get("metadata") or {}
            memories.append(
                MemoryItem(
                    conversation=raw_result.get("input") or [],
                    user_id=user_id,
                    memory=raw_result.get("memory"),
                    tags=raw_result.get("categories") or [],
                    metadata=item_meta,
                )
            )
        return memories

    async def remove_items(self, **kwargs: Any) -> None:
        has_memory_id = "memory_id" in kwargs
        user_id = kwargs.get("user_id")
        if not has_memory_id and not user_id:
            return

        memory_id = kwargs.get("memory_id") if has_memory_id else None
        with trace_memory_operation(
            "delete_memory",
            store_name=self._store_name,
            store_id=self._store_id,
            attributes=self._memory_span_attributes(
                user_id=user_id,
                extra={
                    **({"gen_ai.memory.record.id": memory_id} if memory_id else {}),
                    **({"memory.delete_all": True} if user_id and not has_memory_id else {}),
                },
            ),
        ):
            await self._remove_items_impl(**kwargs)

    async def _remove_items_impl(self, **kwargs: Any) -> None:
        if "memory_id" in kwargs:
            await self._mem0.delete(kwargs.pop("memory_id"))
            return
        user_id = kwargs.pop("user_id", None)
        if user_id:
            await self._mem0.delete_all(user_id=user_id)

is_memory_editor_configured

is_memory_editor_configured(editor: MemoryEditor) -> bool

Return False when dr_mem0_memory yielded an unconfigured no-op editor.

Source code in datarobot_genai/dragent/plugins/datarobot_mem0_memory.py

def is_memory_editor_configured(editor: MemoryEditor) -> bool:
    """Return ``False`` when ``dr_mem0_memory`` yielded an unconfigured no-op editor."""
    return not isinstance(editor, UnconfiguredMemoryEditor)