Merge pull request #11 from InterfazeAI/fix/stream-tag-filter

fix: strip <think>/<precontext> from live stream deltas (#4)

Author: Khurdhula-Harshavardhan

README.md

@@ -78,18 +78,23 @@ res = interfaze.chat.completions.create(
 
 ## Streaming
 
+For live rendering, iterate `text_deltas()` — it yields visible text only, stripping Interfaze's
+inline `<think>`/`<precontext>` side-channels:
+
 ```python
 stream = interfaze.chat.completions.stream(
     messages=[{"role": "user", "content": "Tell me a story."}],
 )
-for chunk in stream:
-    print(chunk.choices[0].delta.content or "", end="")
+for text in stream.text_deltas():
+    print(text, end="")
 final = stream.get_final_completion()
 print(final.reasoning, final.precontext)
 ```
 
-> Plain `create(stream=True)` also works and returns the raw chunk iterator; `.stream()` adds
-> accumulation and surfaces `reasoning`/`precontext`.
+> Iterating the stream directly (`for chunk in stream`) or the plain `create(stream=True)` path
+> yields **raw** chunks whose `delta.content` still contains the `<think>`/`<precontext>` tags — use
+> `text_deltas()` for anything user-facing. `.stream()` also accumulates and surfaces
+> `reasoning`/`precontext` on `get_final_completion()`.
 
 ## Inputs
 

src/interfaze/_stream.py

@@ -30,6 +30,67 @@ def strip_side_channels(content: str) -> Tuple[str, Optional[str], Optional[List
     return text.strip(), ("\n".join(thinks) if thinks else None), (pre or None)
 
 
+_SIDE_OPEN = ("<think>", "<precontext>")
+_SIDE_CLOSE = {"<think>": "</think>", "<precontext>": "</precontext>"}
+
+
+def _suffix_prefix_len(s: str, tag: str) -> int:
+    for k in range(min(len(s), len(tag) - 1), 0, -1):
+        if s[-k:] == tag[:k]:
+            return k
+    return 0
+
+
+class _SideChannelFilter:
+    """Strip inline ``<think>``/``<precontext>`` blocks from streamed content, chunk by chunk.
+
+    Buffers a trailing partial that may be a split tag; never withholds text that cannot be a tag.
+    """
+
+    def __init__(self) -> None:
+        self._buf = ""
+        self._close: Optional[str] = None
+
+    def feed(self, text: str) -> str:
+        self._buf += text
+        out: List[str] = []
+        while self._buf:
+            if self._close is None:
+                lt = self._buf.find("<")
+                if lt == -1:
+                    out.append(self._buf)
+                    self._buf = ""
+                    break
+                if lt:
+                    out.append(self._buf[:lt])
+                    self._buf = self._buf[lt:]
+                opened = next((t for t in _SIDE_OPEN if self._buf.startswith(t)), None)
+                if opened:
+                    self._close = _SIDE_CLOSE[opened]
+                    self._buf = self._buf[len(opened) :]
+                    continue
+                if any(t.startswith(self._buf) for t in _SIDE_OPEN):
+                    break
+                out.append("<")
+                self._buf = self._buf[1:]
+            else:
+                end = self._buf.find(self._close)
+                if end == -1:
+                    keep = _suffix_prefix_len(self._buf, self._close)
+                    self._buf = self._buf[len(self._buf) - keep :] if keep else ""
+                    break
+                self._buf = self._buf[end + len(self._close) :]
+                self._close = None
+        return "".join(out)
+
+    def flush(self) -> str:
+        if self._close is not None:
+            self._buf = ""
+            return ""
+        rest, self._buf = self._buf, ""
+        return rest
+
+
 class _State:
     def __init__(self) -> None:
         self.content = ""
@@ -117,6 +178,29 @@ def __iter__(self) -> "Iterator[ChatCompletionChunk]":
             yield chunk
         self._done = True
 
+    def text_deltas(self) -> "Iterator[str]":
+        """Yield visible text only, stripping ``<think>``/``<precontext>`` across chunk boundaries.
+
+        Use this (not raw ``create(stream=True)`` deltas) for live rendering. ``reasoning`` and
+        ``precontext`` remain available on ``get_final_completion()``.
+        """
+        if self._started:
+            raise InterfazeError("This stream has already been consumed.")
+        self._started = True
+        filt = _SideChannelFilter()
+        for chunk in self._client.chat.completions.create(stream=True, **self._kwargs):
+            self._state.accumulate(chunk)
+            if chunk.choices:
+                delta = chunk.choices[0].delta
+                if delta and isinstance(delta.content, str) and delta.content:
+                    visible = filt.feed(delta.content)
+                    if visible:
+                        yield visible
+        tail = filt.flush()
+        if tail:
+            yield tail
+        self._done = True
+
     @property
     def text(self) -> str:
         return strip_side_channels(self._state.content)[0]
@@ -160,6 +244,30 @@ async def __aiter__(self) -> "AsyncIterator[ChatCompletionChunk]":
             yield chunk
         self._done = True
 
+    async def text_deltas(self) -> "AsyncIterator[str]":
+        """Yield visible text only, stripping ``<think>``/``<precontext>`` across chunk boundaries.
+
+        Use this (not raw ``create(stream=True)`` deltas) for live rendering. ``reasoning`` and
+        ``precontext`` remain available on ``get_final_completion()``.
+        """
+        if self._started:
+            raise InterfazeError("This stream has already been consumed.")
+        self._started = True
+        filt = _SideChannelFilter()
+        stream = await self._client.chat.completions.create(stream=True, **self._kwargs)
+        async for chunk in stream:
+            self._state.accumulate(chunk)
+            if chunk.choices:
+                delta = chunk.choices[0].delta
+                if delta and isinstance(delta.content, str) and delta.content:
+                    visible = filt.feed(delta.content)
+                    if visible:
+                        yield visible
+        tail = filt.flush()
+        if tail:
+            yield tail
+        self._done = True
+
     async def get_final_completion(self) -> InterfazeChatCompletion:
         if not self._started:
             self._started = True

tests/test_stream.py

@@ -3,7 +3,7 @@
 import asyncio
 
 import respx
-from conftest import STREAM_CHUNKS, STREAM_THINK, mock_sse
+from conftest import STREAM_CHUNKS, STREAM_THINK, _chunk, mock_sse
 
 from interfaze import AsyncInterfaze, Interfaze
 
@@ -46,3 +46,61 @@ async def go():
     n, final = asyncio.run(go())
     assert n == len(STREAM_CHUNKS)
     assert final.precontext and final.precontext[0].name == "ocr"
+
+
+@respx.mock
+def test_text_deltas_strips_precontext():
+    mock_sse(STREAM_CHUNKS)
+    stream = Interfaze(api_key="t").chat.completions.stream(messages=[{"role": "user", "content": "x"}])
+    text = "".join(stream.text_deltas())
+    assert text == "Total is $12.34"
+    assert "<precontext>" not in text
+
+
+@respx.mock
+def test_text_deltas_strips_think():
+    mock_sse(STREAM_THINK)
+    stream = Interfaze(api_key="t").chat.completions.stream(messages=[{"role": "user", "content": "x"}])
+    text = "".join(stream.text_deltas())
+    assert text == "The sky is blue."
+    assert "<think>" not in text
+
+
+@respx.mock
+def test_text_deltas_handles_tag_split_across_chunks():
+    chunks = [
+        _chunk({"content": "Hello <pre"}),
+        _chunk({"content": 'context>[{"name":"ocr","result":1}]</precon'}),
+        _chunk({"content": "text> world"}),
+        _chunk({}, finish_reason="stop"),
+    ]
+    mock_sse(chunks)
+    stream = Interfaze(api_key="t").chat.completions.stream(messages=[{"role": "user", "content": "x"}])
+    text = "".join(stream.text_deltas())
+    assert "precontext" not in text and "ocr" not in text
+    assert text == "Hello  world"
+
+
+@respx.mock
+def test_text_deltas_preserves_literal_lt():
+    chunks = [
+        _chunk({"content": "a < b and c "}),
+        _chunk({"content": "< d"}),
+        _chunk({}, finish_reason="stop"),
+    ]
+    mock_sse(chunks)
+    stream = Interfaze(api_key="t").chat.completions.stream(messages=[{"role": "user", "content": "x"}])
+    assert "".join(stream.text_deltas()) == "a < b and c < d"
+
+
+@respx.mock
+def test_async_text_deltas():
+    mock_sse(STREAM_CHUNKS)
+
+    async def go():
+        stream = AsyncInterfaze(api_key="t").chat.completions.stream(
+            messages=[{"role": "user", "content": "x"}]
+        )
+        return "".join([t async for t in stream.text_deltas()])
+
+    assert asyncio.run(go()) == "Total is $12.34"