BaseTextAggregator: make functions async

This commit is contained in:
Aleix Conchillo Flaqué
2025-05-19 09:50:02 -07:00
parent 54388c0d9b
commit 54b1d7fcc1
10 changed files with 59 additions and 63 deletions

View File

@@ -68,6 +68,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Changed
- `BaseTextAggregator` methods `aggregate()`, `handle_interruption()` and
`reset()` are now async.
- The API version for `CartesiaTTSService` and `CartesiaHttpTTSService` has
been updated. Also, the `cartesia` dependency has been updated to 2.x.

View File

@@ -99,21 +99,14 @@ async def run_bot(webrtc_connection: SmallWebRTCConnection, _: argparse.Namespac
)
# Register handler for voice switching
def on_voice_tag(match: PatternMatch):
async def on_voice_tag(match: PatternMatch):
voice_name = match.content.strip().lower()
if voice_name in VOICE_IDS:
voice_id = VOICE_IDS[voice_name]
# Create task to reset the TTS context after voice change
async def change_voice():
# First flush any existing audio to finish the current context
await tts.flush_audio()
# Then set the new voice
tts.set_voice(voice_id)
logger.info(f"Switched to {voice_name} voice")
# Schedule the voice change task
asyncio.create_task(change_voice())
# First flush any existing audio to finish the current context
await tts.flush_audio()
# Then set the new voice
tts.set_voice(VOICE_IDS[voice_name])
logger.info(f"Switched to {voice_name} voice")
else:
logger.warning(f"Unknown voice: {voice_name}")

View File

@@ -183,7 +183,7 @@ class TTSService(AIService):
await self._maybe_pause_frame_processing()
sentence = self._text_aggregator.text
self._text_aggregator.reset()
await self._text_aggregator.reset()
self._processing_text = False
await self._push_tts_frames(sentence)
if isinstance(frame, LLMFullResponseEndFrame):
@@ -234,7 +234,7 @@ class TTSService(AIService):
async def _handle_interruption(self, frame: StartInterruptionFrame, direction: FrameDirection):
self._processing_text = False
self._text_aggregator.handle_interruption()
await self._text_aggregator.handle_interruption()
for filter in self._text_filters:
filter.handle_interruption()
@@ -251,7 +251,7 @@ class TTSService(AIService):
if not self._aggregate_sentences:
text = frame.text
else:
text = self._text_aggregator.aggregate(frame.text)
text = await self._text_aggregator.aggregate(frame.text)
if text:
await self._push_tts_frames(text)

View File

@@ -25,7 +25,7 @@ class BaseTextAggregator(ABC):
pass
@abstractmethod
def aggregate(self, text: str) -> Optional[str]:
async def aggregate(self, text: str) -> Optional[str]:
"""Aggregates the specified text with the currently accumulated text.
This method should be implemented to define how the new text contributes
@@ -43,7 +43,7 @@ class BaseTextAggregator(ABC):
pass
@abstractmethod
def handle_interruption(self):
async def handle_interruption(self):
"""Handles interruptions. When an interruption occurs it is possible
that we might want to discard the aggregated text or do some internal
modifications to the aggregated text.
@@ -52,6 +52,6 @@ class BaseTextAggregator(ABC):
pass
@abstractmethod
def reset(self):
async def reset(self):
"""Clears the internally aggregated text."""
pass

View File

@@ -5,7 +5,7 @@
#
import re
from typing import Callable, Optional, Tuple
from typing import Awaitable, Callable, Optional, Tuple
from loguru import logger
@@ -106,7 +106,7 @@ class PatternPairAggregator(BaseTextAggregator):
return self
def on_pattern_match(
self, pattern_id: str, handler: Callable[[PatternMatch], None]
self, pattern_id: str, handler: Callable[[PatternMatch], Awaitable[None]]
) -> "PatternPairAggregator":
"""Register a handler for when a pattern pair is matched.
@@ -124,7 +124,7 @@ class PatternPairAggregator(BaseTextAggregator):
self._handlers[pattern_id] = handler
return self
def _process_complete_patterns(self, text: str) -> Tuple[str, bool]:
async def _process_complete_patterns(self, text: str) -> Tuple[str, bool]:
"""Process all complete pattern pairs in the text.
Searches for all complete pattern pairs in the text, calls the
@@ -167,7 +167,7 @@ class PatternPairAggregator(BaseTextAggregator):
# Call the appropriate handler if registered
if pattern_id in self._handlers:
try:
self._handlers[pattern_id](pattern_match)
await self._handlers[pattern_id](pattern_match)
except Exception as e:
logger.error(f"Error in pattern handler for {pattern_id}: {e}")
@@ -204,7 +204,7 @@ class PatternPairAggregator(BaseTextAggregator):
return False
def aggregate(self, text: str) -> Optional[str]:
async def aggregate(self, text: str) -> Optional[str]:
"""Aggregate text and process pattern pairs.
This method adds the new text to the buffer, processes any complete pattern
@@ -223,7 +223,7 @@ class PatternPairAggregator(BaseTextAggregator):
self._text += text
# Process any complete patterns in the buffer
processed_text, modified = self._process_complete_patterns(self._text)
processed_text, modified = await self._process_complete_patterns(self._text)
# Only update the buffer if modifications were made
if modified:
@@ -245,7 +245,7 @@ class PatternPairAggregator(BaseTextAggregator):
# No complete sentence found yet
return None
def handle_interruption(self):
async def handle_interruption(self):
"""Handle interruptions by clearing the buffer.
Called when an interruption occurs in the processing pipeline,
@@ -253,7 +253,7 @@ class PatternPairAggregator(BaseTextAggregator):
"""
self._text = ""
def reset(self):
async def reset(self):
"""Clear the internally aggregated text.
Resets the aggregator to its initial state, discarding any

