datarobot_genai.dragent.plugins.streaming_memory_agent

streaming_memory_agent

Streaming counterpart to NAT's auto_memory_agent.

NAT's upstream auto_memory_agent calls inner_agent_fn.ainvoke(...) in its inner_agent_node, which collapses the inner stream into a single string. For AG-UI consumers that means only one CustomEvent("DEFAULT_NAT_RESPONSE") ever reaches the frontend; the token deltas and tool-call deltas the inner agent's stream_fn would have emitted are lost.

streaming_memory_agent performs the same mem0 capture/retrieve operations (save user message → retrieve and inject as system context → save AI response) but astreams the inner agent and yields its native DRAgentEventResponse events straight through, so token deltas and tool-call deltas surface as proper AG-UI TextMessage* / ToolCall* events.

Registered with register_per_user_function so the wrapper itself builds lazily inside a PerUserWorkflowBuilder. Without that, the wrapper's builder.get_function(inner_agent_name) call would run at workflow-build time and miss per-user inner agents (which are built lazily on first user invocation). Per-user inner agents end up in the per-user builder's cache before the workflow is built, and shared inner agents still resolve via fall-through to the shared builder — both cases work transparently.

The configuration surface (memory_name, inner_agent_name, save_user_messages_to_memory, retrieve_memory_for_every_response, save_ai_messages_to_memory, search_params, add_params) is inherited from upstream AutoMemoryAgentConfig so the two wrappers can be swapped without reauthoring workflow.yaml. When the referenced memory backend is unconfigured (dr_mem0_memory with no credentials), the wrapper passthroughs to the inner agent without memory operations.

StreamingMemoryAgentConfig

Bases: AutoMemoryAgentConfig

Streaming, per-user variant of auto_memory_agent.

Inherits memory_name, inner_agent_name, save_user_messages_to_memory, retrieve_memory_for_every_response, save_ai_messages_to_memory, search_params, and add_params from :class:AutoMemoryAgentConfig. Only the _type discriminator differs, so a workflow can switch between the two by renaming _type. When the referenced memory backend is unconfigured, the wrapper passthroughs to the inner agent without memory capture or retrieval.

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

class StreamingMemoryAgentConfig(  # type: ignore[call-arg, misc]
    AutoMemoryAgentConfig,
    name="streaming_memory_agent",
):
    """Streaming, per-user variant of ``auto_memory_agent``.

    Inherits ``memory_name``, ``inner_agent_name``,
    ``save_user_messages_to_memory``, ``retrieve_memory_for_every_response``,
    ``save_ai_messages_to_memory``, ``search_params``, and ``add_params``
    from :class:`AutoMemoryAgentConfig`.  Only the ``_type`` discriminator
    differs, so a workflow can switch between the two by renaming ``_type``.
    When the referenced memory backend is unconfigured, the wrapper passthroughs
    to the inner agent without memory capture or retrieval.
    """

streaming_memory_agent async

streaming_memory_agent(config: StreamingMemoryAgentConfig, builder: Builder) -> AsyncGenerator[Any, None]

Build the streaming memory agent workflow.

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

@register_per_user_function(  # type: ignore[untyped-decorator]
    config_type=StreamingMemoryAgentConfig,
    input_type=RunAgentInput,
    streaming_output_type=DRAgentEventResponse,
)
async def streaming_memory_agent(
    config: StreamingMemoryAgentConfig, builder: Builder
) -> AsyncGenerator[Any, None]:
    """Build the streaming memory agent workflow."""
    inner_agent_fn = await builder.get_function(config.inner_agent_name)

    async def _stream_to_str(
        responses: AsyncGenerator[DRAgentEventResponse],
    ) -> str:
        aggregated = aggregate_dragent_event_responses([r async for r in responses])
        return convert_dragent_event_response_to_str(aggregated)

    memory_editor = await builder.get_memory_client(config.memory_name)

    if not is_memory_editor_configured(memory_editor):

        async def _passthrough_stream_fn(
            input_message: RunAgentInput,
        ) -> Annotated[
            AsyncGenerator[DRAgentEventResponse, None],
            Streaming(convert=aggregate_dragent_event_responses),
        ]:
            async for event_response in inner_agent_fn.astream(input_message):
                yield event_response

        yield FunctionInfo.create(
            stream_fn=_passthrough_stream_fn,
            stream_to_single_fn=_stream_to_str,
            description=config.description,
        )
        return

    async def _stream_fn(
        input_message: RunAgentInput,
    ) -> Annotated[
        AsyncGenerator[DRAgentEventResponse, None],
        Streaming(convert=aggregate_dragent_event_responses),
    ]:
        user_text = _last_user_text(input_message.messages)
        user_id = _user_id_from_context()

        # 1. Capture the user's latest message.
        if config.save_user_messages_to_memory and user_text:
            try:
                await memory_editor.add_items(
                    [
                        MemoryItem(
                            conversation=[{"role": "user", "content": user_text}],
                            user_id=user_id,
                        )
                    ],
                    **config.add_params,
                )
            except Exception as exc:  # noqa: BLE001
                logger.warning("memory.add_items(user) failed: %s", exc)

        # 2. Retrieve relevant memory and inject it as a system message.
        messages = list(input_message.messages)
        if config.retrieve_memory_for_every_response and user_text:
            try:
                memory_items = await memory_editor.search(
                    query=user_text,
                    user_id=user_id,
                    **config.search_params,
                )
                texts = [item.memory for item in memory_items if item.memory]
                if texts:
                    messages = _with_memory_context(messages, "\n".join(texts))
            except Exception as exc:  # noqa: BLE001
                logger.warning("memory.search failed: %s", exc)

        inner_input = input_message.model_copy(update={"messages": messages})

        # 3. astream the inner agent. Inner per-user agents (or shared inner
        #    agents) yield DRAgentEventResponse natively; we pass them straight
        #    through to the caller while collecting them for the post-stream
        #    assistant memory save.
        collected: list[DRAgentEventResponse] = []
        try:
            async for event_response in inner_agent_fn.astream(inner_input):
                collected.append(event_response)
                yield event_response
        finally:
            # 4. Persist the assistant response. Run regardless of partial
            #    errors so partial output still lands in memory.
            if config.save_ai_messages_to_memory and collected:
                aggregated = aggregate_dragent_event_responses(collected)
                ai_text = convert_dragent_event_response_to_str(aggregated)
                if ai_text:
                    try:
                        await memory_editor.add_items(
                            [
                                MemoryItem(
                                    conversation=[{"role": "assistant", "content": ai_text}],
                                    user_id=user_id,
                                )
                            ],
                            **config.add_params,
                        )
                    except Exception as exc:  # noqa: BLE001
                        logger.warning("memory.add_items(assistant) failed: %s", exc)

    yield FunctionInfo.create(
        stream_fn=_stream_fn,
        stream_to_single_fn=_stream_to_str,
        description=config.description,
    )