-
Notifications
You must be signed in to change notification settings - Fork 3.5k
conversation-aware STT recognition (keyterms + chat context) #6039
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
59f1a58
ebbe1b5
8234757
73eb033
a6096af
12672e0
65cd307
2ca544a
ffd105c
e7b7c92
cb0b308
7780b6f
603417b
ef23cc3
d734805
d140b9e
abf645f
7cf1d68
86f4f81
2e82b55
0a93cb8
8ba5370
04751d1
b442136
a4a0012
5328bbd
a7acb41
1f6f803
c744533
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -196,6 +196,46 @@ def _diarization_enabled(extra_kwargs: dict[str, Any] | None) -> bool: | |
| return False | ||
|
|
||
|
|
||
| def _keyterms_extra_for_model( | ||
| model: NotGivenOr[str], | ||
| *, | ||
| extra_kwargs: dict[str, Any] | None = None, | ||
| session_keyterms: list[str] | None = None, | ||
| ) -> dict[str, Any] | None: | ||
| """Return the provider's keyterm ``extra`` entry: user keyterms (from ``extra_kwargs``) | ||
| merged with the framework ``session_keyterms``. | ||
|
|
||
| None if the model has no keyterm prompting, so ``_keyterms_extra_for_model(model) is not | ||
| None`` is also the capability check. | ||
| """ | ||
| if not (is_given(model) and isinstance(model, str)): | ||
| return None | ||
|
|
||
| extra_kwargs = extra_kwargs or {} | ||
| session_keyterms = session_keyterms or [] | ||
|
|
||
| if model.startswith("speechmatics/"): | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. would it make sense to add STTCapabilities.keyterms to the speechmatics plugin? |
||
| # keep existing entries as-is (they may carry sounds_like etc.); append new session terms | ||
| existing = list(extra_kwargs.get("additional_vocab", [])) | ||
| seen = {v["content"] for v in existing} | ||
| additions = set(session_keyterms) - seen | ||
| return {"additional_vocab": existing + [{"content": term} for term in additions]} | ||
|
|
||
| key: str | None = None | ||
| if model.startswith("deepgram/"): | ||
| key = "keyterm" | ||
| elif model.startswith("assemblyai/"): | ||
| key = "keyterms_prompt" | ||
|
|
||
| if key is None: | ||
| return None | ||
| # deepgram's keyterm may be a bare string; wrap it so it isn't splat char-by-char | ||
| existing = extra_kwargs.get(key, []) | ||
| if isinstance(existing, str): | ||
| existing = [existing] | ||
| return {key: list(dict.fromkeys([*existing, *session_keyterms]))} | ||
|
|
||
|
|
||
| STTLanguages = Literal["multi", "en", "de", "es", "fr", "ja", "pt", "zh", "hi"] | ||
|
|
||
|
|
||
|
|
@@ -516,6 +556,7 @@ def __init__( | |
| diarization=diarization_enabled, | ||
| aligned_transcript="word", | ||
| offline_recognize=False, | ||
| keyterms=_keyterms_extra_for_model(model) is not None, | ||
| ), | ||
| ) | ||
|
|
||
|
|
@@ -559,6 +600,7 @@ def __init__( | |
|
|
||
| self._session = http_session | ||
| self._vad = vad | ||
| self._session_keyterms: list[str] = [] # framework-managed; merged into extra_kwargs | ||
| self._streams = weakref.WeakSet[SpeechStream]() | ||
|
|
||
| @classmethod | ||
|
|
@@ -633,6 +675,10 @@ def update_options( | |
|
|
||
| self._opts.model = model | ||
| self._vad = _resolve_vad_for_model(model, self._vad) | ||
| self._capabilities = replace( | ||
| self._capabilities, | ||
| keyterms=_keyterms_extra_for_model(self._opts.model) is not None, | ||
| ) | ||
| if is_given(language): | ||
| self._opts.language = LanguageCode(language) | ||
| if is_given(extra): | ||
|
|
@@ -641,10 +687,37 @@ def update_options( | |
| self._capabilities, | ||
| diarization=_diarization_enabled(self._opts.extra_kwargs), | ||
| ) | ||
| # re-merge the active session keyterms so a user extra update doesn't drop them | ||
| keyterm_extra = _keyterms_extra_for_model( | ||
| self._opts.model, | ||
| extra_kwargs=self._opts.extra_kwargs, | ||
| session_keyterms=self._session_keyterms, | ||
| ) | ||
| if keyterm_extra is not None: | ||
| extra = {**extra, **keyterm_extra} | ||
|
|
||
| for stream in self._streams: | ||
| stream.update_options(model=model, language=language, extra=extra) | ||
|
|
||
| def _update_session_keyterms(self, keyterms: list[str]) -> None: | ||
| if keyterms == self._session_keyterms: | ||
| return | ||
| keyterm_extra = _keyterms_extra_for_model( | ||
| self._opts.model, extra_kwargs=self._opts.extra_kwargs, session_keyterms=keyterms | ||
| ) | ||
| if keyterm_extra is None: | ||
| super()._update_session_keyterms(keyterms) # warn-and-skip for unsupported models | ||
| return | ||
|
|
||
| self._session_keyterms = list(keyterms) | ||
| # inference applies extra live via session.update; defer to END_OF_SPEECH since the | ||
| # gateway may reconnect upstream when the keyterms change | ||
| for stream in self._streams: | ||
| if stream._speaking: | ||
| stream._pending_extra = keyterm_extra | ||
| else: | ||
| stream.update_options(extra=keyterm_extra) | ||
|
longcw marked this conversation as resolved.
|
||
|
|
||
| def _sanitize_options( | ||
| self, *, language: NotGivenOr[STTLanguages | str] = NOT_GIVEN | ||
| ) -> STTOptions: | ||
|
|
@@ -673,6 +746,9 @@ def __init__( | |
| self._request_id = str(utils.shortuuid("stt_request_")) | ||
|
|
||
| self._speaking = False | ||
| # keyterm extra set while the user is speaking; applied at END_OF_SPEECH (latest wins). | ||
| # inference applies live, but the gateway may reconnect upstream, so defer to a calm moment. | ||
| self._pending_extra: dict[str, Any] | None = None | ||
| self._speech_duration: float = 0 | ||
| self._ws: aiohttp.ClientWebSocketResponse | None = None | ||
| self._vad: vad.VAD | None = vad_instance | ||
|
|
@@ -696,6 +772,7 @@ def update_options( | |
| self._opts.language = LanguageCode(language) | ||
| if is_given(extra): | ||
| self._opts.extra_kwargs.update(extra) | ||
| self._pending_extra = None | ||
|
|
||
| has_update = is_given(model) or is_given(language) or is_given(extra) | ||
| if has_update and self._ws is not None and not self._ws.closed: | ||
|
|
@@ -712,6 +789,11 @@ def update_options( | |
| } | ||
| asyncio.ensure_future(self._send_session_update(update_msg)) | ||
|
|
||
| def _on_end_of_speech(self) -> None: | ||
| if self._pending_extra is not None: | ||
| self.update_options(extra=self._pending_extra) | ||
| self._pending_extra = None | ||
|
|
||
| async def _send_session_update(self, msg: dict[str, Any]) -> None: | ||
| try: | ||
| if self._ws is not None and not self._ws.closed: | ||
|
|
@@ -846,7 +928,18 @@ async def _connect_ws( | |
| "settings": { | ||
| "sample_rate": str(self._opts.sample_rate), | ||
| "encoding": self._opts.encoding, | ||
| "extra": self._opts.extra_kwargs, | ||
| # merge the framework session keyterms into the user's extra_kwargs keyterm key | ||
| "extra": { | ||
| **self._opts.extra_kwargs, | ||
| **( | ||
| _keyterms_extra_for_model( | ||
| self._opts.model, | ||
| extra_kwargs=self._opts.extra_kwargs, | ||
| session_keyterms=self._stt._session_keyterms, | ||
| ) | ||
| or {} | ||
| ), | ||
| }, | ||
| }, | ||
| } | ||
|
|
||
|
|
@@ -973,6 +1066,7 @@ def _process_transcript(self, data: dict, is_final: bool) -> None: | |
| self._speaking = False | ||
| end_event = stt.SpeechEvent(type=stt.SpeechEventType.END_OF_SPEECH) | ||
| self._event_ch.send_nowait(end_event) | ||
| self._on_end_of_speech() | ||
| else: | ||
| event = stt.SpeechEvent( | ||
| type=stt.SpeechEventType.INTERIM_TRANSCRIPT, | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2,13 +2,16 @@ | |
|
|
||
| import asyncio | ||
| from collections.abc import AsyncIterable | ||
| from typing import Any | ||
| from typing import TYPE_CHECKING, Any | ||
|
|
||
| from .. import utils | ||
| from ..types import DEFAULT_API_CONNECT_OPTIONS, NOT_GIVEN, APIConnectOptions, NotGivenOr | ||
| from ..vad import VAD, VADEventType | ||
| from .stt import STT, RecognizeStream, SpeechEvent, SpeechEventType, STTCapabilities | ||
|
|
||
| if TYPE_CHECKING: | ||
| from ..voice.events import ConversationItemAddedEvent | ||
|
|
||
| # already a retry mechanism in STT.recognize, don't retry in stream adapter | ||
| DEFAULT_STREAM_ADAPTER_API_CONNECT_OPTIONS = APIConnectOptions( | ||
| max_retry=0, timeout=DEFAULT_API_CONNECT_OPTIONS.timeout | ||
|
|
@@ -22,6 +25,8 @@ def __init__(self, *, stt: STT, vad: VAD) -> None: | |
| streaming=True, | ||
| interim_results=False, | ||
| diarization=False, # diarization requires streaming STT | ||
| keyterms=stt.capabilities.keyterms, | ||
| chat_context=stt.capabilities.chat_context, | ||
| ) | ||
| ) | ||
| self._vad = vad | ||
|
|
@@ -42,6 +47,12 @@ def model(self) -> str: | |
| def provider(self) -> str: | ||
| return self._stt.provider | ||
|
|
||
| def _update_session_keyterms(self, keyterms: list[str]) -> None: | ||
| self._stt._update_session_keyterms(keyterms) | ||
|
|
||
| def _push_conversation_item(self, item: ConversationItemAddedEvent) -> None: | ||
| self._stt._push_conversation_item(item) | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should it be a
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. using push because each call carries a single chat item as it's added and feeds it incrementally, |
||
|
|
||
| async def _recognize_impl( | ||
| self, | ||
| buffer: utils.AudioBuffer, | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.