wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

transports(websocket): base class from BaseInputTransport

Browse Source
This commit is contained in:
Aleix Conchillo Flaqué
2024-05-31 11:29:37 -07:00
parent 54fccd2e25
commit 428c8af77e
5 changed files with 78 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/pipecat/transports/base_input.py
View File

@@ -21,7 +21,7 @@ from pipecat.frames.frames import (
UserStartedSpeakingFrame,
UserStoppedSpeakingFrame)
from pipecat.transports.base_transport import TransportParams
from pipecat.vad.vad_analyzer import VADState
from pipecat.vad.vad_analyzer import VADAnalyzer, VADState
from loguru import logger
@@ -74,10 +74,10 @@ class BaseInputTransport(FrameProcessor):
self._push_frame_task.cancel()
def vad_analyze(self, audio_frames: bytes) -> VADState:
pass
def vad_analyzer(self) -> VADAnalyzer | None:
return self._params.vad_analyzer
def read_raw_audio_frames(self, frame_count: int) -> bytes:
def read_next_audio_frame(self) -> AudioRawFrame | None:
pass
#
@@ -146,8 +146,15 @@ class BaseInputTransport(FrameProcessor):
# Audio input
#
def _vad_analyze(self, audio_frames: bytes) -> VADState:
state = VADState.QUIET
vad_analyzer = self.vad_analyzer()
if vad_analyzer:
state = vad_analyzer.analyze_audio(audio_frames)
return state
def _handle_vad(self, audio_frames: bytes, vad_state: VADState):
new_vad_state = self.vad_analyze(audio_frames)
new_vad_state = self._vad_analyze(audio_frames)
if new_vad_state != vad_state and new_vad_state != VADState.STARTING and new_vad_state != VADState.STOPPING:
frame = None
if new_vad_state == VADState.SPEAKING:
@@ -165,19 +172,11 @@ class BaseInputTransport(FrameProcessor):
def _audio_thread_handler(self):
vad_state: VADState = VADState.QUIET
sample_rate = self._params.audio_in_sample_rate
num_channels = self._params.audio_in_channels
num_frames = int(sample_rate / 100) # 10ms of audio
while self._running:
try:
audio_frames = self.read_raw_audio_frames(num_frames)
if len(audio_frames) > 0:
frame = AudioRawFrame(
audio=audio_frames,
sample_rate=sample_rate,
num_channels=num_channels)
frame = self.read_next_audio_frame()
if frame:
audio_passthrough = True
# Check VAD and push event if necessary. We just care about

12
src/pipecat/transports/local/audio.py
View File

@@ -6,7 +6,7 @@
import asyncio
from pipecat.frames.frames import StartFrame
from pipecat.frames.frames import AudioRawFrame, StartFrame
from pipecat.processors.frame_processor import FrameProcessor
from pipecat.transports.base_input import BaseInputTransport
from pipecat.transports.base_output import BaseOutputTransport
@@ -35,8 +35,14 @@ class LocalAudioInputTransport(BaseInputTransport):
frames_per_buffer=params.audio_in_sample_rate,
input=True)
def read_raw_audio_frames(self, frame_count: int) -> bytes:
return self._in_stream.read(frame_count, exception_on_overflow=False)
def read_next_audio_frame(self) -> AudioRawFrame | None:
sample_rate = self._params.audio_in_sample_rate
num_channels = self._params.audio_in_channels
num_frames = int(sample_rate / 100) # 10ms of audio
audio = self._in_stream.read(num_frames, exception_on_overflow=False)
return AudioRawFrame(audio=audio, sample_rate=sample_rate, num_channels=num_channels)
async def start(self, frame: StartFrame):
await super().start(frame)

12
src/pipecat/transports/local/tk.py
View File

