Refactoring to remove "Reset" and "TTSStoppedFrame" from word.

This commit is contained in:
filipi87
2026-03-25 17:47:24 -03:00
parent c4253a7d98
commit b4096f9a11
9 changed files with 35 additions and 69 deletions

View File

@@ -611,8 +611,6 @@ class AzureTTSService(TTSService, AzureBaseTTSService):
await super().push_frame(frame, direction)
if isinstance(frame, (TTSStoppedFrame, InterruptionFrame)):
self._reset_state()
if isinstance(frame, TTSStoppedFrame) and self._current_context_id:
await self.add_word_timestamps([("Reset", 0)], self._current_context_id)
def _reset_state(self):
"""Reset TTS state between turns."""

View File

@@ -606,7 +606,7 @@ class CartesiaTTSService(WebsocketTTSService):
ctx_id = msg["context_id"]
if msg["type"] == "done":
await self.stop_ttfb_metrics()
await self.add_word_timestamps([("TTSStoppedFrame", 0), ("Reset", 0)], ctx_id)
await self.append_to_audio_context(ctx_id, TTSStoppedFrame(context_id=ctx_id))
await self.remove_audio_context(ctx_id)
elif msg["type"] == "timestamps":
# Process the timestamps based on language before adding them

View File

@@ -620,18 +620,6 @@ class ElevenLabsTTSService(WebsocketTTSService):
msg = {"context_id": flush_id, "flush": True}
await self._websocket.send(json.dumps(msg))
async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM):
"""Push a frame and handle state changes.
Args:
frame: The frame to push.
direction: The direction to push the frame.
"""
await super().push_frame(frame, direction)
if isinstance(frame, (TTSStoppedFrame, InterruptionFrame)):
if isinstance(frame, TTSStoppedFrame):
await self.add_word_timestamps([("Reset", 0)], self.get_active_audio_context_id())
async def _connect(self):
await super()._connect()
@@ -1130,10 +1118,6 @@ class ElevenLabsHttpTTSService(TTSService):
if isinstance(frame, (InterruptionFrame, TTSStoppedFrame)):
# Reset timing on interruption or stop
self._reset_state()
if isinstance(frame, TTSStoppedFrame):
await self.add_word_timestamps([("Reset", 0)])
elif isinstance(frame, LLMFullResponseEndFrame):
# End of turn - reset previous text
self._previous_text = ""

View File

@@ -342,7 +342,7 @@ class GradiumTTSService(WebsocketTTSService):
elif msg["type"] == "end_of_stream":
if ctx_id and self.audio_context_available(ctx_id):
await self.add_word_timestamps([("TTSStoppedFrame", 0), ("Reset", 0)], ctx_id)
await self.append_to_audio_context(ctx_id, TTSStoppedFrame(context_id=ctx_id))
await self.remove_audio_context(ctx_id)
if ctx_id:
self._setup_context_ids.discard(ctx_id)

View File

@@ -235,9 +235,6 @@ class HumeTTSService(TTSService):
# Reset timing on interruption or stop
self._reset_state()
if isinstance(frame, TTSStoppedFrame):
await self.add_word_timestamps([("Reset", 0)])
async def update_setting(self, key: str, value: Any) -> None:
"""Runtime updates via key/value pair.

View File

@@ -248,8 +248,6 @@ class InworldHttpTTSService(TTSService):
await super().push_frame(frame, direction)
if isinstance(frame, (InterruptionFrame, TTSStoppedFrame)):
self._cumulative_time = 0.0
if isinstance(frame, TTSStoppedFrame):
await self.add_word_timestamps([("Reset", 0)])
def _calculate_word_times(
self,
@@ -728,8 +726,6 @@ class InworldTTSService(WebsocketTTSService):
)
self._cumulative_time = 0.0
self._generation_end_time = 0.0
if isinstance(frame, TTSStoppedFrame):
await self.add_word_timestamps([("Reset", 0)])
async def on_turn_context_created(self, context_id: str):
"""Eagerly open the context on the server when a new turn starts.
@@ -996,7 +992,7 @@ class InworldTTSService(WebsocketTTSService):
if "contextClosed" in result:
logger.trace(f"{self}: Context closed on server: {ctx_id}")
await self.stop_ttfb_metrics()
await self.add_word_timestamps([("TTSStoppedFrame", 0), ("Reset", 0)], ctx_id)
await self.append_to_audio_context(ctx_id, TTSStoppedFrame(context_id=ctx_id))
await self.remove_audio_context(ctx_id)
async def _keepalive_task_handler(self):

View File

@@ -387,7 +387,9 @@ class ResembleAITTSService(WebsocketTTSService):
if request_id in self._request_id_to_context:
del self._request_id_to_context[request_id]
await self.add_word_timestamps([("TTSStoppedFrame", 0), ("Reset", 0)], context_id)
await self.append_to_audio_context(
context_id, TTSStoppedFrame(context_id=context_id)
)
await self.remove_audio_context(context_id)
elif msg_type == "error":

View File

