fix: avoid mid-word splits in text_transforms.replace

livekit-agents/livekit/agents/voice/transcription/text_transforms.py

@@ -67,34 +67,59 @@ def replace(
 ) -> Callable[[AsyncIterable[str]], AsyncIterable[str]]:
     """Create a text transform that replaces words with new words.
 
-    Buffers enough text to handle terms that might be split across token boundaries
-    during streaming.
+    Buffers only a trailing key-prefix while streaming, so a key split across
+    chunks is still matched but ordinary words are never split.
 
     Args:
         replacements: A dictionary mapping words to their new words.
         case_sensitive: Whether to match the case of the words.
     Returns:
         A callable that can be used as a ``TextTransforms`` entry.
     """
-    tail_len = max(len(k) for k in replacements) - 1 if replacements else 0
     flags = re.IGNORECASE if not case_sensitive else 0
+    keys = list(replacements)
+    lookup = replacements if case_sensitive else {k.lower(): v for k, v in replacements.items()}
+    # single-pass alternation, longest key first so overlaps prefer the longer match
+    pattern = (
+        re.compile("|".join(re.escape(k) for k in sorted(keys, key=len, reverse=True)), flags)
+        if keys
+        else None
+    )
+    # longest end-anchored run that is a proper prefix of a key (a later chunk may complete it)
+    max_prefix = max((len(k) - 1 for k in keys), default=0)
+    prefixes = {k[:n] for k in keys for n in range(1, len(k))}
+    holdback_re = (
+        re.compile("(?:" + "|".join(re.escape(p) for p in prefixes) + r")\Z", flags)
+        if prefixes
+        else None
+    )
+
+    def _apply(text: str) -> str:
+        if pattern is None:
+            return text
+        # map each match to its value literally (no backreference expansion)
+        return pattern.sub(
+            lambda m: lookup[m.group(0) if case_sensitive else m.group(0).lower()], text
+        )
+
+    def _holdback(buffer: str) -> int:
+        if holdback_re is None:
+            return 0
+        # only the last max_prefix chars can begin a match
+        m = holdback_re.search(buffer[-max_prefix:])
+        return len(m.group()) if m else 0
 
     async def _transform(text: AsyncIterable[str]) -> AsyncIterable[str]:
         buffer = ""
         async for chunk in text:
-            buffer += chunk
-            if len(buffer) <= tail_len:
-                continue
-            for old, new in replacements.items():
-                buffer = re.sub(re.escape(old), lambda _, r=new: r, buffer, flags=flags)
-
-            flush_to = len(buffer) - tail_len
-            yield buffer[:flush_to]
-            buffer = buffer[flush_to:]
+            # substitute complete matches, then hold back a trailing partial-key run
+            buffer = _apply(buffer + chunk)
+            flush_to = len(buffer) - _holdback(buffer)
+            if flush_to > 0:
+                yield buffer[:flush_to]
+                buffer = buffer[flush_to:]
 
         if buffer:
-            for old, new in replacements.items():
-                buffer = re.sub(re.escape(old), lambda _, r=new: r, buffer, flags=flags)
             yield buffer
 
     return _transform

tests/test_transcription_filter.py

@@ -324,6 +324,10 @@ async def _collect(stream) -> str:
     return result
 
 
+async def _collect_chunks(stream) -> list[str]:
+    return [chunk async for chunk in stream]
+
+
 @pytest.mark.parametrize("chunk_size", [1, 2, 5, 11, 50])
 async def test_replace_across_chunk_sizes(chunk_size: int):
     """Test replacement with multiple keys, case insensitivity, and boundary spanning."""
@@ -367,6 +371,39 @@ async def test_replace_edge_cases():
     assert await _collect(transform(_stream_text("a word here", 2))) == r"a \1 \n \t here"
 
 
+async def test_replace_flushes_non_prefix_text_immediately():
+    """Text that cannot begin any key is emitted whole, not held or split mid-word."""
+    transform = replace({"LiveKit": "Lyve Kit"})
+    chunks = await _collect_chunks(transform(_stream_text("you connect.", 100)))
+    assert chunks == ["you connect."]
+
+
+async def test_replace_holds_only_potential_prefix():
+    """A trailing run that is a prefix of a key is held until it can be resolved."""
+    transform = replace({"LiveKit": "Lyve Kit"})
+    chunks = await _collect_chunks(transform(_stream_text("visit Live", 100)))
+    # "visit " flushes immediately; only "Live" (a LiveKit prefix) is held to the end
+    assert chunks[0] == "visit "
+    assert "".join(chunks) == "visit Live"
+
+
+async def test_replace_prefers_longest_overlapping_key():
+    """Overlapping keys resolve to the longest match, not dict/insertion order.
+
+    Longest-match holds when the longer key is buffered together; a key split
+    mid-token across chunks falls back to the shorter match (same as before —
+    correcting that needs incremental leftmost-longest matching).
+    """
+    transform = replace({"a": "X", "ab": "Y"})
+    assert await _collect(transform(_stream_text("ab", 100))) == "Y"
+
+
+async def test_replace_does_not_cascade():
+    """A replacement's output is not re-matched against other keys (single pass)."""
+    transform = replace({"a": "b", "b": "c"})
+    assert await _collect(transform(_stream_text("a", 100))) == "b"
+
+
 async def test_apply_text_transforms_with_callable():
     """Test _apply_text_transforms with callable, mixed transforms, and invalid input."""
     # callable only