feat: SSE streaming for service-mode answer generation

nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py

@@ -14,6 +14,7 @@
 
 import logging
 import time
+from collections.abc import AsyncIterator
 from typing import Any, Optional
 
 from nemo_retriever.common.params.models import (
@@ -27,6 +28,7 @@
     _build_rag_prompt as _task_build_rag_prompt,
     _deep_merge_dicts,
 )
+from nemo_retriever.models.llm.text_utils import ThinkTagStreamFilter
 from nemo_retriever.models.llm.types import (
     GenerationResult,
     UnsupportedTextResponseError,
@@ -129,16 +131,15 @@ def from_kwargs(
         )
         return cls(transport=transport, sampling=sampling)
 
-    def complete(
+    def _build_call_kwargs(
         self,
         messages: list[dict],
         max_tokens: Optional[int] = None,
         extra_params: Optional[dict[str, Any]] = None,
-    ) -> tuple[str, float]:
-        """Raw litellm completion call. Returns (content_text, latency_s)."""
+    ) -> dict[str, Any]:
+        """Build provider-neutral kwargs shared by sync and streaming completion."""
         validate_llm_extra_params(self.transport.extra_params, source="LLMRemoteClientParams.extra_params")
         validate_llm_extra_params(extra_params, source="GenerationRequest.extra_params")
-        import litellm
 
         sampling_kwargs = self.sampling.to_sampling_kwargs()
         if max_tokens is not None:
@@ -160,6 +161,39 @@ def complete(
             if resolved_api_key:
                 call_kwargs["api_key"] = resolved_api_key
         call_kwargs.update(_deep_merge_dicts(self.transport.extra_params, extra_params or {}))
+        return call_kwargs
+
+    @staticmethod
+    def _delta_from_stream_chunk(chunk: object) -> str | None:
+        """Extract a text delta from one LiteLLM/OpenAI-compatible stream chunk."""
+        choices = _field(chunk, "choices")
+        if not isinstance(choices, (list, tuple)) or not choices:
+            return None
+        choice = choices[0]
+        delta = _field(choice, "delta")
+        if delta is None:
+            message = _field(choice, "message")
+            if message is not None:
+                delta = message
+        if delta is None:
+            return None
+        content = _field(delta, "content")
+        if content is None:
+            return None
+        if not isinstance(content, str) or not content:
+            return None
+        return content
+
+    def complete(
+        self,
+        messages: list[dict],
+        max_tokens: Optional[int] = None,
+        extra_params: Optional[dict[str, Any]] = None,
+    ) -> tuple[str, float]:
+        """Raw litellm completion call. Returns (content_text, latency_s)."""
+        import litellm
+
+        call_kwargs = self._build_call_kwargs(messages, max_tokens, extra_params)
 
         t0 = time.monotonic()
         try:
@@ -201,6 +235,110 @@ def complete(
         content = content.strip()
         return content, latency
 
+    async def stream_complete(
+        self,
+        messages: list[dict],
+        max_tokens: Optional[int] = None,
+        extra_params: Optional[dict[str, Any]] = None,
+    ) -> AsyncIterator[str]:
+        """Yield raw text deltas from ``litellm.acompletion(..., stream=True)``."""
+        import litellm
+
+        call_kwargs = self._build_call_kwargs(messages, max_tokens, extra_params)
+        call_kwargs["stream"] = True
+
+        try:
+            response = await litellm.acompletion(**call_kwargs)
+        except Exception as exc:
+            err = str(exc)
+            if "temperature" in err and "top_p" in err:
+                logger.error(
+                    "Model %s rejected the request because both `temperature` "
+                    "and `top_p` were specified. Some providers (e.g. Bedrock) "
+                    "only accept one. Either remove `top_p` from the model "
+                    "config or set `temperature` to null. Sent: "
+                    "temperature=%s, top_p=%s",
+                    self.transport.model,
+                    call_kwargs.get("temperature"),
+                    call_kwargs.get("top_p"),
+                )
+            raise
+
+        async for chunk in response:
+            delta = self._delta_from_stream_chunk(chunk)
+            if delta is not None:
+                yield delta
+
+    async def stream_generate(
+        self,
+        query: str,
+        chunks: list[str],
+        *,
+        reasoning_enabled: Optional[bool] = None,
+    ) -> AsyncIterator[tuple[str, dict[str, Any]]]:
+        """Stream visible RAG answer tokens and final generation metadata.
+
+        Yields ``("token", {"delta": str, "index": int})`` events followed by
+        ``("complete", {"answer": str, "latency_s": float, "model": str, "error": str | None})``.
+        """
+        effective_reasoning_enabled = (
+            self.transport.reasoning_enabled if reasoning_enabled is None else reasoning_enabled
+        )
+        task = RagAnswerTask(
+            system_prompt=self.transport.rag_system_prompt,
+            system_prompt_prefix=self.transport.rag_system_prompt_prefix,
+            reasoning_enabled=effective_reasoning_enabled,
+        )
+        request = task.build_request(
+            query=query,
+            chunks=chunks,
+            reasoning_enabled=reasoning_enabled,
+        )
+        think_filter = ThinkTagStreamFilter() if effective_reasoning_enabled is not False else None
+
+        t0 = time.monotonic()
+        ttft_s: float | None = None
+        raw_parts: list[str] = []
+        token_index = 0
+
+        try:
+            async for delta in self.stream_complete(
+                request.messages,
+                max_tokens=request.max_tokens,
+                extra_params=request.extra_params,
+            ):
+                raw_parts.append(delta)
+                visible_deltas = [delta] if think_filter is None else think_filter.feed(delta)
+                for visible in visible_deltas:
+                    if ttft_s is None:
+                        ttft_s = time.monotonic() - t0
+                        yield "metrics", {"ttft_s": ttft_s}
+                    yield "token", {"delta": visible, "index": token_index}
+                    token_index += 1
+        except Exception as exc:
+            yield "complete", {
+                "answer": "",
+                "latency_s": time.monotonic() - t0,
+                "model": self.model,
+                "error": str(exc),
+                "ttft_s": ttft_s,
+            }
+            return
+
+        raw_text = "".join(raw_parts)
+        parsed = task.parse(raw_text)
+        latency_s = time.monotonic() - t0
+        error = None if parsed else task.empty_output_error
+        if ttft_s is not None:
+            yield "metrics", {"ttft_s": ttft_s, "generation_latency_s": latency_s}
+        yield "complete", {
+            "answer": parsed,
+            "latency_s": latency_s,
+            "model": self.model,
+            "error": error,
+            "ttft_s": ttft_s,
+        }
+
     def generate(
         self,
         query: str,

nemo_retriever/src/nemo_retriever/models/llm/text_utils.py

@@ -14,6 +14,58 @@
 
 import re
 
+_THINK_OPEN = "<think>"
+_THINK_CLOSE = "</think>"
+
+
+class ThinkTagStreamFilter:
+    """Incrementally suppress ``<think>`` blocks from a token stream."""
+
+    def __init__(self) -> None:
+        self._pending = ""
+        self._in_thinking = False
+
+    def feed(self, chunk: str) -> list[str]:
+        """Return zero or more visible answer deltas from one streamed chunk."""
+        if not chunk:
+            return []
+
+        self._pending += chunk
+        emitted: list[str] = []
+
+        while self._pending:
+            if self._in_thinking:
+                close_idx = self._pending.find(_THINK_CLOSE)
+                if close_idx == -1:
+                    self._pending = ""
+                    break
+                self._pending = self._pending[close_idx + len(_THINK_CLOSE) :]
+                self._in_thinking = False
+                continue
+
+            open_idx = self._pending.find(_THINK_OPEN)
+            if open_idx == -1:
+                safe, self._pending = _split_safe_suffix(self._pending, _THINK_OPEN)
+                if safe:
+                    emitted.append(safe)
+                break
+
+            if open_idx > 0:
+                emitted.append(self._pending[:open_idx])
+            self._pending = self._pending[open_idx + len(_THINK_OPEN) :]
+            self._in_thinking = True
+
+        return emitted
+
+
+def _split_safe_suffix(text: str, sentinel: str) -> tuple[str, str]:
+    """Split *text* into (safe_to_emit, suffix_that_may_prefix_sentinel)."""
+    max_prefix = min(len(text), len(sentinel) - 1)
+    for prefix_len in range(max_prefix, 0, -1):
+        if sentinel.startswith(text[-prefix_len:]):
+            return text[:-prefix_len], text[-prefix_len:]
+    return text, ""
+
 
 def strip_think_tags(text: str) -> str:
     """Remove ``<think>...</think>`` reasoning blocks from model output.

nemo_retriever/src/nemo_retriever/service/client.py

@@ -245,6 +245,58 @@ def query(self, query: str | list[str], *, top_k: int) -> list[list[dict[str, An
         except (ValidationError, ValueError) as exc:
             raise RuntimeError(f"Service query returned invalid response: {exc}") from exc
 
+    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 = ""
+
     # ------------------------------------------------------------------
     # Job lifecycle
     # ------------------------------------------------------------------