open source
This commit is contained in:
Ajay Raj
parent
70b6e17c69
commit
a93bfc1ec9
61 changed files with 4013 additions and 126 deletions
101
vocode/streaming/transcriber/assembly_ai_transcriber.py
Normal file
101
vocode/streaming/transcriber/assembly_ai_transcriber.py
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
import websockets
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from vocode.streaming.models.transcriber import AssemblyAITranscriberConfig
|
||||
from vocode.streaming.models.websocket import AudioMessage
|
||||
from vocode.streaming.transcriber.base_transcriber import (
|
||||
BaseTranscriber,
|
||||
Transcription,
|
||||
)
|
||||
from vocode.streaming.models.audio_encoding import AudioEncoding
|
||||
|
||||
load_dotenv()
|
||||
|
||||
ASSEMBLY_AI_API_KEY = os.environ.get("ASSEMBLY_AI_API_KEY")
|
||||
ASSEMBLY_AI_URL = "wss://api.assemblyai.com/v2/realtime/ws"
|
||||
|
||||
|
||||
class AssemblyAITranscriber(BaseTranscriber):
|
||||
def __init__(
|
||||
self,
|
||||
transcriber_config: AssemblyAITranscriberConfig,
|
||||
logger: logging.Logger = None,
|
||||
):
|
||||
super().__init__(transcriber_config)
|
||||
self._ended = False
|
||||
self.is_ready = False
|
||||
self.logger = logger or logging.getLogger(__name__)
|
||||
if self.transcriber_config.should_warmup_model:
|
||||
raise Exception("AssemblyAI model warmup not supported yet")
|
||||
elif self.transcriber_config.endpointing_config:
|
||||
raise Exception("Assembly AI endpointing config not supported yet")
|
||||
|
||||
async def ready(self):
|
||||
# while not self.warmed_up:
|
||||
# await asyncio.sleep(0.1)
|
||||
# return self.is_ready
|
||||
return True
|
||||
|
||||
async def run(self):
|
||||
await self.process()
|
||||
|
||||
def send_audio(self, chunk):
|
||||
self.audio_queue.put_nowait(chunk)
|
||||
|
||||
def terminate(self):
|
||||
terminate_msg = json.dumps({"terminate_session": True})
|
||||
self.audio_queue.put_nowait(terminate_msg)
|
||||
self._ended = True
|
||||
|
||||
def get_assembly_ai_url(self):
|
||||
return ASSEMBLY_AI_URL + f"?sample_rate={self.transcriber_config.sampling_rate}"
|
||||
|
||||
async def process(self):
|
||||
self.audio_queue = asyncio.Queue()
|
||||
URL = self.get_assembly_ai_url()
|
||||
|
||||
async with websockets.connect(
|
||||
URL,
|
||||
extra_headers=(("Authorization", ASSEMBLY_AI_API_KEY),),
|
||||
ping_interval=5,
|
||||
ping_timeout=20,
|
||||
) as ws:
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
async def sender(ws): # sends audio to websocket
|
||||
while not self._ended:
|
||||
try:
|
||||
data = await asyncio.wait_for(self.audio_queue.get(), 5)
|
||||
except asyncio.exceptions.TimeoutError:
|
||||
break
|
||||
await ws.send(
|
||||
json.dumps({"audio_data": AudioMessage.from_bytes(data).data})
|
||||
)
|
||||
self.logger.debug("Terminating AssemblyAI transcriber sender")
|
||||
|
||||
async def receiver(ws):
|
||||
while not self._ended:
|
||||
try:
|
||||
result_str = await ws.recv()
|
||||
except websockets.exceptions.ConnectionClosedError as e:
|
||||
self.logger.debug(e)
|
||||
break
|
||||
except Exception as e:
|
||||
assert False, "Not a websocket 4008 error"
|
||||
|
||||
data = json.loads(result_str)
|
||||
is_final = (
|
||||
"message_type" in data
|
||||
and data["message_type"] == "FinalTranscript"
|
||||
)
|
||||
if "text" in data and data["text"]:
|
||||
await self.on_response(
|
||||
Transcription(data["text"], data["confidence"], is_final)
|
||||
)
|
||||
|
||||
await asyncio.gather(sender(ws), receiver(ws))
|
||||
59
vocode/streaming/transcriber/base_transcriber.py
Normal file
59
vocode/streaming/transcriber/base_transcriber.py
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
from dotenv import load_dotenv
|
||||
from typing import Callable, Optional, Awaitable
|
||||
|
||||
from vocode.streaming.utils import convert_wav
|
||||
from vocode.streaming.models.transcriber import EndpointingConfig, TranscriberConfig
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
class Transcription:
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
confidence: float,
|
||||
is_final: bool,
|
||||
is_interrupt: bool = False,
|
||||
):
|
||||
self.message = message
|
||||
self.confidence = confidence
|
||||
self.is_final = is_final
|
||||
self.is_interrupt = is_interrupt
|
||||
|
||||
def __str__(self):
|
||||
return f"Transcription({self.message}, {self.confidence}, {self.is_final})"
|
||||
|
||||
|
||||
class BaseTranscriber:
|
||||
def __init__(
|
||||
self,
|
||||
transcriber_config: TranscriberConfig,
|
||||
):
|
||||
self.transcriber_config = transcriber_config
|
||||
self.on_response: Optional[Callable[[Transcription], Awaitable]] = None
|
||||
|
||||
def get_transcriber_config(self) -> TranscriberConfig:
|
||||
return self.transcriber_config
|
||||
|
||||
def set_on_response(self, on_response: Callable[[Transcription], Awaitable]):
|
||||
self.on_response = on_response
|
||||
|
||||
def get_warmup_bytes(self):
|
||||
sampling_rate = self.transcriber_config.sampling_rate
|
||||
return convert_wav(
|
||||
"convo/audio/ajay.wav",
|
||||
sampling_rate,
|
||||
self.transcriber_config.audio_encoding,
|
||||
)
|
||||
|
||||
async def ready(self):
|
||||
return True
|
||||
|
||||
async def run(self):
|
||||
pass
|
||||
|
||||
def send_audio(self, chunk):
|
||||
pass
|
||||
|
||||
def terminate(self):
|
||||
pass
|
||||
230
vocode/streaming/transcriber/deepgram_transcriber.py
Normal file
230
vocode/streaming/transcriber/deepgram_transcriber.py
Normal file
|
|
@ -0,0 +1,230 @@
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from dotenv import load_dotenv
|
||||
import websockets
|
||||
from websockets.client import WebSocketClientProtocol
|
||||
import audioop
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from vocode.streaming.transcriber.base_transcriber import (
|
||||
BaseTranscriber,
|
||||
Transcription,
|
||||
)
|
||||
from vocode.streaming.models.transcriber import (
|
||||
DeepgramTranscriberConfig,
|
||||
EndpointingConfig,
|
||||
EndpointingType,
|
||||
)
|
||||
from vocode.streaming.models.audio_encoding import AudioEncoding
|
||||
|
||||
load_dotenv()
|
||||
|
||||
DEEPGRAM_API_KEY = os.environ.get("DEEPGRAM_API_KEY")
|
||||
PUNCTUATION_TERMINATORS = [".", "!", "?"]
|
||||
NUM_RESTARTS = 5
|
||||
|
||||
|
||||
class DeepgramTranscriber(BaseTranscriber):
|
||||
def __init__(
|
||||
self,
|
||||
transcriber_config: DeepgramTranscriberConfig,
|
||||
logger: logging.Logger = None,
|
||||
):
|
||||
super().__init__(transcriber_config)
|
||||
self.transcriber_config = transcriber_config
|
||||
self._ended = False
|
||||
self.warmed_up = False
|
||||
self.is_ready = False
|
||||
self.logger = logger or logging.getLogger(__name__)
|
||||
|
||||
def create_warmup_chunks(self):
|
||||
warmup_chunks = []
|
||||
warmup_bytes = self.get_warmup_bytes()
|
||||
chunk_size = self.transcriber_config.chunk_size
|
||||
for i in range(len(warmup_bytes) // chunk_size):
|
||||
warmup_chunks.append(warmup_bytes[i * chunk_size : (i + 1) * chunk_size])
|
||||
return warmup_chunks
|
||||
|
||||
async def ready(self):
|
||||
while not self.warmed_up:
|
||||
await asyncio.sleep(0.1)
|
||||
return self.is_ready
|
||||
|
||||
async def run(self):
|
||||
# warmup_chunks = await self.create_warmup_chunks()
|
||||
restarts = 0
|
||||
while not self._ended and restarts < NUM_RESTARTS:
|
||||
await self.process(self.transcriber_config.should_warmup_model)
|
||||
restarts += 1
|
||||
self.logger.debug(
|
||||
"Deepgram connection died, restarting, num_restarts: %s", restarts
|
||||
)
|
||||
|
||||
def send_audio(self, chunk):
|
||||
if (
|
||||
self.transcriber_config.downsampling
|
||||
and self.transcriber_config.audio_encoding == AudioEncoding.LINEAR16
|
||||
):
|
||||
chunk, _ = audioop.ratecv(
|
||||
chunk,
|
||||
2,
|
||||
1,
|
||||
self.transcriber_config.sampling_rate
|
||||
* self.transcriber_config.downsampling,
|
||||
self.transcriber_config.sampling_rate,
|
||||
None,
|
||||
)
|
||||
self.audio_queue.put_nowait(chunk)
|
||||
|
||||
def terminate(self):
|
||||
terminate_msg = json.dumps({"type": "CloseStream"})
|
||||
self.audio_queue.put_nowait(terminate_msg)
|
||||
self._ended = True
|
||||
|
||||
def get_deepgram_url(self):
|
||||
if self.transcriber_config.audio_encoding == AudioEncoding.LINEAR16:
|
||||
encoding = "linear16"
|
||||
elif self.transcriber_config.audio_encoding == AudioEncoding.MULAW:
|
||||
encoding = "mulaw"
|
||||
url_params = {
|
||||
"encoding": encoding,
|
||||
"sample_rate": self.transcriber_config.sampling_rate,
|
||||
"channels": 1,
|
||||
"interim_results": "true",
|
||||
}
|
||||
extra_params = {}
|
||||
if self.transcriber_config.model:
|
||||
extra_params["model"] = self.transcriber_config.model
|
||||
if self.transcriber_config.tier:
|
||||
extra_params["tier"] = self.transcriber_config.tier
|
||||
if self.transcriber_config.version:
|
||||
extra_params["version"] = self.transcriber_config.version
|
||||
if (
|
||||
self.transcriber_config.endpointing_config
|
||||
and self.transcriber_config.endpointing_config.type
|
||||
== EndpointingType.PUNCTUATION_BASED
|
||||
):
|
||||
extra_params["punctuate"] = "true"
|
||||
url_params.update(extra_params)
|
||||
return f"wss://api.deepgram.com/v1/listen?{urlencode(url_params)}"
|
||||
|
||||
def is_speech_final(
|
||||
self, current_buffer: str, deepgram_response: dict, time_silent: float
|
||||
):
|
||||
transcript = deepgram_response["channel"]["alternatives"][0]["transcript"]
|
||||
|
||||
# if it is not time based, then return true if speech is final and there is a transcript
|
||||
if not self.transcriber_config.endpointing_config:
|
||||
return transcript and deepgram_response["speech_final"]
|
||||
elif (
|
||||
self.transcriber_config.endpointing_config.type
|
||||
== EndpointingType.TIME_BASED
|
||||
):
|
||||
# if it is time based, then return true if there is no transcript
|
||||
# and there is some speech to send
|
||||
# and the time_silent is greater than the cutoff
|
||||
return (
|
||||
not transcript
|
||||
and current_buffer
|
||||
and (time_silent + deepgram_response["duration"])
|
||||
> self.transcriber_config.endpointing_config.time_cutoff_seconds
|
||||
)
|
||||
elif (
|
||||
self.transcriber_config.endpointing_config.type
|
||||
== EndpointingType.PUNCTUATION_BASED
|
||||
):
|
||||
return (
|
||||
transcript
|
||||
and deepgram_response["speech_final"]
|
||||
and transcript.strip()[-1] in PUNCTUATION_TERMINATORS
|
||||
) or (
|
||||
not transcript
|
||||
and current_buffer
|
||||
and (time_silent + deepgram_response["duration"])
|
||||
> self.transcriber_config.endpointing_config.time_cutoff_seconds
|
||||
)
|
||||
raise Exception("Endpointing config not supported")
|
||||
|
||||
def calculate_time_silent(self, data: dict):
|
||||
end = data["start"] + data["duration"]
|
||||
words = data["channel"]["alternatives"][0]["words"]
|
||||
if words:
|
||||
return end - words[-1]["end"]
|
||||
return data["duration"]
|
||||
|
||||
async def process(self, warmup=True):
|
||||
extra_headers = {"Authorization": f"Token {DEEPGRAM_API_KEY}"}
|
||||
self.audio_queue = asyncio.Queue()
|
||||
|
||||
async with websockets.connect(
|
||||
self.get_deepgram_url(), extra_headers=extra_headers
|
||||
) as ws:
|
||||
|
||||
async def warmup_sender(ws: WebSocketClientProtocol):
|
||||
if warmup:
|
||||
warmup_chunks = self.create_warmup_chunks()
|
||||
for chunk in warmup_chunks:
|
||||
await ws.send(chunk)
|
||||
await asyncio.sleep(5)
|
||||
self.warmed_up = True
|
||||
self.is_ready = True
|
||||
|
||||
async def sender(ws: WebSocketClientProtocol): # sends audio to websocket
|
||||
while not self._ended:
|
||||
try:
|
||||
data = await asyncio.wait_for(self.audio_queue.get(), 5)
|
||||
except asyncio.exceptions.TimeoutError:
|
||||
break
|
||||
await ws.send(data)
|
||||
self.logger.debug("Terminating Deepgram transcriber sender")
|
||||
|
||||
async def receiver(ws: WebSocketClientProtocol):
|
||||
buffer = ""
|
||||
time_silent = 0
|
||||
while not self._ended:
|
||||
try:
|
||||
msg = await ws.recv()
|
||||
except Exception as e:
|
||||
self.logger.debug(f"Got error {e} in Deepgram receiver")
|
||||
break
|
||||
data = json.loads(msg)
|
||||
if (
|
||||
not "is_final" in data
|
||||
): # means we've finished receiving transcriptions
|
||||
break
|
||||
is_final = data["is_final"]
|
||||
speech_final = self.is_speech_final(buffer, data, time_silent)
|
||||
top_choice = data["channel"]["alternatives"][0]
|
||||
confidence = top_choice["confidence"]
|
||||
|
||||
if (
|
||||
top_choice["transcript"]
|
||||
and confidence > 0.0
|
||||
and self.warmed_up
|
||||
and is_final
|
||||
):
|
||||
buffer = f"{buffer} {top_choice['transcript']}"
|
||||
|
||||
if speech_final:
|
||||
await self.on_response(Transcription(buffer, confidence, True))
|
||||
buffer = ""
|
||||
time_silent = 0
|
||||
elif (
|
||||
top_choice["transcript"] and confidence > 0.0 and self.warmed_up
|
||||
):
|
||||
await self.on_response(
|
||||
Transcription(
|
||||
buffer,
|
||||
confidence,
|
||||
False,
|
||||
)
|
||||
)
|
||||
time_silent = self.calculate_time_silent(data)
|
||||
else:
|
||||
time_silent += data["duration"]
|
||||
|
||||
self.logger.debug("Terminating Deepgram transcriber receiver")
|
||||
|
||||
await asyncio.gather(warmup_sender(ws), sender(ws), receiver(ws))
|
||||
145
vocode/streaming/transcriber/google_transcriber.py
Normal file
145
vocode/streaming/transcriber/google_transcriber.py
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
import asyncio
|
||||
import time
|
||||
import queue
|
||||
from google.cloud import speech
|
||||
import threading
|
||||
|
||||
from vocode.streaming.models.audio_encoding import AudioEncoding
|
||||
from vocode.streaming.transcriber.base_transcriber import (
|
||||
BaseTranscriber,
|
||||
Transcription,
|
||||
)
|
||||
from vocode.streaming.models.transcriber import GoogleTranscriberConfig
|
||||
from vocode.streaming.utils import create_loop_in_thread
|
||||
|
||||
|
||||
class GoogleTranscriber(BaseTranscriber):
|
||||
def __init__(self, transcriber_config: GoogleTranscriberConfig):
|
||||
super().__init__(transcriber_config)
|
||||
self._queue = queue.Queue()
|
||||
self._ended = False
|
||||
self.google_streaming_config = self.create_google_streaming_config()
|
||||
self.client = speech.SpeechClient()
|
||||
self.warmed_up = False
|
||||
self.is_ready = False
|
||||
if self.transcriber_config.endpointing_config:
|
||||
raise Exception("Google endpointing config not supported yet")
|
||||
self.event_loop = asyncio.new_event_loop()
|
||||
self.thread = threading.Thread(
|
||||
name="google_transcriber",
|
||||
target=create_loop_in_thread,
|
||||
args=(self.event_loop, self.process()),
|
||||
)
|
||||
|
||||
def create_google_streaming_config(self):
|
||||
extra_params = {}
|
||||
if self.transcriber_config.model:
|
||||
extra_params["model"] = self.transcriber_config.model
|
||||
extra_params["use_enhanced"] = True
|
||||
|
||||
if self.transcriber_config.audio_encoding == AudioEncoding.LINEAR16:
|
||||
google_audio_encoding = speech.RecognitionConfig.AudioEncoding.LINEAR16
|
||||
elif self.transcriber_config.audio_encoding == AudioEncoding.MULAW:
|
||||
google_audio_encoding = speech.RecognitionConfig.AudioEncoding.MULAW
|
||||
|
||||
return speech.StreamingRecognitionConfig(
|
||||
config=speech.RecognitionConfig(
|
||||
encoding=google_audio_encoding,
|
||||
sample_rate_hertz=self.transcriber_config.sampling_rate,
|
||||
language_code="en-US",
|
||||
**extra_params
|
||||
),
|
||||
interim_results=True,
|
||||
)
|
||||
|
||||
async def ready(self):
|
||||
if not self.transcriber_config.should_warmup_model:
|
||||
return True
|
||||
while not self.warmed_up:
|
||||
await asyncio.sleep(0.1)
|
||||
return self.is_ready
|
||||
|
||||
def warmup(self):
|
||||
warmup_bytes = self.get_warmup_bytes()
|
||||
|
||||
def stream():
|
||||
chunk_size = self.transcriber_config.sampling_rate * 2
|
||||
for i in range(len(warmup_bytes) // chunk_size):
|
||||
yield speech.StreamingRecognizeRequest(
|
||||
audio_content=warmup_bytes[i * chunk_size : (i + 1) * chunk_size]
|
||||
)
|
||||
time.sleep(0.01)
|
||||
|
||||
for _ in self.client.streaming_recognize(
|
||||
self.google_streaming_config, stream()
|
||||
):
|
||||
pass
|
||||
self.warmed_up = True
|
||||
self.is_ready = True
|
||||
|
||||
async def run(self):
|
||||
self.thread.start()
|
||||
|
||||
async def process(self):
|
||||
if self.transcriber_config.should_warmup_model:
|
||||
self.warmup()
|
||||
stream = self.generator()
|
||||
requests = (
|
||||
speech.StreamingRecognizeRequest(audio_content=content)
|
||||
for content in stream
|
||||
)
|
||||
responses = self.client.streaming_recognize(
|
||||
self.google_streaming_config, requests
|
||||
)
|
||||
await self.process_responses_loop(responses)
|
||||
|
||||
def terminate(self):
|
||||
self._ended = True
|
||||
|
||||
def send_audio(self, chunk: bytes):
|
||||
self._queue.put(chunk, block=False)
|
||||
|
||||
async def process_responses_loop(self, responses):
|
||||
for response in responses:
|
||||
await self._on_response(response)
|
||||
|
||||
if self._ended:
|
||||
break
|
||||
|
||||
async def _on_response(self, response):
|
||||
if not response.results:
|
||||
return
|
||||
|
||||
result = response.results[0]
|
||||
if not result.alternatives:
|
||||
return
|
||||
|
||||
top_choice = result.alternatives[0]
|
||||
message = top_choice.transcript
|
||||
confidence = top_choice.confidence
|
||||
|
||||
return await self.on_response(
|
||||
Transcription(message, confidence, result.is_final)
|
||||
)
|
||||
|
||||
def generator(self):
|
||||
while not self._ended:
|
||||
# Use a blocking get() to ensure there's at least one chunk of
|
||||
# data, and stop iteration if the chunk is None, indicating the
|
||||
# end of the audio stream.
|
||||
chunk = self._queue.get()
|
||||
if chunk is None:
|
||||
return
|
||||
data = [chunk]
|
||||
|
||||
# Now consume whatever other data's still buffered.
|
||||
while True:
|
||||
try:
|
||||
chunk = self._queue.get(block=False)
|
||||
if chunk is None:
|
||||
return
|
||||
data.append(chunk)
|
||||
except queue.Empty:
|
||||
break
|
||||
|
||||
yield b"".join(data)
|
||||
Loading…
Add table
Add a link
Reference in a new issue