Refactor text chunking implementation and update configuration

- Moved SentenceTextChunker and SentenceTextChunkerConfig to the utils module for better organization. - Updated pytest.ini to include the current directory in the Python path. - Added a new utils module with shared utility helpers. - Adjusted import paths in the test files to reflect the new location of text chunking classes.
2026-06-18 10:15:07 +08:00
parent a6777a827b
commit cf0a8b71fd
5 changed files with 5 additions and 3 deletions
--- a/src/utils/text_chunker.py
+++ b/src/utils/text_chunker.py
@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+SENTENCE_ENDING_PUNCTUATION = frozenset(".!?;。！？；")
+SOFT_BREAK_PUNCTUATION = frozenset(",，、:：")
+CLOSING_PUNCTUATION = frozenset("\"'”’)]}）】》」』")
+
+
+@dataclass(frozen=True)
+class SentenceTextChunkerConfig:
+    """Configuration for streaming text chunks sent to TTS."""
+
+    min_chars: int = 1
+    max_chars: int = 80
+    use_soft_breaks: bool = True
+
+
+class SentenceTextChunker:
+    """Lightweight Pipecat-style sentence chunker for streaming TTS text.
+
+    The chunker waits for one non-whitespace lookahead character after sentence
+    punctuation before emitting a sentence. This avoids splitting too early when
+    a punctuation mark arrives at the end of a stream token.
+    """
+
+    def __init__(self, config: SentenceTextChunkerConfig | None = None) -> None:
+        self._config = config or SentenceTextChunkerConfig()
+        self._buffer = ""
+        self._needs_lookahead = False
+        self._last_soft_break = -1
+
+    @property
+    def text(self) -> str:
+        return self._buffer.strip(" ")
+
+    def feed(self, text: str) -> list[str]:
+        """Append streaming text and return chunks ready for TTS."""
+        chunks: list[str] = []
+        if not text:
+            return chunks
+
+        for char in text:
+            self._buffer += char
+            index = len(self._buffer) - 1
+
+            if char in SOFT_BREAK_PUNCTUATION:
+                self._last_soft_break = index + 1
+
+            if self._needs_lookahead and char.strip():
+                self._needs_lookahead = False
+                chunk = self._pop_sentence_chunk()
+                if chunk:
+                    chunks.append(chunk)
+                    continue
+
+            if char in SENTENCE_ENDING_PUNCTUATION and not self._is_decimal_point(index):
+                self._needs_lookahead = True
+
+            chunk = self._pop_soft_chunk_if_needed()
+            if chunk:
+                chunks.append(chunk)
+
+        return chunks
+
+    def flush(self) -> str | None:
+        """Return any remaining buffered text at end of stream."""
+        if not self._buffer:
+            return None
+        chunk = self._buffer.strip(" ")
+        self.reset()
+        return chunk or None
+
+    def reset(self) -> None:
+        self._buffer = ""
+        self._needs_lookahead = False
+        self._last_soft_break = -1
+
+    def _pop_sentence_chunk(self) -> str | None:
+        end = self._sentence_end_index()
+        if end is None:
+            return None
+        chunk = self._buffer[:end].strip(" ")
+        if len(chunk) < self._config.min_chars:
+            return None
+        self._buffer = self._buffer[end:]
+        self._last_soft_break = self._find_last_soft_break()
+        return chunk
+
+    def _sentence_end_index(self) -> int | None:
+        index = 0
+        while index < len(self._buffer):
+            char = self._buffer[index]
+            if char in SENTENCE_ENDING_PUNCTUATION and not self._is_decimal_point(index):
+                end = index + 1
+                while end < len(self._buffer) and self._buffer[end] in CLOSING_PUNCTUATION:
+                    end += 1
+                return end
+            index += 1
+        return None
+
+    def _pop_soft_chunk_if_needed(self) -> str | None:
+        if (
+            not self._config.use_soft_breaks
+            or self._config.max_chars <= 0
+            or len(self._buffer) < self._config.max_chars
+            or self._last_soft_break <= 0
+        ):
+            return None
+
+        chunk = self._buffer[: self._last_soft_break].strip(" ")
+        if len(chunk) < self._config.min_chars:
+            return None
+
+        self._buffer = self._buffer[self._last_soft_break :]
+        self._last_soft_break = self._find_last_soft_break()
+        self._needs_lookahead = False
+        return chunk
+
+    def _find_last_soft_break(self) -> int:
+        for index in range(len(self._buffer) - 1, -1, -1):
+            if self._buffer[index] in SOFT_BREAK_PUNCTUATION:
+                return index + 1
+        return -1
+
+    def _is_decimal_point(self, index: int) -> bool:
+        if self._buffer[index] != ".":
+            return False
+        return (
+            index > 0
+            and index + 1 < len(self._buffer)
+            and self._buffer[index - 1].isdigit()
+            and self._buffer[index + 1].isdigit()
+        )