nemo_gym.chat_streaming

Streaming Chat Completions-dialect support shared by every Gym model server.

Blackbox harnesses that speak the OpenAI Chat Completions API over SSE (e.g. the OpenClaw agent PinchBench runs) send requests the strict NeMoGymChatCompletionCreateParamsNonStreaming model rejects: a stream: true flag (its stream field is typed Literal[False], since a Gym model server calls its backend non-streaming and buffers the whole response) plus a stream_options block, which is only meaningful alongside stream: true.

A Gym model server can still serve these clients by computing the complete Chat Completion with its existing non-streaming backend call and re-emitting it as an SSE stream. This module provides:

the request-side sanitizer that maps the streaming wire body onto the strict params shape (drops stream/stream_options, remembers whether usage was requested);
the response-side synthesizer that re-emits a complete NeMoGymChatCompletion as the chat.completion.chunk SSE sequence a streaming client expects, terminated by data: [DONE].

Only the SSE envelope is synthesized — there is no true token-by-token streaming. The backend call completes before the first byte is emitted, so the model server’s retry and error-normalization behavior is fully preserved on this path. This path is intended for eval-only streaming clients: token ids and logprobs from the backend response are not carried in the chat.completion.chunk schema, and a client that does not set stream_options.include_usage gets no usage chunk, so a model-call record reconstructed from this stream will lack token counts.

Module Contents

Functions

Name	Description
`_choice_deltas`	Yield the `chat.completion.chunk` choice deltas for one completed choice.
`_chunk`	Build one `chat.completion.chunk` object sharing the completion’s id/created/model.
`_sse_data`	-
`_wants_usage`	Whether the client asked for a terminal usage chunk (`stream_options.include_usage`).
`sanitize_streaming_chat_body`	Map a streaming-dialect chat body onto the strict non-streaming params shape.
`synthesize_chat_completion_sse`	Re-emit a complete Chat Completion object as a `chat.completion.chunk` SSE stream.

Data

LOG

_PARAM_FIELDS

API

nemo_gym.chat_streaming._choice_deltas(
    index: int,
    choice: dict[str, typing.Any]
) -> typing.Iterator[dict[str, typing.Any]]

Yield the chat.completion.chunk choice deltas for one completed choice.

A role delta opens the choice; reasoning, content, and tool-call deltas follow (each emitted only when present); a terminal delta carries the finish_reason. Splitting the message this way keeps every field a client tracks (role, reasoning, content, tool-call name/arguments, finish reason) in the delta position that client expects, even though it is a single logical chunk sequence rather than incremental tokens.

nemo_gym.chat_streaming._chunk(
    completion: dict[str, typing.Any],
    choices: list[dict[str, typing.Any]],
    usage: typing.Any = None
) -> dict[str, typing.Any]

Build one chat.completion.chunk object sharing the completion’s id/created/model.

nemo_gym.chat_streaming._sse_data(
    payload: dict[str, typing.Any]
) -> str

nemo_gym.chat_streaming._wants_usage(
    stream_options: typing.Any
) -> bool

Whether the client asked for a terminal usage chunk (stream_options.include_usage).

nemo_gym.chat_streaming.sanitize_streaming_chat_body(
    body: dict[str, typing.Any]
) -> tuple[dict[str, typing.Any], bool]

Map a streaming-dialect chat body onto the strict non-streaming params shape.

Returns the cleaned body dict (ready for NeMoGymChatCompletionCreateParamsNonStreaming validation) and whether a terminal usage chunk was requested via stream_options.include_usage.

stream and stream_options are removed: the params model’s stream field is typed Literal[False], and stream_options is only meaningful with stream: true, so it has no effect on the non-streaming backend call. Remaining fields are filtered to the known params fields, so a harness’s extra bookkeeping never reaches the backend.

nemo_gym.chat_streaming.synthesize_chat_completion_sse(
    completion: dict[str, typing.Any],
    include_usage: bool = False
) -> typing.Iterator[str]

Re-emit a complete Chat Completion object as a chat.completion.chunk SSE stream.

Emits, per choice, a role chunk -> optional reasoning/content/tool-call chunks -> a terminal chunk carrying finish_reason. When include_usage is set and the completion reports usage, a final choices: [] chunk carries the usage block (OpenAI’s contract). The stream always ends with the data: [DONE] sentinel streaming clients treat as terminal.

nemo_gym.chat_streaming.LOG = logging.getLogger(__name__)

nemo_gym.chat_streaming._PARAM_FIELDS = frozenset(NeMoGymChatCompletionCreateParamsNonStreaming.model_fields)