@@ -9,7 +9,7 @@ import asyncio
import numpy as np
import tkinter as tk
from pipecat.frames.frames import ImageRawFrame, StartFrame
from pipecat.frames.frames import AudioRawFrame, ImageRawFrame, StartFrame
from pipecat.processors.frame_processor import FrameProcessor
from pipecat.transports.base_input import BaseInputTransport
from pipecat.transports.base_output import BaseOutputTransport
@@ -45,8 +45,14 @@ class TkInputTransport(BaseInputTransport):
frames_per_buffer=params.audio_in_sample_rate,
input=True)
def read_raw_audio_frames(self, frame_count: int) -> bytes:
return self._in_stream.read(frame_count, exception_on_overflow=False)
def read_next_audio_frame(self) -> AudioRawFrame | None:
sample_rate = self._params.audio_in_sample_rate
num_channels = self._params.audio_in_channels
num_frames = int(sample_rate / 100) # 10ms of audio
audio = self._in_stream.read(num_frames, exception_on_overflow=False)
return AudioRawFrame(audio=audio, sample_rate=sample_rate, num_channels=num_channels)
async def start(self, frame: StartFrame):
await super().start(frame)

94
src/pipecat/transports/network/websocket_server.py
View File

@@ -7,23 +7,18 @@
import asyncio
import io
import queue
import wave
import websockets
from typing import Awaitable, Callable
from pydantic.main import BaseModel
from pipecat.frames.frames import (
AudioRawFrame,
CancelFrame,
EndFrame,
Frame,
StartFrame,
TTSStartedFrame,
TTSStoppedFrame)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.frames.frames import AudioRawFrame, StartFrame
from pipecat.processors.frame_processor import FrameProcessor
from pipecat.serializers.base_serializer import FrameSerializer
from pipecat.serializers.protobuf import ProtobufFrameSerializer
from pipecat.transports.base_input import BaseInputTransport
from pipecat.transports.base_output import BaseOutputTransport
from pipecat.transports.base_transport import BaseTransport, TransportParams
@@ -40,7 +35,7 @@ class WebsocketServerCallbacks(BaseModel):
on_connection: Callable[[websockets.WebSocketServerProtocol], Awaitable[None]]
class WebsocketServerInputTransport(FrameProcessor):
class WebsocketServerInputTransport(BaseInputTransport):
def __init__(
self,
@@ -48,7 +43,7 @@ class WebsocketServerInputTransport(FrameProcessor):
port: int,
params: WebsocketServerParams,
callbacks: WebsocketServerCallbacks):
super().__init__()
super().__init__(params)
self._host = host
self._port = port
@@ -57,25 +52,23 @@ class WebsocketServerInputTransport(FrameProcessor):
self._websocket: websockets.WebSocketServerProtocol | None = None
self._client_audio_queue = queue.Queue()
self._stop_server_event = asyncio.Event()
# Create push frame task. This is the task that will push frames in
# order. We also guarantee that all frames are pushed in the same task.
self._create_push_task()
async def start(self, frame: StartFrame):
self._server_task = self.get_event_loop().create_task(self._server_task_handler())
await super().start(frame)
async def process_frame(self, frame: Frame, direction: FrameDirection):
if isinstance(frame, CancelFrame):
await self._stop()
# We don't queue a CancelFrame since we want to stop ASAP.
await self.push_frame(frame, direction)
elif isinstance(frame, StartFrame):
await self._start()
await self._internal_push_frame(frame, direction)
elif isinstance(frame, EndFrame):
await self._stop()
await self._internal_push_frame(frame, direction)
else:
await self._internal_push_frame(frame, direction)
async def stop(self):
self._stop_server_event.set()
await self._server_task
await super().stop()
def read_next_audio_frame(self) -> AudioRawFrame | None:
try:
return self._client_audio_queue.get(timeout=1)
except queue.Empty:
return None
async def _server_task_handler(self):
logger.info(f"Starting websocket server on {self._host}:{self._port}")
@@ -96,44 +89,16 @@ class WebsocketServerInputTransport(FrameProcessor):
# Handle incoming messages
async for message in websocket:
frame = self._params.serializer.deserialize(message)
await self._internal_push_frame(frame)
if isinstance(frame, AudioRawFrame) and self._params.audio_in_enabled:
self._client_audio_queue.put_nowait(frame)
else:
await self._internal_push_frame(frame)
await self._websocket.close()
self._websocket = None
logger.info(f"Client {websocket.remote_address} disconnected")
async def _start(self):
loop = self.get_event_loop()
self._server_task = loop.create_task(self._server_task_handler())
async def _stop(self):
self._stop_server_event.set()
self._push_frame_task.cancel()
await self._server_task
#
# Push frames task
#
def _create_push_task(self):
loop = self.get_event_loop()
self._push_frame_task = loop.create_task(self._push_frame_task_handler())
self._push_queue = asyncio.Queue()
async def _internal_push_frame(
self,
frame: Frame | None,
direction: FrameDirection | None = FrameDirection.DOWNSTREAM):
await self._push_queue.put((frame, direction))
async def _push_frame_task_handler(self):
running = True
while running:
try:
(frame, direction) = await self._push_queue.get()
await self.push_frame(frame, direction)
running = not isinstance(frame, EndFrame)
except asyncio.CancelledError:
break
class WebsocketServerOutputTransport(BaseOutputTransport):
@@ -177,7 +142,10 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
frame = wav_frame
proto = self._params.serializer.serialize(frame)
asyncio.run_coroutine_threadsafe(self._websocket.send(proto), self.get_event_loop())
future = asyncio.run_coroutine_threadsafe(
self._websocket.send(proto), self.get_event_loop())
future.result()
self._audio_buffer = self._audio_buffer[self._params.audio_frame_size:]