@@ -604,18 +604,6 @@ class RimeTTSService(WebsocketTTSService):
await self.push_error(error_msg=f"Error: {msg['message']}")
self.reset_active_audio_context()
async def push_frame(self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM):
"""Push frame and handle end-of-turn conditions.
Args:
frame: The frame to push.
direction: The direction to push the frame.
"""
await super().push_frame(frame, direction)
if isinstance(frame, (TTSStoppedFrame, InterruptionFrame)):
if isinstance(frame, TTSStoppedFrame):
await self.add_word_timestamps([("Reset", 0)])
@traced_tts
async def run_tts(self, text: str, context_id: str) -> AsyncGenerator[Frame, None]:
"""Generate speech from text using Rime's streaming API.

View File

@@ -342,7 +342,7 @@ class TTSService(AIService):
self._initial_word_timestamp: int = -1
self._initial_word_times: List[Tuple[str, float, Optional[str]]] = []
# PTS of the last word frame pushed via _add_word_timestamps, used to assign
# correct PTS to sentinel frames ("TTSStoppedFrame", "Reset") that follow.
# correct PTS to TTSStoppedFrame and LLMFullResponseEndFrame.
self._word_last_pts: int = 0
self._llm_response_started: bool = False
self._reuse_context_id_within_turn: bool = reuse_context_id_within_turn
@@ -1166,33 +1166,20 @@ class TTSService(AIService):
called (i.e. when the first audio chunk is received).
"""
for word, timestamp in word_times:
if word == "Reset" and timestamp == 0:
await self.reset_word_timestamps()
if self._llm_response_started:
self._llm_response_started = False
frame = LLMFullResponseEndFrame()
frame.pts = self._word_last_pts
await self.push_frame(frame)
elif word == "TTSStoppedFrame" and timestamp == 0:
frame = TTSStoppedFrame(context_id=context_id)
frame.pts = self._word_last_pts
frame.context_id = context_id
await self.push_frame(frame)
ts_ns = seconds_to_nanoseconds(timestamp)
if self._initial_word_timestamp == -1:
# Cache until we have audio and can compute PTS.
self._initial_word_times.append((word, timestamp, context_id))
else:
ts_ns = seconds_to_nanoseconds(timestamp)
if self._initial_word_timestamp == -1:
# Cache until we have audio and can compute PTS.
self._initial_word_times.append((word, timestamp, context_id))
else:
# Assumption: word-by-word text frames don't include spaces, so
# we can rely on the default includes_inter_frame_spaces=False
frame = TTSTextFrame(word, aggregated_by=AggregationType.WORD)
frame.pts = self._initial_word_timestamp + ts_ns
frame.context_id = context_id
if context_id in self._tts_contexts:
frame.append_to_context = self._tts_contexts[context_id].append_to_context
self._word_last_pts = frame.pts
await self.push_frame(frame)
# Assumption: word-by-word text frames don't include spaces, so
# we can rely on the default includes_inter_frame_spaces=False
frame = TTSTextFrame(word, aggregated_by=AggregationType.WORD)
frame.pts = self._initial_word_timestamp + ts_ns
frame.context_id = context_id
if context_id in self._tts_contexts:
frame.append_to_context = self._tts_contexts[context_id].append_to_context
self._word_last_pts = frame.pts
await self.push_frame(frame)
#
# Audio context methods (active when using websocket-based TTS with context management)
@@ -1223,7 +1210,8 @@ class TTSService(AIService):
if self.audio_context_available(context_id):
logger.trace(f"{self} appending audio {frame} to audio context {context_id}")
await self._audio_contexts[context_id].put(frame)
elif context_id == self._turn_context_id:
# In case the frame is None, we should not recreate the context.
elif context_id == self._turn_context_id and frame:
# Sometimes the HTTP service can take more than 3 seconds without sending any audio
# So we are now recreating the context id while we are in the same turn
logger.debug(f"{self} recreating audio context {context_id}")
@@ -1348,6 +1336,15 @@ class TTSService(AIService):
self._serialization_queue.task_done()
async def _maybe_reset_word_timestamps(self):
await self.reset_word_timestamps()
# If self._push_text_frames is True, we have already pushed the original LLMFullResponseEndFrame
if self._llm_response_started and not self._push_text_frames:
self._llm_response_started = False
frame = LLMFullResponseEndFrame()
frame.pts = self._word_last_pts
await self.push_frame(frame)
async def _handle_audio_context(self, context_id: str):
"""Process items from an audio context queue until it is exhausted."""
queue = self._audio_contexts[context_id]
@@ -1382,6 +1379,9 @@ class TTSService(AIService):
should_push_stop_frame = self._push_stop_frames
elif isinstance(frame, TTSStoppedFrame):
should_push_stop_frame = False
# Setting the last word timestamp as the TTSStoppedFrame PTS
if not frame.pts:
frame.pts = self._word_last_pts
if isinstance(frame, ErrorFrame):
await self.push_error_frame(frame)
@@ -1392,6 +1392,7 @@ class TTSService(AIService):
logger.trace(f"{self} time out on audio context {context_id}")
if should_push_stop_frame and self._push_stop_frames:
await self.push_frame(TTSStoppedFrame(context_id=context_id))
await self._maybe_reset_word_timestamps()
break
if should_push_stop_frame and self._push_stop_frames: