LLMUserAggregator: add user turn stopped message argument
It is now possible to get the user aggregation when a `on_user_turn_stopped` event is emitted.
This commit is contained in:
@@ -102,6 +102,23 @@ class LLMAssistantAggregatorParams:
|
||||
expect_stripped_words: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class UserTurnStoppedMessage:
|
||||
"""A user turn stopped message containing a user transcript update.
|
||||
|
||||
A message in a conversation transcript containing the user content. This is
|
||||
the aggregated transcript that is then used in the context.
|
||||
|
||||
Parameters:
|
||||
content: The message content/text.
|
||||
user_id: Optional identifier for the user.
|
||||
|
||||
"""
|
||||
|
||||
content: str
|
||||
user_id: Optional[str] = None
|
||||
|
||||
|
||||
class LLMContextAggregator(FrameProcessor):
|
||||
"""Base LLM aggregator that uses an LLMContext for conversation storage.
|
||||
|
||||
@@ -205,8 +222,12 @@ class LLMContextAggregator(FrameProcessor):
|
||||
self._aggregation = []
|
||||
|
||||
@abstractmethod
|
||||
async def push_aggregation(self):
|
||||
"""Push the current aggregation downstream."""
|
||||
async def push_aggregation(self) -> str:
|
||||
"""Push the current aggregation downstream.
|
||||
|
||||
Returns:
|
||||
The pushed aggregation.
|
||||
"""
|
||||
pass
|
||||
|
||||
def aggregation_string(self) -> str:
|
||||
@@ -243,7 +264,7 @@ class LLMUserAggregator(LLMContextAggregator):
|
||||
...
|
||||
|
||||
@aggregator.event_handler("on_user_turn_stopped")
|
||||
async def on_user_turn_stopped(aggregator, strategy: BaseUserTurnStopStrategy):
|
||||
async def on_user_turn_stopped(aggregator, strategy: BaseUserTurnStopStrategy, message: UserTurnStoppedMessage):
|
||||
...
|
||||
|
||||
@aggregator.event_handler("on_user_turn_stop_timeout")
|
||||
@@ -348,16 +369,18 @@ class LLMUserAggregator(LLMContextAggregator):
|
||||
|
||||
await self._user_turn_controller.process_frame(frame)
|
||||
|
||||
async def push_aggregation(self):
|
||||
async def push_aggregation(self) -> str:
|
||||
"""Push the current aggregation."""
|
||||
if len(self._aggregation) == 0:
|
||||
return
|
||||
return ""
|
||||
|
||||
aggregation = self.aggregation_string()
|
||||
await self.reset()
|
||||
self._context.add_message({"role": self.role, "content": aggregation})
|
||||
await self.push_context_frame()
|
||||
|
||||
return aggregation
|
||||
|
||||
async def _start(self, frame: StartFrame):
|
||||
await self._user_turn_controller.setup(self.task_manager)
|
||||
|
||||
@@ -493,9 +516,10 @@ class LLMUserAggregator(LLMContextAggregator):
|
||||
await self.broadcast_frame(UserStoppedSpeakingFrame)
|
||||
|
||||
# Always push context frame.
|
||||
await self.push_aggregation()
|
||||
aggregation = await self.push_aggregation()
|
||||
|
||||
await self._call_event_handler("on_user_turn_stopped", strategy)
|
||||
message = UserTurnStoppedMessage(content=aggregation)
|
||||
await self._call_event_handler("on_user_turn_stopped", strategy, message)
|
||||
|
||||
async def _on_user_turn_stop_timeout(self, controller):
|
||||
await self._call_event_handler("on_user_turn_stop_timeout")
|
||||
@@ -633,10 +657,10 @@ class LLMAssistantAggregator(LLMContextAggregator):
|
||||
else:
|
||||
await self.push_frame(frame, direction)
|
||||
|
||||
async def push_aggregation(self):
|
||||
async def push_aggregation(self) -> str:
|
||||
"""Push the current assistant aggregation with timestamp."""
|
||||
if not self._aggregation:
|
||||
return
|
||||
return ""
|
||||
|
||||
aggregation = self.aggregation_string()
|
||||
await self.reset()
|
||||
@@ -651,6 +675,8 @@ class LLMAssistantAggregator(LLMContextAggregator):
|
||||
timestamp_frame = LLMContextAssistantTimestampFrame(timestamp=time_now_iso8601())
|
||||
await self.push_frame(timestamp_frame)
|
||||
|
||||
return aggregation
|
||||
|
||||
async def _handle_llm_run(self, frame: LLMRunFrame):
|
||||
await self.push_context_frame(FrameDirection.UPSTREAM)
|
||||
|
||||
|
||||
@@ -143,6 +143,7 @@ class TestLLMUserAggregator(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
should_start = None
|
||||
should_stop = None
|
||||
stop_message = None
|
||||
|
||||
@user_aggregator.event_handler("on_user_turn_started")
|
||||
async def on_user_turn_started(aggregator, strategy):
|
||||
@@ -150,9 +151,10 @@ class TestLLMUserAggregator(unittest.IsolatedAsyncioTestCase):
|
||||
should_start = True
|
||||
|
||||
@user_aggregator.event_handler("on_user_turn_stopped")
|
||||
async def on_user_turn_stopped(aggregator, strategy):
|
||||
nonlocal should_stop
|
||||
async def on_user_turn_stopped(aggregator, strategy, message):
|
||||
nonlocal should_stop, stop_message
|
||||
should_stop = True
|
||||
stop_message = message
|
||||
|
||||
pipeline = Pipeline([user_aggregator])
|
||||
|
||||
@@ -177,6 +179,7 @@ class TestLLMUserAggregator(unittest.IsolatedAsyncioTestCase):
|
||||
)
|
||||
self.assertTrue(should_start)
|
||||
self.assertTrue(should_stop)
|
||||
self.assertEqual(stop_message.content, "Hello!")
|
||||
|
||||
async def test_user_turn_stop_timeout_no_transcription(self):
|
||||
context = LLMContext()
|
||||
@@ -196,7 +199,7 @@ class TestLLMUserAggregator(unittest.IsolatedAsyncioTestCase):
|
||||
should_start = True
|
||||
|
||||
@user_aggregator.event_handler("on_user_turn_stopped")
|
||||
async def on_user_turn_stopped(aggregator, strategy):
|
||||
async def on_user_turn_stopped(aggregator, strategy, message):
|
||||
nonlocal should_stop
|
||||
should_stop = True
|
||||
|
||||
@@ -236,6 +239,7 @@ class TestLLMUserAggregator(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
should_start = None
|
||||
should_stop = None
|
||||
stop_message = None
|
||||
timeout = None
|
||||
|
||||
@user_aggregator.event_handler("on_user_turn_started")
|
||||
@@ -244,9 +248,10 @@ class TestLLMUserAggregator(unittest.IsolatedAsyncioTestCase):
|
||||
should_start = True
|
||||
|
||||
@user_aggregator.event_handler("on_user_turn_stopped")
|
||||
async def on_user_turn_stopped(aggregator, strategy):
|
||||
nonlocal should_stop
|
||||
async def on_user_turn_stopped(aggregator, strategy, message):
|
||||
nonlocal should_stop, stop_message
|
||||
should_stop = True
|
||||
stop_message = message
|
||||
|
||||
@user_aggregator.event_handler("on_user_turn_stop_timeout")
|
||||
async def on_user_turn_stop_timeout(aggregator):
|
||||
@@ -271,6 +276,7 @@ class TestLLMUserAggregator(unittest.IsolatedAsyncioTestCase):
|
||||
# The transcription strategy should kick-in before the user turn end timeout.
|
||||
self.assertTrue(should_start)
|
||||
self.assertTrue(should_stop)
|
||||
self.assertEqual(stop_message.content, "Hello!")
|
||||
self.assertFalse(timeout)
|
||||
|
||||
async def test_user_mute_strategies(self):
|
||||
|
||||
Reference in New Issue
Block a user