26
src/pipecat/transports/services/daily.py
View File

@@ -188,15 +188,22 @@ class DailyTransportClient(EventHandler):
def send_message(self, frame: DailyTransportMessageFrame):
self._client.send_app_message(frame.message, frame.participant_id)
def read_raw_audio_frames(self, frame_count: int) -> bytes:
def read_next_audio_frame(self) -> AudioRawFrame | None:
sample_rate = self._params.audio_in_sample_rate
num_channels = self._params.audio_in_channels
if self._other_participant_has_joined:
return self._speaker.read_frames(frame_count)
num_frames = int(sample_rate / 100) # 10ms of audio
audio = self._speaker.read_frames(num_frames)
return AudioRawFrame(audio=audio, sample_rate=sample_rate, num_channels=num_channels)
else:
# If no one has ever joined the meeting `read_frames()` would block,
# instead we just wait a bit. daily-python should probably return
# silence instead.
time.sleep(0.01)
return b''
return None
def write_raw_audio_frames(self, frames: bytes):
self._mic.write_frames(frames)
@@ -467,7 +474,7 @@ class DailyInputTransport(BaseInputTransport):
self._video_renderers = {}
self._camera_in_queue = queue.Queue()
self._vad_analyzer = params.vad_analyzer
self._vad_analyzer: VADAnalyzer | None = params.vad_analyzer
if params.vad_enabled and not params.vad_analyzer:
self._vad_analyzer = WebRTCVADAnalyzer(
sample_rate=self._params.audio_in_sample_rate,
@@ -498,14 +505,11 @@ class DailyInputTransport(BaseInputTransport):
await super().cleanup()
await self._client.cleanup()
def vad_analyze(self, audio_frames: bytes) -> VADState:
state = VADState.QUIET
if self._vad_analyzer:
state = self._vad_analyzer.analyze_audio(audio_frames)
return state
def vad_analyzer(self) -> VADAnalyzer | None:
return self._vad_analyzer
def read_raw_audio_frames(self, frame_count: int) -> bytes:
return self._client.read_raw_audio_frames(frame_count)
def read_next_audio_frame(self) -> AudioRawFrame | None:
return self._client.read_next_audio_frame()
#
# FrameProcessor