View File

@@ -23,7 +23,7 @@ class SimpleTextAggregator(BaseTextAggregator):
def text(self) -> str:
return self._text
def aggregate(self, text: str) -> Optional[str]:
async def aggregate(self, text: str) -> Optional[str]:
result: Optional[str] = None
self._text += text
@@ -35,8 +35,8 @@ class SimpleTextAggregator(BaseTextAggregator):
return result
def handle_interruption(self):
async def handle_interruption(self):
self._text = ""
def reset(self):
async def reset(self):
self._text = ""

View File

@@ -43,7 +43,7 @@ class SkipTagsAggregator(BaseTextAggregator):
"""
return self._text
def aggregate(self, text: str) -> Optional[str]:
async def aggregate(self, text: str) -> Optional[str]:
"""Aggregate text and process pattern pairs.
This method adds the new text to the buffer, processes any complete pattern
@@ -77,7 +77,7 @@ class SkipTagsAggregator(BaseTextAggregator):
# No complete sentence found yet
return None
def handle_interruption(self):
async def handle_interruption(self):
"""Handle interruptions by clearing the buffer.
Called when an interruption occurs in the processing pipeline,
@@ -85,7 +85,7 @@ class SkipTagsAggregator(BaseTextAggregator):
"""
self._text = ""
def reset(self):
async def reset(self):
"""Clear the internally aggregated text.
Resets the aggregator to its initial state, discarding any

View File

@@ -5,7 +5,7 @@
#
import unittest
from unittest.mock import Mock
from unittest.mock import AsyncMock
from pipecat.utils.text.pattern_pair_aggregator import PatternMatch, PatternPairAggregator
@@ -13,7 +13,7 @@ from pipecat.utils.text.pattern_pair_aggregator import PatternMatch, PatternPair
class TestPatternPairAggregator(unittest.IsolatedAsyncioTestCase):
def setUp(self):
self.aggregator = PatternPairAggregator()
self.test_handler = Mock()
self.test_handler = AsyncMock()
# Add a test pattern
self.aggregator.add_pattern_pair(
@@ -28,12 +28,12 @@ class TestPatternPairAggregator(unittest.IsolatedAsyncioTestCase):
async def test_pattern_match_and_removal(self):
# First part doesn't complete the pattern
result = self.aggregator.aggregate("Hello <test>pattern")
result = await self.aggregator.aggregate("Hello <test>pattern")
self.assertIsNone(result)
self.assertEqual(self.aggregator.text, "Hello <test>pattern")
# Second part completes the pattern and includes an exclamation point
result = self.aggregator.aggregate(" content</test>!")
result = await self.aggregator.aggregate(" content</test>!")
# Verify the handler was called with correct PatternMatch object
self.test_handler.assert_called_once()
@@ -48,7 +48,7 @@ class TestPatternPairAggregator(unittest.IsolatedAsyncioTestCase):
self.assertEqual(result, "Hello !")
# Next sentence should be processed separately
result = self.aggregator.aggregate(" This is another sentence.")
result = await self.aggregator.aggregate(" This is another sentence.")
self.assertEqual(result, " This is another sentence.")
# Buffer should be empty after returning a complete sentence
@@ -56,7 +56,7 @@ class TestPatternPairAggregator(unittest.IsolatedAsyncioTestCase):
async def test_incomplete_pattern(self):
# Add text with incomplete pattern
result = self.aggregator.aggregate("Hello <test>pattern content")
result = await self.aggregator.aggregate("Hello <test>pattern content")
# No complete pattern yet, so nothing should be returned
self.assertIsNone(result)
@@ -68,13 +68,13 @@ class TestPatternPairAggregator(unittest.IsolatedAsyncioTestCase):
self.assertEqual(self.aggregator.text, "Hello <test>pattern content")
# Reset and confirm buffer is cleared
self.aggregator.reset()
await self.aggregator.reset()
self.assertEqual(self.aggregator.text, "")
async def test_multiple_patterns(self):
# Set up multiple patterns and handlers
voice_handler = Mock()
emphasis_handler = Mock()
voice_handler = AsyncMock()
emphasis_handler = AsyncMock()
self.aggregator.add_pattern_pair(
pattern_id="voice", start_pattern="<voice>", end_pattern="</voice>", remove_match=True
@@ -92,7 +92,7 @@ class TestPatternPairAggregator(unittest.IsolatedAsyncioTestCase):
# Test with multiple patterns in one text block
text = "Hello <voice>female</voice> I am <em>very</em> excited to meet you!"
result = self.aggregator.aggregate(text)
result = await self.aggregator.aggregate(text)
# Both handlers should be called with correct data
voice_handler.assert_called_once()
@@ -113,11 +113,11 @@ class TestPatternPairAggregator(unittest.IsolatedAsyncioTestCase):
async def test_handle_interruption(self):
# Start with incomplete pattern
result = self.aggregator.aggregate("Hello <test>pattern")
result = await self.aggregator.aggregate("Hello <test>pattern")
self.assertIsNone(result)
# Simulate interruption
self.aggregator.handle_interruption()
await self.aggregator.handle_interruption()
# Buffer should be cleared
self.assertEqual(self.aggregator.text, "")
@@ -127,13 +127,13 @@ class TestPatternPairAggregator(unittest.IsolatedAsyncioTestCase):
async def test_pattern_across_sentences(self):
# Test pattern that spans multiple sentences
result = self.aggregator.aggregate("Hello <test>This is sentence one.")
result = await self.aggregator.aggregate("Hello <test>This is sentence one.")
# First sentence contains start of pattern but no end, so no complete pattern yet
self.assertIsNone(result)
# Add second part with pattern end
result = self.aggregator.aggregate(" This is sentence two.</test> Final sentence.")
result = await self.aggregator.aggregate(" This is sentence two.</test> Final sentence.")
# Handler should be called with entire content
self.test_handler.assert_called_once()

View File

@@ -14,16 +14,16 @@ class TestSimpleTextAggregator(unittest.IsolatedAsyncioTestCase):
self.aggregator = SimpleTextAggregator()
async def test_reset_aggregations(self):
assert self.aggregator.aggregate("Hello ") == None
assert await self.aggregator.aggregate("Hello ") == None
assert self.aggregator.text == "Hello "
self.aggregator.reset()
await self.aggregator.reset()
assert self.aggregator.text == ""
async def test_simple_sentence(self):
assert self.aggregator.aggregate("Hello ") == None
assert self.aggregator.aggregate("Pipecat!") == "Hello Pipecat!"
assert await self.aggregator.aggregate("Hello ") == None
assert await self.aggregator.aggregate("Pipecat!") == "Hello Pipecat!"
assert self.aggregator.text == ""
async def test_multiple_sentences(self):
assert self.aggregator.aggregate("Hello Pipecat! How are ") == "Hello Pipecat!"
assert self.aggregator.aggregate("you?") == " How are you?"
assert await self.aggregator.aggregate("Hello Pipecat! How are ") == "Hello Pipecat!"
assert await self.aggregator.aggregate("you?") == " How are you?"

View File

@@ -14,41 +14,41 @@ class TestSkipTagsAggregator(unittest.IsolatedAsyncioTestCase):
self.aggregator = SkipTagsAggregator([("<spell>", "</spell>")])
async def test_no_tags(self):
self.aggregator.reset()
await self.aggregator.reset()
# No tags involved, aggregate at end of sentence.
result = self.aggregator.aggregate("Hello Pipecat!")
result = await self.aggregator.aggregate("Hello Pipecat!")
self.assertEqual(result, "Hello Pipecat!")
self.assertEqual(self.aggregator.text, "")
async def test_basic_tags(self):
self.aggregator.reset()
await self.aggregator.reset()
# Tags involved, avoid aggregation during tags.
result = self.aggregator.aggregate("My email is <spell>foo@pipecat.ai</spell>.")
result = await self.aggregator.aggregate("My email is <spell>foo@pipecat.ai</spell>.")
self.assertEqual(result, "My email is <spell>foo@pipecat.ai</spell>.")
self.assertEqual(self.aggregator.text, "")
async def test_streaming_tags(self):
self.aggregator.reset()
await self.aggregator.reset()
# Tags involved, stream small chunk of texts.
result = self.aggregator.aggregate("My email is <sp")
result = await self.aggregator.aggregate("My email is <sp")
self.assertIsNone(result)
self.assertEqual(self.aggregator.text, "My email is <sp")
result = self.aggregator.aggregate("ell>foo.")
result = await self.aggregator.aggregate("ell>foo.")
self.assertIsNone(result)
self.assertEqual(self.aggregator.text, "My email is <spell>foo.")
result = self.aggregator.aggregate("bar@pipecat.")
result = await self.aggregator.aggregate("bar@pipecat.")
self.assertIsNone(result)
self.assertEqual(self.aggregator.text, "My email is <spell>foo.bar@pipecat.")
result = self.aggregator.aggregate("ai</spe")
result = await self.aggregator.aggregate("ai</spe")
self.assertIsNone(result)
self.assertEqual(self.aggregator.text, "My email is <spell>foo.bar@pipecat.ai</spe")
result = self.aggregator.aggregate("ll>.")
result = await self.aggregator.aggregate("ll>.")
self.assertEqual(result, "My email is <spell>foo.bar@pipecat.ai</spell>.")
self.assertEqual(self.aggregator.text, "")