|
4 | 4 | import json |
5 | 5 | import time |
6 | 6 | import weakref |
| 7 | +from collections import deque |
7 | 8 | from dataclasses import dataclass, replace |
8 | 9 | from typing import TYPE_CHECKING, Any, Literal |
9 | 10 |
|
@@ -470,6 +471,24 @@ class _PendingCommit: |
470 | 471 | acknowledged_at: float | None = None |
471 | 472 |
|
472 | 473 |
|
| 474 | +@dataclass(eq=False) |
| 475 | +class _Boundary: |
| 476 | + """What the server reported about where an utterance ended. |
| 477 | +
|
| 478 | + A stop carries the engine's end position when the gateway exposes one. |
| 479 | + Without it, `sent_at_ms` still records how much audio had been uploaded |
| 480 | + when the report arrived: anything sent afterwards is beyond what the |
| 481 | + server had classified, so it may belong to a later utterance. |
| 482 | + """ |
| 483 | + |
| 484 | + position_ms: float | None |
| 485 | + #: Audio uploaded when this boundary was recorded. |
| 486 | + sent_at_ms: float |
| 487 | + #: The manual commit that established this boundary, when any. Its item |
| 488 | + #: window bounds how long a final is still expected for it. |
| 489 | + commit: _PendingCommit | None = None |
| 490 | + |
| 491 | + |
473 | 492 | class RecognizeStream(stt.RecognizeStream): |
474 | 493 | def __init__( |
475 | 494 | self, |
@@ -509,8 +528,15 @@ def __init__( |
509 | 528 | # sent in this attempt; both restart from zero per connection. |
510 | 529 | self._replay_base_ms = 0.0 |
511 | 530 | self._sent_ms = 0.0 |
512 | | - # Last audio position the server reported as finalized, if any. |
513 | | - self._settled_ms: float | None = None |
| 531 | + # Finalization reports awaiting the final that owns them, oldest |
| 532 | + # first. The server finalizes utterances in order, so the oldest |
| 533 | + # report belongs to the next completed transcript. |
| 534 | + self._pending_boundaries: deque[_Boundary] = deque() |
| 535 | + # Audio uploaded when the last finalization report arrived, kept |
| 536 | + # after its boundary is consumed or pruned. A final that arrives |
| 537 | + # once its own report is gone is still bounded by it: audio uploaded |
| 538 | + # afterwards may belong to a later utterance. |
| 539 | + self._last_report_ms: float | None = None |
514 | 540 | self._replay_ambiguous = False |
515 | 541 |
|
516 | 542 | self._audio_duration = 0.0 |
@@ -547,7 +573,8 @@ async def _run_impl(self) -> None: |
547 | 573 | self._finals_received = 0 |
548 | 574 | self._replay_base_ms = 0.0 |
549 | 575 | self._sent_ms = 0.0 |
550 | | - self._settled_ms = None |
| 576 | + self._pending_boundaries.clear() |
| 577 | + self._last_report_ms = None |
551 | 578 | # Per attempt: a reconnect with nothing to resend has no upload to wait |
552 | 579 | # for. Replayed or new frames set it again. |
553 | 580 | self._consumed_audio = False |
@@ -727,19 +754,100 @@ def _remember(self, data: rtc.AudioFrame | stt.RecognizeStream._FlushSentinel) - |
727 | 754 | self._replay.append((data, seconds * 1000)) |
728 | 755 |
|
729 | 756 | def _settle(self, boundary_ms: float) -> None: |
730 | | - """Drop retained audio the server finalized up to `boundary_ms`.""" |
| 757 | + """Drop retained audio the server finalized up to `boundary_ms`. |
| 758 | +
|
| 759 | + A boundary is measured in milliseconds and can fall inside a retained |
| 760 | + frame, so that frame keeps only its unfinalized suffix; resending the |
| 761 | + whole frame would repeat speech the server already transcribed. |
| 762 | + """ |
731 | 763 | kept: list[tuple[Any, float]] = [] |
732 | | - position = self._replay_base_ms |
| 764 | + base: float | None = None |
| 765 | + cursor = self._replay_base_ms |
| 766 | + trimmed = False |
733 | 767 | for item, duration_ms in self._replay: |
734 | | - if position + duration_ms <= boundary_ms + 0.001: |
735 | | - position += duration_ms |
736 | | - self._replay_seconds = max(0.0, self._replay_seconds - duration_ms / 1000) |
| 768 | + end = cursor + duration_ms |
| 769 | + if end <= boundary_ms + 0.001: |
| 770 | + cursor = end |
737 | 771 | continue |
| 772 | + if not trimmed: |
| 773 | + trimmed = True |
| 774 | + if isinstance(item, rtc.AudioFrame) and cursor < boundary_ms: |
| 775 | + rate = item.sample_rate |
| 776 | + cut = min( |
| 777 | + int(round((boundary_ms - cursor) * rate / 1000)), |
| 778 | + item.samples_per_channel, |
| 779 | + ) |
| 780 | + if cut > 0: |
| 781 | + stride = item.num_channels * 2 |
| 782 | + item = rtc.AudioFrame( |
| 783 | + data=item.data.tobytes()[cut * stride :], |
| 784 | + sample_rate=rate, |
| 785 | + num_channels=item.num_channels, |
| 786 | + samples_per_channel=item.samples_per_channel - cut, |
| 787 | + ) |
| 788 | + duration_ms -= cut * 1000 / rate |
| 789 | + cursor = boundary_ms |
| 790 | + if base is None: |
| 791 | + base = cursor |
738 | 792 | kept.append((item, duration_ms)) |
739 | | - self._replay_base_ms = position |
| 793 | + cursor += duration_ms |
740 | 794 | self._replay = kept |
| 795 | + self._replay_base_ms = base if base is not None else cursor |
| 796 | + self._replay_seconds = sum(duration for _, duration in kept) / 1000 |
741 | 797 | self._audio_since_final = self._replay_seconds |
742 | 798 |
|
| 799 | + def _prune_boundaries(self, now: float) -> None: |
| 800 | + """Drop finalization reports whose final is no longer expected. |
| 801 | +
|
| 802 | + A manual commit whose item window elapsed produced no final, so it |
| 803 | + committed silence: its boundary must not be consumed by the final of |
| 804 | + a later utterance. Speech-stopped reports carry no such bound, since |
| 805 | + the gateway pairs every stop with a final. |
| 806 | + """ |
| 807 | + while self._pending_boundaries: |
| 808 | + commit = self._pending_boundaries[0].commit |
| 809 | + if ( |
| 810 | + commit is None |
| 811 | + or commit.acknowledged_at is None |
| 812 | + or now < commit.acknowledged_at + _COMMIT_ITEM_WINDOW |
| 813 | + ): |
| 814 | + break |
| 815 | + self._pending_boundaries.popleft() |
| 816 | + |
| 817 | + def _settle_replay(self) -> None: |
| 818 | + """Release the replay audio a completed transcript finalized. |
| 819 | +
|
| 820 | + Each report is consumed by the final that owns it, so a boundary |
| 821 | + always describes the utterance that reported it. A report without a |
| 822 | + position still bounds what the final can cover by how much audio had |
| 823 | + been uploaded when it arrived: if more was sent afterwards, that |
| 824 | + audio may belong to a later utterance, so it stays retained and a |
| 825 | + retry is refused rather than allowed to discard it. A final whose own |
| 826 | + report is already gone is bounded the same way by the last one. |
| 827 | + """ |
| 828 | + self._prune_boundaries(time.monotonic()) |
| 829 | + boundary = self._pending_boundaries.popleft() if self._pending_boundaries else None |
| 830 | + |
| 831 | + if boundary is not None and boundary.position_ms is not None: |
| 832 | + self._replay_ambiguous = False |
| 833 | + self._replay_overflow = False |
| 834 | + self._settle(boundary.position_ms) |
| 835 | + return |
| 836 | + |
| 837 | + reported_ms = boundary.sent_at_ms if boundary is not None else self._last_report_ms |
| 838 | + if reported_ms is not None and self._sent_ms > reported_ms: |
| 839 | + # Audio arrived after the server finalized up to here, so the |
| 840 | + # final cannot be shown to cover it. |
| 841 | + self._replay_ambiguous = True |
| 842 | + self._audio_since_final = self._replay_seconds |
| 843 | + return |
| 844 | + |
| 845 | + # Nothing indicates audio beyond this final, so it covers everything |
| 846 | + # uploaded so far and there is nothing left to resend. |
| 847 | + self._replay_ambiguous = False |
| 848 | + self._replay_overflow = False |
| 849 | + self._settle(self._sent_ms) |
| 850 | + |
743 | 851 | async def _send_item( |
744 | 852 | self, |
745 | 853 | ws: _WebSocket, |
@@ -935,19 +1043,25 @@ async def _maybe_commit(self, ws: _WebSocket) -> None: |
935 | 1043 | if pending.acknowledged_at is None |
936 | 1044 | or now < pending.acknowledged_at + _COMMIT_ITEM_WINDOW |
937 | 1045 | ] |
| 1046 | + self._prune_boundaries(now) |
938 | 1047 | commit = _PendingCommit() |
939 | 1048 | # Register before yielding: the receiver may process the ack while |
940 | 1049 | # send_str is suspended. Roll back only if still unacked. |
941 | 1050 | self._commits.append(commit) |
| 1051 | + # A commit finalizes everything sent so far, so its boundary is the |
| 1052 | + # position at send time. Recorded before the await for the same |
| 1053 | + # reason as the commit itself: the final can be processed while the |
| 1054 | + # send is still draining. A failed upload takes the whole attempt |
| 1055 | + # down with it, so a boundary left behind is never used. |
| 1056 | + self._last_report_ms = self._sent_ms |
| 1057 | + self._pending_boundaries.append(_Boundary(self._sent_ms, self._sent_ms, commit)) |
942 | 1058 | try: |
943 | 1059 | await self._send_commit(ws) |
944 | 1060 | except BaseException: |
945 | 1061 | if commit.acknowledged_at is None: |
946 | 1062 | self._commits.remove(commit) |
947 | 1063 | raise |
948 | 1064 | self._uncommitted_audio = False |
949 | | - # A commit finalizes everything sent so far. |
950 | | - self._settled_ms = self._sent_ms |
951 | 1065 |
|
952 | 1066 | # -- downstream -------------------------------------------------------- |
953 | 1067 |
|
@@ -1024,8 +1138,13 @@ def _handle_event(self, event: dict[str, Any]) -> None: |
1024 | 1138 | self._speech_active = False |
1025 | 1139 | # The gateway exposes the engine's end position as audio_start_ms. |
1026 | 1140 | position = event.get("audio_end_ms", event.get("audio_start_ms")) |
1027 | | - if isinstance(position, (int, float)): |
1028 | | - self._settled_ms = max(self._settled_ms or 0.0, float(position)) |
| 1141 | + self._last_report_ms = self._sent_ms |
| 1142 | + self._pending_boundaries.append( |
| 1143 | + _Boundary( |
| 1144 | + float(position) if isinstance(position, (int, float)) else None, |
| 1145 | + self._sent_ms, |
| 1146 | + ) |
| 1147 | + ) |
1029 | 1148 | self._event_ch.send_nowait( |
1030 | 1149 | stt.SpeechEvent( |
1031 | 1150 | type=stt.SpeechEventType.END_OF_SPEECH, |
@@ -1103,15 +1222,7 @@ def _handle_completed(self, event: dict[str, Any]) -> None: |
1103 | 1222 | self._open_items.discard(item_id) |
1104 | 1223 | self._expected_finals = max(0, self._expected_finals - 1) |
1105 | 1224 | self._finals_received += 1 |
1106 | | - if self._settled_ms is None: |
1107 | | - # No boundary for this final: later audio may be buffered, so a |
1108 | | - # retry could discard it. Keep everything and refuse the retry. |
1109 | | - self._replay_ambiguous = True |
1110 | | - self._audio_since_final = self._replay_seconds |
1111 | | - else: |
1112 | | - self._replay_ambiguous = False |
1113 | | - self._replay_overflow = False |
1114 | | - self._settle(self._settled_ms) |
| 1225 | + self._settle_replay() |
1115 | 1226 | # A final closes any unidentified report too: the server finalizes |
1116 | 1227 | # utterances in order, and an id-less start belongs to this one. |
1117 | 1228 | self._anonymous_open = False |
|
0 commit comments