Skip to content

datarobot_genai.dragent.frontends.stream_converter

stream_converter

Convert ChatResponseChunk streams into AG-UI events.

NAT 1.6 added a stream_fn to tool_calling_agent that yields ChatResponseChunk objects (OpenAI-compatible streaming deltas). This module converts that stream into DRAgentEventResponse batches containing AG-UI TextMessage* and ToolCall* events.

See docs/nat-1.6-streaming.md for the full design.

convert_chunks_to_agui_events async 

convert_chunks_to_agui_events(chunks: AsyncGenerator[ChatResponseChunk, None]) -> AsyncGenerator[DRAgentEventResponse, None]

Convert a ChatResponseChunk stream into AG-UI events.

Yields DRAgentEventResponse batches as chunks arrive. On upstream errors, emits RunErrorEvent and stops (does not propagate). On GeneratorExit (client disconnect), exits silently.

Source code in datarobot_genai/dragent/frontends/stream_converter.py 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
async def convert_chunks_to_agui_events(
    chunks: AsyncGenerator[ChatResponseChunk, None],
) -> AsyncGenerator[DRAgentEventResponse, None]:
    """Convert a ChatResponseChunk stream into AG-UI events.

    Yields ``DRAgentEventResponse`` batches as chunks arrive.  On upstream
    errors, emits ``RunErrorEvent`` and stops (does not propagate).  On
    ``GeneratorExit`` (client disconnect), exits silently.
    """
    active_message_id: str | None = None
    # parent_message_id for subsequent tool calls; a synthetic uuid here
    # renders an orphan message stub in the UI.
    last_text_message_id: str | None = None
    seen_tool_calls: bool = False
    tool_index_map: dict[int, str] = {}
    zero = default_usage_metrics()

    error: Exception | None = None
    try:
        async for chunk in chunks:
            if not isinstance(chunk, ChatResponseChunk) or not chunk.choices:
                continue

            delta = chunk.choices[0].delta
            events: list[Event] = []

            if delta and delta.content:
                # Args streaming is complete for all tracked tool calls.
                # Flush any end/result events deferred by the step adaptor.
                for tc_id in tool_index_map.values():
                    events.extend(mark_args_done(tc_id))
                tool_index_map.clear()

                if active_message_id is None:
                    # After a tool call cycle, the LLM response is a new turn
                    # but chunk.id may be reused.  Force a unique messageId.
                    active_message_id = (
                        str(uuid.uuid4()) if seen_tool_calls else (chunk.id or str(uuid.uuid4()))
                    )
                    events.append(TextMessageStartEvent(message_id=active_message_id))
                events.append(
                    TextMessageContentEvent(message_id=active_message_id, delta=delta.content)
                )

            if delta and delta.tool_calls:
                # Close any active text message before starting tool calls.
                if active_message_id is not None:
                    events.append(TextMessageEndEvent(message_id=active_message_id))
                    last_text_message_id = active_message_id
                    active_message_id = None
                seen_tool_calls = True

                for tc in delta.tool_calls:
                    tc_id = tc.id or tool_index_map.get(tc.index)  # type: ignore[assignment]
                    if tc_id is None:
                        logger.warning(
                            "Tool call chunk at index %d has no id and no prior mapping; skipping",
                            tc.index,
                        )
                        continue
                    is_new = tc.id is not None and tc.index not in tool_index_map
                    if is_new:
                        tool_index_map[tc.index] = tc_id
                        tool_name = tc.function.name if tc.function else ""
                        events.append(
                            ToolCallStartEvent(
                                tool_call_id=tc_id,
                                tool_call_name=tool_name,
                                parent_message_id=last_text_message_id or "",
                            )
                        )
                        # Hand the LLM-issued id to the step adaptor.
                        if tool_name:
                            register_tool_call(tool_name, tc_id)
                    arguments = tc.function.arguments if tc.function else None
                    if arguments:
                        events.append(ToolCallArgsEvent(tool_call_id=tc_id, delta=arguments))

            if events:
                yield DRAgentEventResponse(events=events, usage_metrics=zero, original_chunk=chunk)
    except Exception as exc:
        error = exc
    finally:
        if sys.exc_info()[0] is GeneratorExit:
            logger.debug("Client disconnected before end events could be delivered")
            return

    # Emit end/error events after the stream completes (normally or on error).
    # Errors are surfaced to the AG-UI client via RunErrorEvent rather than
    # propagated as exceptions, so NAT's streaming infrastructure stays stable.
    end: list[Event] = []
    # Mark remaining in-flight tool calls as args-done and flush deferred events.
    for tc_id in tool_index_map.values():
        end.extend(mark_args_done(tc_id))
    if active_message_id is not None:
        end.append(TextMessageEndEvent(message_id=active_message_id))
    if error is not None:
        end.append(RunErrorEvent(message=str(error), code="STREAM_ERROR"))
    if end:
        yield DRAgentEventResponse(events=end, usage_metrics=zero)