213 lines
8.1 KiB
Python
213 lines
8.1 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections import Counter, OrderedDict
|
|
from time import perf_counter
|
|
from typing import Any
|
|
|
|
from speech_to_speech.baseHandler import BaseHandler
|
|
from speech_to_speech.pipeline.handler_types import STTIn, STTOut
|
|
from speech_to_speech.pipeline.messages import PartialTranscription, Transcription, VADAudio
|
|
from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BaseSTTHandler(BaseHandler[STTIn, STTOut]):
|
|
"""Base STT handler with speculative-turn stale input filtering."""
|
|
|
|
_MAX_COMPLETED_FINAL_REVISIONS = 2048
|
|
|
|
speculative_turns: SpeculativeTurnTracker | None = None
|
|
final_revision_settle_s: float = 0.0
|
|
|
|
def should_process_input(self, item: STTIn) -> bool:
|
|
mode = getattr(item, "mode", None)
|
|
turn_id = getattr(item, "turn_id", None)
|
|
turn_revision = getattr(item, "turn_revision", None)
|
|
if self._is_completed_final_revision(item):
|
|
queued_drops = self._drop_stale_queued_inputs()
|
|
self._log_stale_turn_item(item, "input-after-final", queued_drops=queued_drops)
|
|
return False
|
|
if mode == "progressive" and self._has_queued_final_for_revision(item):
|
|
self._log_stale_turn_item(item, "progressive-before-final")
|
|
return False
|
|
|
|
wait_for_stability = mode == "final"
|
|
gate_start = perf_counter()
|
|
is_latest = self._is_latest_turn_item(
|
|
item,
|
|
wait_for_pending_reopen=True,
|
|
wait_for_stability=wait_for_stability,
|
|
)
|
|
gate_wait_s = perf_counter() - gate_start
|
|
if gate_wait_s >= 0.05:
|
|
logger.info(
|
|
"%s: STT input gate waited %.3fs for turn=%s rev=%s mode=%s latest=%s age=%.3fs queue=%s",
|
|
self.__class__.__name__,
|
|
gate_wait_s,
|
|
turn_id,
|
|
turn_revision,
|
|
mode,
|
|
is_latest,
|
|
self._item_age_s(item),
|
|
self._safe_qsize(),
|
|
)
|
|
|
|
if not is_latest:
|
|
queued_drops = self._drop_stale_queued_inputs()
|
|
self._log_stale_turn_item(item, "input", queued_drops=queued_drops)
|
|
return False
|
|
return True
|
|
|
|
def should_emit_output(self, output: STTOut) -> bool:
|
|
if isinstance(output, PartialTranscription) and self._is_completed_final_revision(output):
|
|
self._log_stale_turn_item(output, "output-after-final")
|
|
return False
|
|
|
|
if not self._is_latest_turn_item(output, wait_for_pending_reopen=True, wait_for_stability=False):
|
|
self._log_stale_turn_item(output, "output")
|
|
return False
|
|
return True
|
|
|
|
def before_emit_output(self, output: STTOut) -> None:
|
|
if isinstance(output, Transcription):
|
|
self._mark_completed_final_revision(output)
|
|
|
|
def _is_latest_turn_item(
|
|
self,
|
|
item: object,
|
|
*,
|
|
wait_for_pending_reopen: bool,
|
|
wait_for_stability: bool,
|
|
) -> bool:
|
|
if self.speculative_turns is None:
|
|
return True
|
|
turn_id = getattr(item, "turn_id", None)
|
|
turn_revision = getattr(item, "turn_revision", None)
|
|
if turn_id is None or turn_revision is None:
|
|
return True
|
|
|
|
if wait_for_stability:
|
|
is_latest = self.speculative_turns.is_latest_after_stability_window(
|
|
turn_id,
|
|
turn_revision,
|
|
self.final_revision_settle_s,
|
|
)
|
|
elif wait_for_pending_reopen:
|
|
is_latest = self.speculative_turns.is_latest_after_pending_reopen(turn_id, turn_revision)
|
|
else:
|
|
is_latest = self.speculative_turns.is_latest(turn_id, turn_revision)
|
|
return is_latest
|
|
|
|
def _drop_stale_queued_inputs(self) -> int:
|
|
if self.speculative_turns is None or not hasattr(self.queue_in, "mutex") or not hasattr(self.queue_in, "queue"):
|
|
return 0
|
|
|
|
dropped = 0
|
|
with self.queue_in.mutex:
|
|
kept: list[Any] = []
|
|
while self.queue_in.queue:
|
|
queued_item = self.queue_in.queue.popleft()
|
|
if isinstance(queued_item, VADAudio) and (
|
|
self._is_completed_final_revision(queued_item)
|
|
or (queued_item.mode == "progressive" and self._has_queued_final_for_revision_locked(queued_item))
|
|
or not self._is_latest_turn_item(
|
|
queued_item,
|
|
wait_for_pending_reopen=False,
|
|
wait_for_stability=False,
|
|
)
|
|
):
|
|
dropped += 1
|
|
else:
|
|
kept.append(queued_item)
|
|
self.queue_in.queue.extend(kept)
|
|
if dropped:
|
|
self.queue_in.not_full.notify_all()
|
|
return dropped
|
|
|
|
def _log_stale_turn_item(self, item: object, stage: str, *, queued_drops: int = 0) -> None:
|
|
turn_id = getattr(item, "turn_id", None)
|
|
turn_revision = getattr(item, "turn_revision", None)
|
|
if turn_id is None or turn_revision is None:
|
|
return
|
|
|
|
if not hasattr(self, "_stale_drop_counts"):
|
|
self._stale_drop_counts: Counter[tuple[str, str, int]] = Counter()
|
|
key = (stage, turn_id, turn_revision)
|
|
self._stale_drop_counts[key] += 1
|
|
|
|
message = "%s: dropping stale STT %s for turn=%s rev=%s age=%.3fs"
|
|
args: tuple[object, ...] = (
|
|
self.__class__.__name__,
|
|
stage,
|
|
turn_id,
|
|
turn_revision,
|
|
self._item_age_s(item),
|
|
)
|
|
if queued_drops:
|
|
message += " (+%d queued)"
|
|
args = (*args, queued_drops)
|
|
|
|
if self._stale_drop_counts[key] == 1:
|
|
logger.info(message, *args)
|
|
else:
|
|
logger.debug(message, *args)
|
|
|
|
def _item_age_s(self, item: object) -> float:
|
|
created_at_s = getattr(item, "created_at_s", None)
|
|
if not isinstance(created_at_s, float):
|
|
return 0.0
|
|
return max(0.0, perf_counter() - created_at_s)
|
|
|
|
def _safe_qsize(self) -> int | str:
|
|
try:
|
|
return self.queue_in.qsize()
|
|
except NotImplementedError:
|
|
return "unknown"
|
|
|
|
def _has_queued_final_for_revision(self, item: object) -> bool:
|
|
if not hasattr(self.queue_in, "mutex") or not hasattr(self.queue_in, "queue"):
|
|
return False
|
|
with self.queue_in.mutex:
|
|
return self._has_queued_final_for_revision_locked(item)
|
|
|
|
def _has_queued_final_for_revision_locked(self, item: object) -> bool:
|
|
key = self._revision_key(item)
|
|
if key is None:
|
|
return False
|
|
return any(
|
|
isinstance(queued_item, VADAudio) and queued_item.mode == "final" and self._revision_key(queued_item) == key
|
|
for queued_item in self.queue_in.queue
|
|
)
|
|
|
|
def _revision_key(self, item: object) -> tuple[str, int] | None:
|
|
turn_id = getattr(item, "turn_id", None)
|
|
turn_revision = getattr(item, "turn_revision", None)
|
|
if not isinstance(turn_id, str) or not isinstance(turn_revision, int):
|
|
return None
|
|
return (turn_id, turn_revision)
|
|
|
|
def _completed_final_revisions(self) -> OrderedDict[tuple[str, int], None]:
|
|
if not hasattr(self, "_completed_final_revision_keys"):
|
|
self._completed_final_revision_keys: OrderedDict[tuple[str, int], None] = OrderedDict()
|
|
return self._completed_final_revision_keys
|
|
|
|
def _is_completed_final_revision(self, item: object) -> bool:
|
|
key = self._revision_key(item)
|
|
return key is not None and key in self._completed_final_revisions()
|
|
|
|
def _mark_completed_final_revision(self, output: Transcription) -> None:
|
|
key = self._revision_key(output)
|
|
if key is None:
|
|
return
|
|
completed = self._completed_final_revisions()
|
|
completed[key] = None
|
|
completed.move_to_end(key)
|
|
while len(completed) > self._MAX_COMPLETED_FINAL_REVISIONS:
|
|
completed.popitem(last=False)
|
|
|
|
def on_session_end(self) -> None:
|
|
if hasattr(self, "_completed_final_revision_keys"):
|
|
self._completed_final_revision_keys.clear()
|