Merge pull request #3080 from pipecat-ai/pk/assistant-aggregator-handles-mixed-includes-inter-frame-spaces-text

`LLMAssistantAggregator` now properly aggregates text that might be a…

src/pipecat/processors/aggregators/llm_response_universal.py

@@ -66,7 +66,7 @@ from pipecat.processors.aggregators.llm_response import (
     LLMUserAggregatorParams,
 )
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
-from pipecat.utils.string import concatenate_aggregated_text
+from pipecat.utils.string import TextPartForConcatenation, concatenate_aggregated_text
 from pipecat.utils.time import time_now_iso8601
 
 
@@ -90,15 +90,7 @@ class LLMContextAggregator(FrameProcessor):
         self._context = context
         self._role = role
 
-        self._aggregation: List[str] = []
+        self._aggregation: List[TextPartForConcatenation] = []
-
-        # Whether to add spaces between text parts.
-        # (Currently only used by LLMAssistantAggregator, but could be expanded
-        # to LLMUserAggregator in the future if needed; that would require
-        # additional work since LLMUserAggregator currently trims spaces from
-        # incoming frames before determining whether it "really" received any
-        # text).
-        self._add_spaces = True
 
     @property
     def messages(self) -> List[LLMContextMessage]:
@@ -191,7 +183,7 @@ class LLMContextAggregator(FrameProcessor):
         Returns:
             The concatenated aggregation string.
         """
-        return concatenate_aggregated_text(self._aggregation, self._add_spaces)
+        return concatenate_aggregated_text(self._aggregation)
 
 
 class LLMUserAggregator(LLMContextAggregator):
@@ -441,7 +433,12 @@ class LLMUserAggregator(LLMContextAggregator):
         if not text.strip():
             return
 
-        self._aggregation.append(text)
+        # Transcriptions never include inter-part spaces (so far).
+        self._aggregation.append(
+            TextPartForConcatenation(
+                text, includes_inter_part_spaces=frame.includes_inter_frame_spaces
+            )
+        )
         # We just got a final result, so let's reset interim results.
         self._seen_interim_results = False
         # Reset aggregation timer.
@@ -850,11 +847,11 @@ class LLMAssistantAggregator(LLMContextAggregator):
         if len(frame.text) == 0:
             return
 
-        # Track whether we need to add spaces between text parts
+        self._aggregation.append(
-        # Assumption: we can just keep track of the latest frame's value
+            TextPartForConcatenation(
-        self._add_spaces = not frame.includes_inter_frame_spaces
+                frame.text, includes_inter_part_spaces=frame.includes_inter_frame_spaces
-
+            )
-        self._aggregation.append(frame.text)
+        )
 
     def _context_updated_task_finished(self, task: asyncio.Task):
         self._context_updated_tasks.discard(task)

src/pipecat/processors/transcript_processor.py

@@ -26,7 +26,7 @@ from pipecat.frames.frames import (
     TTSTextFrame,
 )
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
-from pipecat.utils.string import concatenate_aggregated_text
+from pipecat.utils.string import TextPartForConcatenation, concatenate_aggregated_text
 from pipecat.utils.time import time_now_iso8601
 
 
@@ -98,15 +98,9 @@ class AssistantTranscriptProcessor(BaseTranscriptProcessor):
             **kwargs: Additional arguments passed to parent class.
         """
         super().__init__(**kwargs)
-        self._current_text_parts: List[str] = []
+        self._current_text_parts: List[TextPartForConcatenation] = []
         self._aggregation_start_time: Optional[str] = None
 
-        # Whether to add spaces between text parts.
-        # (The use of this could be expanded to the UserTranscriptProcessor in
-        # the future if needed; currently the UserTranscriptProcessor assumes
-        # that user transcription frames do not need aggregation).
-        self._add_spaces = True
-
     async def _emit_aggregated_text(self):
         """Aggregates and emits text fragments as a transcript message.
 
@@ -147,7 +141,7 @@ class AssistantTranscriptProcessor(BaseTranscriptProcessor):
                 Result: "Hello there how are you"
         """
         if self._current_text_parts and self._aggregation_start_time:
-            content = concatenate_aggregated_text(self._current_text_parts, self._add_spaces)
+            content = concatenate_aggregated_text(self._current_text_parts)
             if content:
                 logger.trace(f"Emitting aggregated assistant message: {content}")
                 message = TranscriptionMessage(
@@ -191,11 +185,11 @@ class AssistantTranscriptProcessor(BaseTranscriptProcessor):
             if not self._aggregation_start_time:
                 self._aggregation_start_time = time_now_iso8601()
 
-            # Track whether we need to add spaces between text parts
+            self._current_text_parts.append(
-            # Assumption: we can just keep track of the latest frame's value
+                TextPartForConcatenation(
-            self._add_spaces = not frame.includes_inter_frame_spaces
+                    frame.text, includes_inter_part_spaces=frame.includes_inter_frame_spaces
-
+                )
-            self._current_text_parts.append(frame.text)
+            )
 
             # Push frame.
             await self.push_frame(frame, direction)

src/pipecat/utils/string.py

@@ -18,6 +18,7 @@ Dependencies:
 """
 
 import re
+from dataclasses import dataclass
 from typing import FrozenSet, List, Optional, Sequence, Tuple
 
 import nltk
@@ -198,7 +199,24 @@ def parse_start_end_tags(
     return (None, current_tag_index)
 
 
-def concatenate_aggregated_text(text_parts: List[str], add_spaces: bool) -> str:
+@dataclass
+class TextPartForConcatenation:
+    """Class representing a part of text for concatenation with concatenate_aggregated_text.
+
+    Attributes:
+        text: The text content.
+        includes_inter_part_spaces: Whether any necessary inter-frame
+            (leading/trailing) spaces are already included in the text.
+    """
+
+    text: str
+    includes_inter_part_spaces: bool
+
+    def __str__(self):
+        return f"{self.name}(text: [{self.text}], includes_inter_part_spaces: {self.includes_inter_part_spaces})"
+
+
+def concatenate_aggregated_text(text_parts: List[TextPartForConcatenation]) -> str:
     """Concatenate a list of text parts into a single string.
 
     This function joins the provided list of text parts into a single string,
@@ -208,15 +226,55 @@ def concatenate_aggregated_text(text_parts: List[str], add_spaces: bool) -> str:
     transcription services.
 
     Args:
-        text_parts: A list of strings representing parts of text to concatenate.
+        text_parts: A list of text parts to concatenate.
-        add_spaces: Whether to add spaces between text parts during concatenation.
 
     Returns:
         A single concatenated string.
     """
-    # Concatenate text parts with or without spaces based on the flag
+    result = ""
-    separator = " " if add_spaces else ""
+    last_includes_inter_part_spaces = False
-    result = separator.join(text_parts)
+
+    if not text_parts:
+        return result
+
+    def append_part(part: TextPartForConcatenation):
+        nonlocal result
+        nonlocal last_includes_inter_part_spaces
+        result += part.text
+        last_includes_inter_part_spaces = part.includes_inter_part_spaces
+
+    for part in text_parts:
+        # Part is empty.
+        # Skip.
+        if not part.text:
+            continue
+
+        # Result is as yet empty.
+        # Just append.
+        if not result:
+            append_part(part)
+            continue
+
+        if part.includes_inter_part_spaces and last_includes_inter_part_spaces:
+            # This part is part of an ongoing run that has spaces already included.
+            # Just append.
+            append_part(part)
+        elif not part.includes_inter_part_spaces and not last_includes_inter_part_spaces:
+            # This part is part of an ongoing run that has no spaces included.
+            # Add a space before appending.
+            result += " "
+            append_part(part)
+        else:
+            # This part represents a transition to a new run (spaces -> no spaces, or vice versa).
+            # Add a space if needed, before appending.
+            if not result[-1].isspace() and not part.text[0].isspace():
+                result += " "
+            append_part(part)
+
+    # NOTE: the above logic assumes that runs of text parts with
+    # includes_inter_part_spaces=True are well-formed, i.e. they're not
+    # actually multiple separate runs with a space-less boundary, like
+    # "hello ", "world.", "goodnight ", "moon."
 
     # Clean up any excessive whitespace
     result = result.strip()

tests/test_context_aggregators.py

@@ -1005,3 +1005,53 @@ class TestLLMAssistantAggregator(
     ) -> Optional[LLMAssistantAggregatorParams]:
         kwargs.pop("expect_stripped_words", None)
         return LLMAssistantAggregatorParams(**kwargs) if kwargs else None
+
+    async def test_multiple_text_mixed(self):
+        assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
+        assert self.AGGREGATOR_CLASS is not None, "AGGREGATOR_CLASS must be set in a subclass"
+
+        context = self.CONTEXT_CLASS()
+        aggregator = self.AGGREGATOR_CLASS(
+            context, params=self.create_assistant_aggregator_params(expect_stripped_words=False)
+        )
+
+        # The newer LLMAssistantAggregator expects TextFrames to declare
+        # when they include inter-frame spaces.
+        def make_text_frame(text: str, includes_spaces: bool) -> TextFrame:
+            frame = TextFrame(text=text)
+            frame.includes_inter_frame_spaces = includes_spaces
+            return frame
+
+        frames_to_send = [
+            LLMFullResponseStartFrame(),
+            make_text_frame("Hello ", includes_spaces=True),
+            make_text_frame("Pipecat. ", includes_spaces=True),
+            make_text_frame("Here's some", includes_spaces=True),
+            make_text_frame(
+                " code:", includes_spaces=True
+            ),  # Validates ending includes_inter_frame_spaces run with no space
+            make_text_frame("```python\nprint('Hello, World!')\n```", includes_spaces=False),
+            make_text_frame(
+                "```javascript\nconsole.log('Hello, World!');\n```", includes_spaces=False
+            ),
+            make_text_frame(
+                " And some more: ", includes_spaces=True
+            ),  # Validates starting includes_inter_frame_spaces run with a space and ending it with no space
+            make_text_frame("```html\n<div>Hello, World!</div>\n```", includes_spaces=False),
+            make_text_frame(
+                "Hope that ", includes_spaces=True
+            ),  # Validates starting includes_inter_frame_spaces run with no space
+            make_text_frame("helps!", includes_spaces=True),
+            LLMFullResponseEndFrame(),
+        ]
+        expected_down_frames = [*self.EXPECTED_CONTEXT_FRAMES]
+        await run_test(
+            aggregator,
+            frames_to_send=frames_to_send,
+            expected_down_frames=expected_down_frames,
+        )
+        self.check_message_content(
+            context,
+            0,
+            "Hello Pipecat. Here's some code: ```python\nprint('Hello, World!')\n``` ```javascript\nconsole.log('Hello, World!');\n``` And some more: ```html\n<div>Hello, World!</div>\n``` Hope that helps!",
+        )