feat: accept stream:true on /v1/chat/completions via synthesized SSE#2130
feat: accept stream:true on /v1/chat/completions via synthesized SSE#2130ananthsub wants to merge 1 commit into
Conversation
…quests Blackbox harnesses that speak the OpenAI Chat Completions API over SSE (e.g. the OpenClaw agent PinchBench runs) send stream:true, which the strict NeMoGymChatCompletionCreateParamsNonStreaming model rejects. Mirror the /v1/responses streaming seam (NVIDIA-NeMo#2095) for chat/completions: a shared chat_completions_dispatch route validates non-streaming requests strictly (unchanged 422 behavior) and, for stream:true, sanitizes the wire body, delegates to the server's existing non-streaming chat_completions(), and re-emits the complete response as chat.completion.chunk SSE terminated by data: [DONE]. Backend call stays non-streaming, so retries and token-id capture are unaffected. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Do you have examples of those? I think in @adil-a's PR #2059, he mentioned Streamed HTTP, and I also have seen SSE (from anthropic messages/v1). Did you check the differences between those options? How does vllm handle it? |
The motivation is supporting the OpenClaw agent that PinchBench runs. Currently it points to a non-gym-based model URL so we cannot capture model call records |
Yeah, I was curious about the specific repro tool example, functionally, but it seems it is not in the issue either. I was curious if you used that as a test case |
ffrujeri
left a comment
ffrujeri
left a comment
There was a problem hiding this comment.
Background
- OpenAI Chat Completions streaming (chunk objects,
data: [DONE],stream_options.include_usage): API reference; pinned chunk schema for the installed client version: openai-python v2.7.2chat_completion_chunk.py - Server-sent events: MDN primer
- Related work: #2099 (motivating issue), #2095 (the
/v1/responsessynthesized-SSE analog for Codex), #2059 (MCP streamable-HTTP — a different streaming concern: MCP tool transport, not model-inference SSE)
Where this PR sits
This changes the core model-server base class: SimpleResponsesAPIModel is extended by every model server (openai_model, vllm_model + its two local variants, azure_openai_model, inference_provider), and the /v1/chat/completions route is rebound from the subclass's chat_completions to the new chat_completions_dispatch on the hot request path.
flowchart LR
Harness[Streaming CC harness<br/>OpenClaw agent / PinchBench] -->|"POST /v1/chat/completions<br/>stream: true"| Dispatch[chat_completions_dispatch<br/>nemo_gym/base_responses_api_model.py]
Dispatch -->|sanitize + validate| CS[chat_streaming.py<br/>NEW module]
Dispatch -->|strict params| Impl["subclass chat_completions()<br/>responses_api_models/*"]
Impl -->|non-streaming call| Backend[(OpenAI / vLLM / Azure / NIM)]
Impl -->|complete reply| CS
CS -->|"synthesized SSE chunks<br/>data: [DONE]"| Harness
Dispatch -.->|observed by| Capture[model-call capture<br/>_reconstruct_chat_sse]
classDef touched fill:#ffe08a,stroke:#d48806;
class Dispatch,CS touched
Summary
Nice, well-scoped change — the buffer-then-replay approach mirrors the /v1/messages path and #2095, the new module is thoroughly documented, and the 23-test suite (with the capture reconstructor used as a fidelity oracle) is a great touch. I verified all six concrete model servers are handled by the signature dispatch, the non-streaming 200 response body is byte-for-byte unchanged, and the sync generator in StreamingResponse doesn't block the event loop (Starlette iterates it in a threadpool).
The inline comments cover one small contract regression (an extra url key in 422 details — one-line fix in two places), a question about the reachability of the extra_forbidden pruning machinery, and a capture/observability question about usage accounting on the new streaming path.
Design notes (non-blocking)
- Buffer-then-replay vs true streaming. Since the backend call completes before the first byte is emitted, clients get no incremental token timing — which seems fine for the stated goal (capturing model-call records for OpenClaw/PinchBench). Worth one line in the PR description making explicit that this is deliberately buffer-then-replay, so future readers don't expect token-level streaming semantics. Is it confirmed that OpenClaw doesn't rely on incremental chunk timing (e.g. inter-token timeouts)?
- Synthesized SSE vs proxying vLLM's native stream. vLLM supports
stream: truenatively; synthesizing from a non-streaming call instead keeps retry/error-normalization and stays consistent with/v1/messagesand #2095. Reading the PR description this looks intentional — flagging only to confirm the consistency choice over proxying was deliberate.
Other observations (no action needed)
- The
else dict(completion)branch in_invoke_chat_completions's caller (base_responses_api_model.py:141) is unreached today — all shipped servers return a validatedNeMoGymChatCompletion(whoseid/created/modelare required). A custom subclass returning a raw dict missing those keys would synthesize chunks withnullid/model/created that strict SDK clients reject mid-stream; not worth guarding now, just noting for the record. - Minor optional polish: the streaming path deep-copies the body twice (
sanitize_streaming_chat_bodyandvalidate_streaming_chat_paramseach take a defensive copy), and_invoke_chat_completionsre-inspects thechat_completionssignature per request — both consistent with the existing/v1/messagesprecedent, so fine as is. fern/doesn't document the new streaming support, but the/v1/messagesstreaming dialect is likewise undocumented, so this is consistent with current practice; could be a follow-up.
| try: | ||
| return NeMoGymChatCompletionCreateParamsNonStreaming.model_validate(body) | ||
| except ValidationError as exc: | ||
| raise RequestValidationError([{**error, "loc": ("body", *error["loc"])} for error in exc.errors()]) |
There was a problem hiding this comment.
nemo_gym/base_responses_api_model.py:193
Was byte-for-byte 422 parity with the previous typed-Body() binding the intent here? exc.errors() defaults to include_url=True, so every 422 detail entry now carries an extra "url": "https://errors.pydantic.dev/..." key that FastAPI's native body validation strips — it calls errors(include_url=False) (fastapi 0.135.2, _compat/v2.py:174). I verified the detail entries become byte-identical to the pre-PR shape with the flag added (the existing tests only assert loc[0] == "body", so they don't catch this):
| raise RequestValidationError([{**error, "loc": ("body", *error["loc"])} for error in exc.errors()]) | |
| raise RequestValidationError( | |
| [{**error, "loc": ("body", *error["loc"])} for error in exc.errors(include_url=False)] | |
| ) |
| try: | ||
| params = validate_streaming_chat_params(cleaned) | ||
| except ValidationError as exc: | ||
| raise RequestValidationError([{**error, "loc": ("body", *error["loc"])} for error in exc.errors()]) |
There was a problem hiding this comment.
nemo_gym/base_responses_api_model.py:139
Same include_url point as in _validate_chat_params below — this streaming-path 422 also gains the extra url key:
| raise RequestValidationError([{**error, "loc": ("body", *error["loc"])} for error in exc.errors()]) | |
| raise RequestValidationError( | |
| [{**error, "loc": ("body", *error["loc"])} for error in exc.errors(include_url=False)] | |
| ) |
| return False | ||
|
|
||
|
|
||
| def validate_streaming_chat_params(body: dict[str, Any]) -> NeMoGymChatCompletionCreateParamsNonStreaming: |
There was a problem hiding this comment.
nemo_gym/chat_streaming.py:102
Is the extra_forbidden pruning intended as dormant future-proofing? As far as I can tell it's unreachable with the shipped models: NeMoGymChatCompletionCreateParamsNonStreaming is a plain BaseModel (default extra="ignore", unlike the Responses params model at openai_utils.py:315), and the nested SDK types are TypedDicts that silently drop unknown keys — I probed extras at the top level, inside messages, tools, and tool calls, and none raise extra_forbidden. Top-level extras are already stripped by sanitize_streaming_chat_body anyway. That leaves the retry loop, _delete_loc, and the LOG.warning branch (L87–99, L117–121) uncovered — chat_streaming.py sits at 81% in isolation.
Two easy resolutions: drop the pruning and let validate_streaming_chat_params be a thin strict-validate (simplest), or keep it as future-proofing with a comment noting it only activates if a params model forbids extras, plus a test using a synthetic extra="forbid" model so the branch is covered.
There was a problem hiding this comment.
Is this difference intentional? Could we keep both paths strict and only remove stream and stream_options here?
| yield _sse_data(_chunk(completion, [delta])) | ||
|
|
||
| usage = completion.get("usage") | ||
| if include_usage and usage is not None: |
There was a problem hiding this comment.
nemo_gym/chat_streaming.py:200
Was it intentional that captured streaming records lose token accounting unless the client opts in? The usage chunk is only emitted when the client sets stream_options.include_usage (correct per the OpenAI contract), but model-call capture reconstructs the record from the outbound SSE — so for clients that don't set it (OpenClaw-style harnesses typically don't), the captured record ends up with tokens_in/out/total = None even though the backend returned usage. Since capturing model-call records is the motivating use case, it might be worth capturing usage out-of-band from the backend response rather than from the synthesized stream (always emitting the usage chunk would deviate from the OpenAI contract, so that path seems worse). A documented limitation would also be fine.
Related wording nit: the module docstring's "token-id / logprob capture on the backend response is unaffected" reads a bit stronger than what happens — captured streaming records won't include token_ids/logprobs (they don't ride in the chunk schema, as the docstring itself notes). Maybe rephrase to say this path is intended for eval-only streaming clients that don't consume those fields.
| ``nemo_gym.chat_streaming``), delegated to the same ``chat_completions()``, and the | ||
| complete response is re-emitted as a synthesized ``chat.completion.chunk`` SSE stream. | ||
| """ | ||
| if not body.get("stream"): |
There was a problem hiding this comment.
nemo_gym/base_responses_api_model.py:131
Is truthiness-based routing intended for malformed stream values? stream: "false" (string) or stream: 1 now routes to the SSE path (the sanitizer then drops the field, so the Literal[False] guard never sees it), where the previous typed binding rejected them with 422. Strictly malformed input, so arguably fine — but if body.get("stream") is not True: would keep the historical strictness for everything except a genuine boolean true.
| async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()) -> NeMoGymResponse: | ||
| pass | ||
|
|
||
| async def chat_completions_dispatch(self, request: Request, body: dict = Body()): |
There was a problem hiding this comment.
nemo_gym/base_responses_api_model.py:119
Tiny note for the record (no action needed, and not fixable via include_url): with body: dict = Body(), a non-object JSON body (array/string/number) now fails FastAPI's dict coercion with type=dict_type / "Input should be a valid dictionary", where the previous typed binding produced type=model_attributes_type / "...valid dictionary or object to extract fields from". Exotic malformed input, but together with the url key it means "the same validation errors" isn't quite exact — might be worth a one-line caveat in the PR description.
| 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. | ||
| """ | ||
| for index, choice in enumerate(completion.get("choices") or []): |
There was a problem hiding this comment.
nemo_gym/chat_streaming.py:195
The per-choice synthesis here is correct (and the openai client accumulates by choice.index fine) — but it makes a pre-existing quirk newly reachable: _reconstruct_chat_sse collapses every choice into a single choices[0] when rebuilding a captured stream, so an n>1 streaming request would have its choices' content concatenated in the captured record. n>1 streaming is rare and this is capture-only, so just flagging — a follow-up in the reconstructor (group by index) would close it.
| assert resp.json()["choices"][0]["message"]["content"] == "hi there" | ||
| assert server.last_params.messages[0]["content"] == "hi there" | ||
|
|
||
| def test_non_streaming_request_still_validates_strictly(self) -> None: |
There was a problem hiding this comment.
tests/unit_tests/test_chat_completions_streaming.py:289
Really solid suite — especially reusing the capture reconstructor as a fidelity oracle. A few paths in the new code aren't exercised yet (_delete_loc L87–99, the prune-retry branch L117–121, and system_fingerprint pass-through L141 — chat_streaming.py sits at 81% in isolation), and the 422 assertions only check loc[0] == "body", which is what lets the url-key difference slip through. The block below runs against this branch (13 tests, all passing) and brings chat_streaming.py to 100%; feel free to take any subset.
Imports: change from fastapi import Body, Request → from fastapi import Body, FastAPI, Request, and add _delete_loc to the from nemo_gym.chat_streaming import (...) group. Then append:
def _legacy_app() -> TestClient:
"""A FastAPI app with the pre-PR direct typed-body binding, to compare 422 parity against."""
app = FastAPI()
async def chat_completions(body: NeMoGymChatCompletionCreateParamsNonStreaming = Body()):
return {"ok": True}
app.post("/v1/chat/completions")(chat_completions)
return TestClient(app)
def _strip_url(detail: list[dict]) -> list[dict]:
return [{k: v for k, v in error.items() if k != "url"} for error in detail]
_BAD_BODIES = [
{"model": "x"}, # missing required messages
{"messages": [], "temperature": "not-a-number"}, # wrong scalar type
{"messages": "not-a-list"}, # messages wrong type
{"messages": [{"role": "user"}]}, # user message missing required content
]
class Test422Parity:
"""The dispatch's 422 must match the pre-PR direct typed-body binding (modulo the `url` key)."""
@pytest.mark.parametrize("bad", _BAD_BODIES)
def test_non_streaming_422_matches_legacy_except_url(self, bad) -> None:
legacy, (client, _) = _legacy_app(), _client(_EchoChatModel)
legacy_resp = legacy.post("/v1/chat/completions", json=bad)
new_resp = client.post("/v1/chat/completions", json=bad)
assert legacy_resp.status_code == 422 and new_resp.status_code == 422
assert _strip_url(new_resp.json()["detail"]) == _strip_url(legacy_resp.json()["detail"])
@pytest.mark.parametrize("bad", _BAD_BODIES)
def test_streaming_422_matches_legacy_except_url(self, bad) -> None:
legacy, (client, _) = _legacy_app(), _client(_EchoChatModel)
legacy_resp = legacy.post("/v1/chat/completions", json=bad)
new_resp = client.post("/v1/chat/completions", json={**bad, "stream": True})
assert legacy_resp.status_code == 422 and new_resp.status_code == 422
assert _strip_url(new_resp.json()["detail"]) == _strip_url(legacy_resp.json()["detail"])
class TestDeleteLoc:
def test_deletes_nested_dict_value(self) -> None:
body = {"a": {"b": {"c": 1, "d": 2}}}
assert _delete_loc(body, ("a", "b", "c")) is True
assert body == {"a": {"b": {"d": 2}}}
def test_deletes_through_list_index(self) -> None:
body = {"items": [{"x": 1, "y": 2}]}
assert _delete_loc(body, ("items", 0, "x")) is True
assert body == {"items": [{"y": 2}]}
def test_unwalkable_loc_returns_false(self) -> None:
body = {"items": [{"x": 1}]}
assert _delete_loc(body, ("items", "function", "x")) is False # str where list index expected
assert body == {"items": [{"x": 1}]}
assert _delete_loc({"items": [{"x": 1}]}, ("items", 5, "x")) is False # index out of range
assert _delete_loc({"a": {"b": 1}}, ("a", "missing")) is False # last key absent
assert _delete_loc({"a": 5}, ("a", "b")) is False # non-container mid-walk
class TestSynthesizeSystemFingerprint:
def test_system_fingerprint_propagated_into_every_chunk(self) -> None:
completion = _completion(content="hi", usage=_USAGE).model_dump(mode="json")
completion["system_fingerprint"] = "fp_abc123"
events = _events("".join(synthesize_chat_completion_sse(completion)))
assert events
assert all(event.get("system_fingerprint") == "fp_abc123" for event in events)
class TestValidateStreamingPrune:
def test_prunes_extra_forbidden_field_and_revalidates(self, monkeypatch) -> None:
# The real params model ignores unknown fields (no extra="forbid"), so extra_forbidden never
# fires through it; swap in a forbid-extra model to exercise the prune-and-retry loop.
import nemo_gym.chat_streaming as chat_streaming
class _Forbidding(pydantic.BaseModel):
model_config = pydantic.ConfigDict(extra="forbid")
messages: list
temperature: float | None = None
monkeypatch.setattr(chat_streaming, "NeMoGymChatCompletionCreateParamsNonStreaming", _Forbidding)
params = chat_streaming.validate_streaming_chat_params(
{"messages": [{"role": "user", "content": "hi"}], "temperature": 1.0, "unknown_field": 123}
)
assert params.temperature == 1.0Two notes: Test422Parity tolerates the extra url key on purpose — once the include_url=False change lands, _strip_url(...) == _strip_url(...) can tighten to a plain == for true byte-for-byte parity. TestValidateStreamingPrune needs the monkeypatch because the real params model can't emit extra_forbidden; if you instead decide to drop the prune machinery (see the comment on validate_streaming_chat_params), skip that test and TestDeleteLoc.
|
Tested this PR locally with curl, both with regular responses and forcing the model to use a tool calla, the SSE stream worked as expected. |
There was a problem hiding this comment.
Should we update these the docs now that Gym model servers support streaming for chat completions with this change?
responses_api_agents/pinchbench/README.md
Model wiring: OpenClaw requires a streaming-capable Chat Completions endpoint.
model_base_urlmay point directly to an upstream endpoint or to a sandbox-reachable
Gym model server, which adapts its non-streaming backend response into OpenAI-compatible SSE.
responses_api_agents/pinchbench/app.py
Policy model OpenClaw runs against. May point directly to an upstream streaming endpoint or to a sandbox-reachable Gym model server — see README.
responses_api_agents/pinchbench/configs/pinchbench.yaml
Policy model OpenClaw runs against: a direct streaming endpoint or a sandbox-reachable Gym model server. Secrets are supplied via CLI/env overrides at run time — never committed.
github-actions
Bot
added
the
sla:review-overdue
3 participants
Summary
fixes #2099
Some tools that run a model through Gym only work if the model server accepts streaming requests. They send
stream: trueon/v1/chat/completionsand expect the reply back as a series of small events. Today Gym model servers rejectstream: true— the request fails validation — so these tools can't point at a Gym model server and have to call the model endpoint directly. When they do that, they lose what the Gym model server adds: one consistent request format across providers, and per-call capture. The OpenClaw agent that PinchBench runs is the case that prompted this.This change makes
/v1/chat/completionsacceptstream: true. The model server still calls its backend exactly as it does now: one normal request, wait for the complete reply. Once it has the full reply, it sends it back to the caller as a stream of standardchat.completion.chunkevents ending withdata: [DONE]. This is the same approach already used for/v1/messages(Claude Code) and, in #2095, for/v1/responses(Codex).Because the server still waits for the whole reply before sending anything, this does not make responses arrive sooner — the only goal is to accept the streaming request shape these tools require. Two useful consequences of keeping the backend call non-streaming: the existing retry and error handling still apply (nothing has been sent to the caller yet when a failure happens), and token id / logprob capture is unchanged (streaming here is for eval-only tools, and those fields aren't part of the chunk format anyway).
How it works
stream: true, the server removes the streaming-only fields (stream,stream_options) and any unknown fields the tool added, validates the rest against the normal request model, calls the server's existingchat_completions(), then turns the complete reply into chunks: the role, then reasoning / content / tool calls when present, then a final chunk carrying the finish reason, then a usage chunk if the caller asked for usage, thendata: [DONE].New file
nemo_gym/chat_streaming.pydoes the request cleanup, validation, and chunk generation. The routing change is inSimpleResponsesAPIModel. This covers only/v1/chat/completions; streaming for/v1/responsesis a separate change (#2095).Test plan
tests/unit_tests/test_chat_completions_streaming.py(23 tests): request cleanup, validation, chunk generation (content, reasoning, tool calls, usage, multiple choices, empty choices), and the route over an HTTP test client (non-streaming JSON, strict 422, streaming SSE, dropping unknown fields, invalid params → 422, and servers whosechat_completions()takes a leadingrequestargument). Fidelity is checked by feeding the generated stream back through the capture code that rebuilds a reply from chunks.test_base_responses_api_model.pyandtest_openai_utils.py(71 passed), and theopenai_modelandvllm_modelserver suites (80 passed), which exercise bothchat_completions()argument shapes over real HTTP.openaiPython client withstream: trueandinclude_usagerebuilt the full reply (content, reasoning, tool call, finish reason, usage) from the streamed chunks.