feat: SSE streaming for service-mode answer generation - #2393
Open
jdye64 wants to merge 3 commits into
Open
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>
Contributor
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}
Comment on lines
+37
to
+41
| if self._in_thinking: | ||
| close_idx = self._pending.find(_THINK_CLOSE) | ||
| if close_idx == -1: | ||
| self._pending = "" | ||
| break |
Contributor
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.
Suggested change
| 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.
Comment on lines
+248
to
+298
| 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 = "" |
Contributor
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.
Comment on lines
+291
to
+295
| payload = json.loads(data_buf) | ||
| except json.JSONDecodeError: | ||
| data_buf = "" | ||
| event_type = "" | ||
| continue |
Contributor
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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)