feat: SSE streaming for service-mode answer generation - #2393
Conversation
Introduce POST /v1/answer/stream to forward LLM tokens to clients as Server-Sent Events after VectorDB retrieval, with TTFT metrics and incremental think-tag filtering. Refactor shared answer helpers and add RetrieverServiceClient.aanswer_stream() for async consumption. Co-authored-by: Jeremy Dyer <jdye64@gmail.com>
Fixes pre-commit failure on the answer stream refactor return tuple. Co-authored-by: Jeremy Dyer <jdye64@gmail.com>
Greptile SummaryAdds service-mode answer streaming over SSE.
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py | Adds async completion and RAG token streaming; the final batch result remains correct, but streamed visibility depends on the flawed incremental filter. |
| nemo_retriever/src/nemo_retriever/models/llm/text_utils.py | Adds stateful think-tag filtering that loses closing delimiters split across provider chunks and can suppress all visible tokens. |
| nemo_retriever/src/nemo_retriever/service/client.py | Adds the async SSE client, but its production path lacks direct tests and silently discards malformed event payloads. |
| nemo_retriever/src/nemo_retriever/service/routers/ingest.py | Refactors blocking-answer helpers and adds the streaming route with retrieval, generation, optional judging, and terminal events. |
| nemo_retriever/tests/test_llm_params.py | Verifies LiteLLM streaming invocation and ordinary delta extraction. |
| nemo_retriever/tests/test_service_answer_stream.py | Covers endpoint success, generation failure, disabled configuration, and an unsplit think-tag case, but not split closing tags or the public client. |
Sequence Diagram
sequenceDiagram
participant C as Client
participant API as /v1/answer/stream
participant VDB as VectorDB
participant LLM as LiteLLM
C->>API: POST answer request
API->>VDB: Retrieve top-k chunks
VDB-->>API: Chunks and metadata
API-->>C: retrieval_done
API->>LLM: Async streaming completion
loop Provider deltas
LLM-->>API: Text delta
API-->>C: metrics/token SSE
end
API-->>C: done or error SSE
Prompt To Fix All With AI
Fix the following 3 code review issues. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 3
nemo_retriever/src/nemo_retriever/models/llm/text_utils.py:37-41
**Preserve split closing tags**
When a provider splits `</think>` across streamed deltas, this branch clears the partial delimiter and leaves the filter in thinking mode, suppressing every subsequent visible token even though the final `done` event contains the batch-parsed answer.
```suggestion
if self._in_thinking:
close_idx = self._pending.find(_THINK_CLOSE)
if close_idx == -1:
_, self._pending = _split_safe_suffix(self._pending, _THINK_CLOSE)
break
```
### Issue 2 of 3
nemo_retriever/src/nemo_retriever/service/client.py:248-298
**Test the public stream client**
The new `aanswer_stream` method has no tests exercising its production httpx stream, status handling, or SSE parser; the added endpoint tests use `TestClient` and a separate duplicated parser, allowing client-side regressions to pass CI.
### Issue 3 of 3
nemo_retriever/src/nemo_retriever/service/client.py:291-295
**Surface malformed SSE payloads**
When an SSE event contains malformed JSON, this branch silently discards the event and resets its type, so callers receive neither the event nor an actionable stream error. Raise a descriptive client error instead of continuing with missing stream data.
Reviews (2): Last reviewed commit: "Merge branch 'main' into cursor/sse-answ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if self._in_thinking: | ||
| close_idx = self._pending.find(_THINK_CLOSE) | ||
| if close_idx == -1: | ||
| self._pending = "" | ||
| break |
There was a problem hiding this comment.
When a provider splits </think> across streamed deltas, this branch clears the partial delimiter and leaves the filter in thinking mode, suppressing every subsequent visible token even though the final done event contains the batch-parsed answer.
| if self._in_thinking: | |
| close_idx = self._pending.find(_THINK_CLOSE) | |
| if close_idx == -1: | |
| self._pending = "" | |
| break | |
| if self._in_thinking: | |
| close_idx = self._pending.find(_THINK_CLOSE) | |
| if close_idx == -1: | |
| _, self._pending = _split_safe_suffix(self._pending, _THINK_CLOSE) | |
| break |
Knowledge Base Used: Model Backends
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/models/llm/text_utils.py
Line: 37-41
Comment:
**Preserve split closing tags**
When a provider splits `</think>` across streamed deltas, this branch clears the partial delimiter and leaves the filter in thinking mode, suppressing every subsequent visible token even though the final `done` event contains the batch-parsed answer.
```suggestion
if self._in_thinking:
close_idx = self._pending.find(_THINK_CLOSE)
if close_idx == -1:
_, self._pending = _split_safe_suffix(self._pending, _THINK_CLOSE)
break
```
**Knowledge Base Used:** [Model Backends](https://app.greptile.com/nvidia-public-github/-/custom-context/knowledge-base/nvidia/nemo-retriever/-/docs/models.md)
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
Confirmed independently on the current head, and I consider this merge-blocking.
f = ThinkTagStreamFilter()
f.feed("<think>hidden</thi") # []
f.feed("nk>Visible answer") # []The second call should emit the visible answer, but the first call discarded the partial closing delimiter and left the filter permanently in thinking mode. The focused suite still passes 118 tests, so this boundary is currently unprotected.
Please preserve the longest suffix that can prefix </think> and add parameterized coverage for every split position of both <think> and </think>. The streaming test should also assert that concatenated visible deltas equal the completed answer.
There was a problem hiding this comment.
The fix — replace the two-line self._pending = ""; break with the same _split_safe_suffix call that the non-thinking branch already uses:
| if self._in_thinking: | |
| close_idx = self._pending.find(_THINK_CLOSE) | |
| if close_idx == -1: | |
| self._pending = "" | |
| break | |
| if close_idx == -1: | |
| _, self._pending = _split_safe_suffix(self._pending, _THINK_CLOSE) | |
| break |
New tests to add — add this class to tests/test_llm_params.py:
class TestThinkTagStreamFilter:
"""ThinkTagStreamFilter boundary and concatenation contracts."""
@pytest.mark.parametrize("split_at", range(1, len("<think>")))
def test_open_tag_split_across_chunks(self, split_at):
from nemo_retriever.models.llm.text_utils import ThinkTagStreamFilter
tag = "<think>"
f = ThinkTagStreamFilter()
r1 = f.feed("Prefix" + tag[:split_at])
r2 = f.feed(tag[split_at:] + "hidden</think>Visible")
assert "".join(r1 + r2) == "PrefixVisible"
@pytest.mark.parametrize("split_at", range(1, len("</think>")))
def test_close_tag_split_across_chunks(self, split_at):
from nemo_retriever.models.llm.text_utils import ThinkTagStreamFilter
close_tag = "</think>"
f = ThinkTagStreamFilter()
r1 = f.feed("<think>hidden" + close_tag[:split_at])
r2 = f.feed(close_tag[split_at:] + "Visible answer")
assert "".join(r1 + r2) == "Visible answer"
def test_streamed_deltas_concatenate_to_completed_answer(self):
"""Concatenated visible deltas must equal the batch-stripped answer."""
from nemo_retriever.models.llm.text_utils import ThinkTagStreamFilter, strip_think_tags
raw_chunks = [
"<think>",
"some long\n",
"reasoning chain",
"</think>",
"The ",
"answer ",
"is 42.",
]
full_text = "".join(raw_chunks)
expected = strip_think_tags(full_text)
f = ThinkTagStreamFilter()
deltas = []
for chunk in raw_chunks:
deltas.extend(f.feed(chunk))
assert "".join(deltas) == expectedThe test_close_tag_split_across_chunks parametrization covers all 7 split positions of </think>, including your exact reproduction case ("</thi" | "nk>Visible answer"). The test_open_tag_split_across_chunks covers all 6 split positions of <think>. The concatenation test uses strip_think_tags as the ground-truth oracle so it stays valid if the prompt format evolves.
| async def aanswer_stream( | ||
| self, | ||
| query: str, | ||
| *, | ||
| top_k: int = 5, | ||
| include_chunks: bool = False, | ||
| include_metadata: bool = False, | ||
| reasoning_enabled: bool | None = None, | ||
| reference: str | None = None, | ||
| judge: bool = False, | ||
| ) -> AsyncIterator[dict[str, Any]]: | ||
| """Stream ``POST /v1/answer/stream`` SSE events for one answer request.""" | ||
| url = f"{self._base_url}/v1/answer/stream" | ||
| body: dict[str, Any] = { | ||
| "query": query, | ||
| "top_k": top_k, | ||
| "include_chunks": include_chunks, | ||
| "include_metadata": include_metadata, | ||
| "judge": judge, | ||
| } | ||
| if reasoning_enabled is not None: | ||
| body["reasoning_enabled"] = reasoning_enabled | ||
| if reference is not None: | ||
| body["reference"] = reference | ||
|
|
||
| async with httpx.AsyncClient( | ||
| timeout=httpx.Timeout(None, connect=30.0), | ||
| headers=self._auth_headers, | ||
| ) as client: | ||
| async with client.stream("POST", url, json=body) as response: | ||
| if response.status_code >= 400: | ||
| detail = (await response.aread()).decode(errors="replace")[:500] | ||
| raise RuntimeError(f"Service answer stream failed: HTTP {response.status_code}: {detail}") | ||
|
|
||
| event_type = "" | ||
| data_buf = "" | ||
| async for line in response.aiter_lines(): | ||
| if line.startswith("event:"): | ||
| event_type = line[6:].strip() | ||
| elif line.startswith("data:"): | ||
| data_buf = line[5:].strip() | ||
| elif line == "" and data_buf: | ||
| try: | ||
| payload = json.loads(data_buf) | ||
| except json.JSONDecodeError: | ||
| data_buf = "" | ||
| event_type = "" | ||
| continue | ||
| yield {"event": event_type or "message", **payload} | ||
| data_buf = "" | ||
| event_type = "" |
There was a problem hiding this comment.
The new aanswer_stream method has no tests exercising its production httpx stream, status handling, or SSE parser; the added endpoint tests use TestClient and a separate duplicated parser, allowing client-side regressions to pass CI.
Rule Used: New functionality must include corresponding unit ... (source)
Knowledge Base Used: Service: HTTP/MCP API surface
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/service/client.py
Line: 248-298
Comment:
**Test the public stream client**
The new `aanswer_stream` method has no tests exercising its production httpx stream, status handling, or SSE parser; the added endpoint tests use `TestClient` and a separate duplicated parser, allowing client-side regressions to pass CI.
**Rule Used:** New functionality must include corresponding unit ... ([source](.greptile))
**Knowledge Base Used:** [Service: HTTP/MCP API surface](https://app.greptile.com/nvidia-public-github/-/custom-context/knowledge-base/nvidia/nemo-retriever/-/docs/service-api.md)
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
Affirming this. The endpoint tests parse SSE through separate test code, so they do not protect the behavior callers receive from aanswer_stream().
Direct client tests should cover:
- normal event reconstruction through the production parser;
- non-2xx responses;
- malformed JSON;
- multiple
data:lines in one SSE event; - disconnect/EOF before a terminal event;
- token concatenation matching
done.answer.
In particular, EOF without done or error should not look like successful stream completion.
| payload = json.loads(data_buf) | ||
| except json.JSONDecodeError: | ||
| data_buf = "" | ||
| event_type = "" | ||
| continue |
There was a problem hiding this comment.
When an SSE event contains malformed JSON, this branch silently discards the event and resets its type, so callers receive neither the event nor an actionable stream error. Raise a descriptive client error instead of continuing with missing stream data.
Knowledge Base Used: Service: HTTP/MCP API surface
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/service/client.py
Line: 291-295
Comment:
**Surface malformed SSE payloads**
When an SSE event contains malformed JSON, this branch silently discards the event and resets its type, so callers receive neither the event nor an actionable stream error. Raise a descriptive client error instead of continuing with missing stream data.
**Knowledge Base Used:** [Service: HTTP/MCP API surface](https://app.greptile.com/nvidia-public-github/-/custom-context/knowledge-base/nvidia/nemo-retriever/-/docs/service-api.md)
How can I resolve this? If you propose a fix, please make it concise.There was a problem hiding this comment.
Agree—this is a stream-integrity failure, not something the client can safely skip. If the malformed frame is a token, error, or done event, continuing can return incomplete output or apparent clean EOF.
Please raise a contextual client exception containing the event type and a bounded payload excerpt, preserving the original JSONDecodeError as the cause. I'd pair this with an explicit missing-terminal-event error so all truncated streams fail visibly.
There was a problem hiding this comment.
Here's the updated SSE loop with a contextual raise and a terminal-event check:
| payload = json.loads(data_buf) | |
| except json.JSONDecodeError: | |
| data_buf = "" | |
| event_type = "" | |
| continue | |
| event_type = "" | |
| data_buf = "" | |
| terminal_seen = False | |
| async for line in response.aiter_lines(): | |
| if line.startswith("event:"): | |
| event_type = line[6:].strip() | |
| elif line.startswith("data:"): | |
| data_buf = line[5:].strip() | |
| elif line == "" and data_buf: | |
| try: | |
| payload = json.loads(data_buf) | |
| except json.JSONDecodeError as exc: | |
| excerpt = data_buf[:200] | |
| raise RuntimeError( | |
| f"SSE event '{event_type or 'message'}' contained malformed JSON: {excerpt!r}" | |
| ) from exc | |
| evt = event_type or "message" | |
| if evt in ("done", "error"): | |
| terminal_seen = True | |
| yield {"event": evt, **payload} | |
| data_buf = "" | |
| event_type = "" | |
| if not terminal_seen: | |
| raise RuntimeError( | |
| "SSE stream closed without a terminal 'done' or 'error' event; response may be truncated" | |
| ) |
Key changes:
from exc— chains the originalJSONDecodeErrorso tracebacks show the exact parse failure.- 200-char excerpt — bounded so log lines stay readable; enough to diagnose the malformation.
terminal_seenguard — set ondoneorerrorbefore yielding, then checked after the loop exits. This catches mid-stream connection drops and partial responses that would otherwise look like clean EOF.
Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.
charlesbluca
left a comment
charlesbluca
left a comment
There was a problem hiding this comment.
I ran a live base/head integration comparison using the same ingested corpus, vector DB, and TP1 NVFP4 Nano backend on two RTX PRO 6000 Blackwell GPUs.
After five warmups per arm, 30 rotated sequential measurements showed:
- Base blocking response: 341.7 ms median
- Head SSE first visible token: 189.9 ms median
- Improvement: 151.8 ms / 44.4%
- Bootstrap 95% interval: 130.5–209.3 ms
- Head blocking median versus base: +2.5%
All 30 measured streams completed without errors, and concatenated token events exactly matched each terminal done.answer. This validates the PR's core TTFT motivation without indicating a material blocking-path regression.
I'm requesting changes because the split </think> bug can suppress the entire visible answer, and the stream still has integrity/liveness gaps around malformed events and a silent provider. I've added independent evidence to the existing Greptile threads and one new comment for keepalive/disconnect handling.
|
|
||
| generation: GenerationResult | None = None | ||
| try: | ||
| async for event_type, payload in llm.stream_generate( |
There was a problem hiding this comment.
[P2] Keep the connection alive while awaiting provider output
After retrieval_done, each __anext__() can remain suspended until the provider emits another delta. During that interval the endpoint sends no keepalive and does not execute request.is_disconnected(). A sufficiently slow first token can cross an ingress/proxy idle timeout, while a disconnected request can continue consuming model resources until another delta arrives.
Please consume generation through a producer task/queue with timed waits, following the existing job-SSE keepalive pattern: emit an SSE comment on timeout, check for disconnection independently of provider output, and cancel/close the upstream producer when the client leaves. Add a test with a stalled provider that verifies both keepalives and prompt cancellation after disconnect.
3 participants
Summary
Adds low-latency token streaming for answer generation in service mode via a new SSE endpoint.
POST /v1/answer/stream— streams retrieval + LLM answer tokens as Server-Sent EventsLiteLLMClient.stream_complete()/stream_generate()— async LiteLLM streaming with TTFT metrics and incremental think-tag filteringRetrieverServiceClient.aanswer_stream()— async Python client for consuming the SSE streamSSE event types
retrieval_donechunk_count,retrieval_latency_s, optionalchunks/metadatametricsttft_s, optionalgeneration_latency_stokendelta,indexdoneAnswerResultpayload (same fields asPOST /v1/answer)errordetailUsage
Notes
POST /v1/answeris unchanged for batch/eval workflows.answertool remains blocking; interactive UIs should use the HTTP stream endpoint.reasoning_enabled: falseon the request for faster visible TTFT on Nemotron models.Testing
tests/test_service_answer_stream.pytests/test_llm_params.py::TestLiteLLMStreamingtests/test_service_answer_generation.py(refactored helpers)