260 lines
10 KiB
Python
260 lines
10 KiB
Python
#
|
|
# Copyright (c) 2024-2026, Daily
|
|
#
|
|
# SPDX-License-Identifier: BSD 2-Clause License
|
|
#
|
|
|
|
import unittest
|
|
from unittest.mock import AsyncMock
|
|
|
|
from pipecat.utils.text.pattern_pair_aggregator import (
|
|
MatchAction,
|
|
PatternMatch,
|
|
PatternPairAggregator,
|
|
)
|
|
|
|
|
|
class TestPatternPairAggregator(unittest.IsolatedAsyncioTestCase):
|
|
def setUp(self):
|
|
self.aggregator = PatternPairAggregator()
|
|
self.test_handler = AsyncMock()
|
|
self.code_handler = AsyncMock()
|
|
|
|
# Add a test pattern
|
|
self.aggregator.add_pattern(
|
|
type="test_pattern",
|
|
start_pattern="<test>",
|
|
end_pattern="</test>",
|
|
)
|
|
self.aggregator.add_pattern(
|
|
type="code_pattern",
|
|
start_pattern="<code>",
|
|
end_pattern="</code>",
|
|
action=MatchAction.AGGREGATE,
|
|
)
|
|
|
|
# Register the mock handler
|
|
self.aggregator.on_pattern_match("test_pattern", self.test_handler)
|
|
self.aggregator.on_pattern_match("code_pattern", self.code_handler)
|
|
|
|
async def test_pattern_match_and_removal(self):
|
|
text = "Hello <test>pattern content</test>!"
|
|
results = [result async for result in self.aggregator.aggregate(text)]
|
|
|
|
# Verify the handler was called with correct PatternMatch object
|
|
self.test_handler.assert_called_once()
|
|
call_args = self.test_handler.call_args[0][0]
|
|
self.assertIsInstance(call_args, PatternMatch)
|
|
self.assertEqual(call_args.type, "test_pattern")
|
|
self.assertEqual(call_args.full_match, "<test>pattern content</test>")
|
|
self.assertEqual(call_args.text, "pattern content")
|
|
|
|
# No results yet (waiting for lookahead after "!")
|
|
self.assertEqual(len(results), 0)
|
|
|
|
# Next sentence should provide the lookahead and trigger the previous sentence
|
|
async for result in self.aggregator.aggregate(" This is another sentence."):
|
|
results.append(result)
|
|
|
|
# First result should be "Hello !" triggered by the space lookahead
|
|
self.assertEqual(len(results), 1)
|
|
self.assertEqual(results[0].text, "Hello !")
|
|
self.assertEqual(results[0].type, "sentence")
|
|
|
|
# Now flush to get the remaining sentence
|
|
result = await self.aggregator.flush()
|
|
self.assertEqual(result.text, "This is another sentence.")
|
|
|
|
# Buffer should be empty after returning a complete sentence
|
|
self.assertEqual(self.aggregator.text.text, "")
|
|
|
|
async def test_pattern_match_and_aggregate(self):
|
|
text = "Here is code <code>pattern content</code> This is another sentence."
|
|
|
|
results = [result async for result in self.aggregator.aggregate(text)]
|
|
|
|
# First result should be "Here is code" when pattern starts
|
|
self.assertEqual(results[0].text, "Here is code")
|
|
self.assertEqual(results[0].type, "sentence")
|
|
|
|
# Second result should be the code pattern content
|
|
self.assertEqual(results[1].text, "pattern content")
|
|
self.assertEqual(results[1].type, "code_pattern")
|
|
|
|
# Verify the handler was called with correct PatternMatch object
|
|
self.code_handler.assert_called_once()
|
|
call_args = self.code_handler.call_args[0][0]
|
|
self.assertIsInstance(call_args, PatternMatch)
|
|
self.assertEqual(call_args.type, "code_pattern")
|
|
self.assertEqual(call_args.full_match, "<code>pattern content</code>")
|
|
self.assertEqual(call_args.text, "pattern content")
|
|
|
|
# Last sentence needs flush (waiting for lookahead after ".")
|
|
result = await self.aggregator.flush()
|
|
self.assertEqual(result.text, "This is another sentence.")
|
|
self.assertEqual(result.type, "sentence")
|
|
|
|
# Buffer should be empty after returning a complete sentence
|
|
self.assertEqual(self.aggregator.text.text, "")
|
|
|
|
async def test_incomplete_pattern(self):
|
|
text = "Hello <test>pattern content"
|
|
results = [result async for result in self.aggregator.aggregate(text)]
|
|
# No complete pattern yet, so nothing should be returned
|
|
self.assertEqual(len(results), 0)
|
|
|
|
# The handler should not be called yet
|
|
self.test_handler.assert_not_called()
|
|
|
|
# Buffer should contain the incomplete text
|
|
self.assertEqual(self.aggregator.text.text, "Hello <test>pattern content")
|
|
self.assertEqual(self.aggregator.text.type, "test_pattern")
|
|
|
|
# Reset and confirm buffer is cleared
|
|
await self.aggregator.reset()
|
|
self.assertEqual(self.aggregator.text.text, "")
|
|
|
|
async def test_multiple_patterns(self):
|
|
# Set up multiple patterns and handlers
|
|
voice_handler = AsyncMock()
|
|
emphasis_handler = AsyncMock()
|
|
|
|
self.aggregator.add_pattern(
|
|
type="voice",
|
|
start_pattern="<voice>",
|
|
end_pattern="</voice>",
|
|
action=MatchAction.REMOVE,
|
|
)
|
|
|
|
self.aggregator.add_pattern(
|
|
type="emphasis",
|
|
start_pattern="<em>",
|
|
end_pattern="</em>",
|
|
action=MatchAction.KEEP, # Keep emphasis tags
|
|
)
|
|
|
|
self.aggregator.on_pattern_match("voice", voice_handler)
|
|
self.aggregator.on_pattern_match("emphasis", emphasis_handler)
|
|
|
|
text = "Hello <voice>female</voice> I am <em>very</em> excited to meet you!"
|
|
results = [result async for result in self.aggregator.aggregate(text)]
|
|
|
|
# Both handlers should be called with correct data
|
|
voice_handler.assert_called_once()
|
|
voice_match = voice_handler.call_args[0][0]
|
|
self.assertEqual(voice_match.type, "voice")
|
|
self.assertEqual(voice_match.text, "female")
|
|
|
|
emphasis_handler.assert_called_once()
|
|
emphasis_match = emphasis_handler.call_args[0][0]
|
|
self.assertEqual(emphasis_match.type, "emphasis")
|
|
self.assertEqual(emphasis_match.text, "very")
|
|
|
|
# With lookahead, we need to flush to get the final sentence
|
|
self.assertEqual(len(results), 0) # Waiting for lookahead after "!"
|
|
|
|
result = await self.aggregator.flush()
|
|
# Voice pattern should be removed, emphasis pattern should remain
|
|
self.assertEqual(result.text, "Hello I am <em>very</em> excited to meet you!")
|
|
|
|
# Buffer should be empty
|
|
self.assertEqual(self.aggregator.text.text, "")
|
|
|
|
async def test_handle_interruption(self):
|
|
text = "Hello <test>pattern"
|
|
results = [result async for result in self.aggregator.aggregate(text)]
|
|
self.assertEqual(len(results), 0)
|
|
|
|
# Simulate interruption
|
|
await self.aggregator.handle_interruption()
|
|
|
|
# Buffer should be cleared
|
|
self.assertEqual(self.aggregator.text.text, "")
|
|
|
|
# Handler should not have been called
|
|
self.test_handler.assert_not_called()
|
|
|
|
async def test_pattern_across_sentences(self):
|
|
text = "Hello <test>This is sentence one. This is sentence two.</test> Final sentence."
|
|
results = [result async for result in self.aggregator.aggregate(text)]
|
|
|
|
# Handler should be called with entire content
|
|
self.test_handler.assert_called_once()
|
|
call_args = self.test_handler.call_args[0][0]
|
|
self.assertEqual(call_args.text, "This is sentence one. This is sentence two.")
|
|
|
|
# With lookahead, we need to flush to get the final sentence
|
|
self.assertEqual(len(results), 0) # Waiting for lookahead after "."
|
|
|
|
result = await self.aggregator.flush()
|
|
# Pattern should be removed, resulting in text with sentences merged
|
|
self.assertEqual(result.text, "Hello Final sentence.")
|
|
|
|
# Buffer should be empty
|
|
self.assertEqual(self.aggregator.text.text, "")
|
|
|
|
|
|
class TestPatternPairAggregatorTokenMode(unittest.IsolatedAsyncioTestCase):
|
|
def setUp(self):
|
|
from pipecat.utils.text.base_text_aggregator import AggregationType
|
|
|
|
self.aggregator = PatternPairAggregator(aggregation_type=AggregationType.TOKEN)
|
|
self.handler = AsyncMock()
|
|
self.aggregator.add_pattern(
|
|
type="think",
|
|
start_pattern="<think>",
|
|
end_pattern="</think>",
|
|
action=MatchAction.REMOVE,
|
|
)
|
|
self.aggregator.on_pattern_match("think", self.handler)
|
|
|
|
async def test_token_no_patterns(self):
|
|
"""Non-pattern text passes through as TOKEN, one per aggregate call."""
|
|
results = []
|
|
for token in ["Hello", " world", "."]:
|
|
async for r in self.aggregator.aggregate(token):
|
|
results.append(r)
|
|
|
|
self.assertEqual(len(results), 3)
|
|
self.assertEqual(results[0].text, "Hello")
|
|
self.assertEqual(results[1].text, " world")
|
|
self.assertEqual(results[2].text, ".")
|
|
for r in results:
|
|
self.assertEqual(r.type, "token")
|
|
|
|
async def test_token_pattern_detection(self):
|
|
"""Pattern detection still works with word-by-word token delivery."""
|
|
results = []
|
|
for token in ["Hi ", "<think>", "secret", "</think>", " bye"]:
|
|
async for r in self.aggregator.aggregate(token):
|
|
results.append(r)
|
|
|
|
# Handler called once when the pattern completes
|
|
self.handler.assert_called_once()
|
|
call_args = self.handler.call_args[0][0]
|
|
self.assertEqual(call_args.text, "secret")
|
|
|
|
# "Hi " yields before pattern starts, pattern is removed, " bye" yields after
|
|
self.assertEqual(len(results), 2)
|
|
self.assertEqual(results[0].text, "Hi ")
|
|
self.assertEqual(results[0].type, "token")
|
|
self.assertEqual(results[1].text, " bye")
|
|
self.assertEqual(results[1].type, "token")
|
|
|
|
async def test_token_incomplete_pattern_buffers(self):
|
|
"""Incomplete pattern is buffered across calls, not leaked to output."""
|
|
results = []
|
|
for token in ["Hi ", "<think>", "partial"]:
|
|
async for r in self.aggregator.aggregate(token):
|
|
results.append(r)
|
|
|
|
# Only "Hi " should be yielded; "<think>partial" stays buffered
|
|
self.assertEqual(len(results), 1)
|
|
self.assertEqual(results[0].text, "Hi ")
|
|
self.assertEqual(results[0].type, "token")
|
|
self.handler.assert_not_called()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|