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

Remove run_to_queue and run from AIService class

Browse Source
This commit is contained in:
Moishe Lettvin
2024-03-15 11:04:22 -04:00
parent 358166f347
commit b8b35db89c
15 changed files with 182 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
src/dailyai/pipeline/merge_pipeline.py Normal file
View File

@@ -0,0 +1,21 @@
from typing import List
from dailyai.pipeline.frames import EndFrame, EndPipeFrame
from dailyai.pipeline.pipeline import Pipeline
class SequentialMergePipeline(Pipeline):
"""This class merges the sink queues from a list of pipelines. Frames from
each pipeline's sink are merged in the order of pipelines in the list."""
def __init__(self, pipelines:List[Pipeline]):
super().__init__([])
self.pipelines = pipelines
async def run_pipeline(self):
for pipeline in self.pipelines:
while True:
frame = await pipeline.sink.get()
if isinstance(frame, EndFrame) or isinstance(frame, EndPipeFrame):
break
await self.sink.put(frame)
await self.sink.put(EndFrame())

74
src/dailyai/pipeline/pipeline.py
View File

@@ -1,5 +1,5 @@
import asyncio
from typing import AsyncGenerator, List
from typing import AsyncGenerator, AsyncIterable, Iterable, List
from dailyai.pipeline.frame_processor import FrameProcessor
from dailyai.pipeline.frames import EndPipeFrame, EndFrame, Frame
@@ -17,17 +17,17 @@ class Pipeline:
self,
processors: List[FrameProcessor],
source: asyncio.Queue | None = None,
sink: asyncio.Queue[Frame] | None = None,
sink: asyncio.Queue[Frame] | None = None
):
"""Create a new pipeline. By default neither the source nor sink
queues are set, so you'll need to pass them to this constructor or
call set_source and set_sink before using the pipeline. Note that
the transport's run_*_pipeline methods will set the source and sink
queues on the pipeline for you.
"""Create a new pipeline. By default we create the sink and source queues
if they're not provided, but these can be overridden to point to other
queues. If this pipeline is run by a transport, its sink and source queues
will be overridden.
"""
self.processors = processors
self.source: asyncio.Queue[Frame] | None = source
self.sink: asyncio.Queue[Frame] | None = sink
self.processors: List[FrameProcessor] = processors
self.source: asyncio.Queue[Frame] = source or asyncio.Queue()
self.sink: asyncio.Queue[Frame] = sink or asyncio.Queue()
def set_source(self, source: asyncio.Queue[Frame]):
"""Set the source queue for this pipeline. Frames from this queue
@@ -44,21 +44,24 @@ class Pipeline:
"""Convenience function to get the next frame from the source queue. This
lets us consistently have an AsyncGenerator yield frames, from either the
source queue or a frame_processor."""
if self.source is None:
raise ValueError("Source queue not set")
yield await self.source.get()
async def run_pipeline_recursively(
self, initial_frame: Frame, processors: List[FrameProcessor]
) -> AsyncGenerator[Frame, None]:
if processors:
async for frame in processors[0].process_frame(initial_frame):
async for final_frame in self.run_pipeline_recursively(
frame, processors[1:]
):
yield final_frame
async def queue_frames(
self,
frames: Iterable[Frame] | AsyncIterable[Frame],
) -> None:
"""Insert frames directly into a pipeline. This is typically used inside a transport
participant_joined callback to prompt a bot to start a conversation, for example."""
if isinstance(frames, AsyncIterable):
async for frame in frames:
await self.source.put(frame)
elif isinstance(frames, Iterable):
for frame in frames:
await self.source.put(frame)
else:
yield initial_frame
raise Exception("Frames must be an iterable or async iterable")
async def run_pipeline(self):
"""Run the pipeline. Take each frame from the source queue, pass it to
@@ -73,13 +76,10 @@ class Pipeline:
if it's not the last frame yielded by the last frame_processor in the pipeline..
"""
if self.source is None or self.sink is None:
raise ValueError("Source or sink queue not set")
try:
while True:
initial_frame = await self.source.get()
async for frame in self.run_pipeline_recursively(
async for frame in self._run_pipeline_recursively(
initial_frame, self.processors
):
await self.sink.put(frame)
@@ -94,11 +94,17 @@ class Pipeline:
await processor.interrupted()
pass
async def queue_frames(self, frames: Frame | List[Frame]):
"""Insert frames directly into a pipeline. This is typically used inside a transport
participant_joined callback to prompt a bot to start a conversation, for example.
"""
if not isinstance(frames, List):
frames = [frames]
for f in frames:
await self.source.put(f)
async def _run_pipeline_recursively(
self, initial_frame: Frame, processors: List[FrameProcessor]
) -> AsyncGenerator[Frame, None]:
"""Internal function to add frames to the pipeline as they're yielded
by each processor."""
if processors:
async for frame in processors[0].process_frame(initial_frame):
async for final_frame in self._run_pipeline_recursively(
frame, processors[1:]
):
yield final_frame
else:
yield initial_frame

51
src/dailyai/services/ai_services.py
View File

@@ -8,6 +8,7 @@ from dailyai.pipeline.frame_processor import FrameProcessor
from dailyai.pipeline.frames import (
AudioFrame,
EndFrame,
EndPipeFrame,
ImageFrame,
LLMMessagesQueueFrame,
LLMResponseEndFrame,
@@ -20,53 +21,13 @@ from dailyai.pipeline.frames import (
)
from abc import abstractmethod
from typing import AsyncGenerator, AsyncIterable, BinaryIO, Iterable, List
from typing import AsyncGenerator, BinaryIO
class AIService(FrameProcessor):
def __init__(self):
self.logger = logging.getLogger("dailyai")
def stop(self):
pass
async def run_to_queue(
self, queue: asyncio.Queue, frames, add_end_of_stream=False
) -> None:
async for frame in self.run(frames):
await queue.put(frame)
if add_end_of_stream:
await queue.put(EndFrame())
async def run(
self,
frames: Iterable[Frame] | AsyncIterable[Frame] | asyncio.Queue[Frame],
) -> AsyncGenerator[Frame, None]:
try:
if isinstance(frames, AsyncIterable):
async for frame in frames:
async for output_frame in self.process_frame(frame):
yield output_frame
elif isinstance(frames, Iterable):
for frame in frames:
async for output_frame in self.process_frame(frame):
yield output_frame
elif isinstance(frames, asyncio.Queue):
while True:
frame = await frames.get()
async for output_frame in self.process_frame(frame):
yield output_frame
if isinstance(frame, EndFrame):
break
else:
raise Exception("Frames must be an iterable or async iterable")
except Exception as e:
self.logger.error("Exception occurred while running AI service", e)
raise e
class LLMService(AIService):
"""This class is a no-op but serves as a base class for LLM services."""
@@ -92,7 +53,7 @@ class TTSService(AIService):
yield bytes()
async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]:
if isinstance(frame, EndFrame):
if isinstance(frame, EndFrame) or isinstance(frame, EndPipeFrame):
if self.current_sentence:
async for audio_chunk in self.run_tts(self.current_sentence):
yield AudioFrame(audio_chunk)
@@ -118,12 +79,6 @@ class TTSService(AIService):
# note we pass along the text frame *after* the audio, so the text frame is completed after the audio is processed.
yield TextFrame(text)
# Convenience function to send the audio for a sentence to the given queue
async def say(self, sentence, queue: asyncio.Queue):
await self.run_to_queue(
queue, [LLMResponseStartFrame(), TextFrame(sentence), LLMResponseEndFrame()]
)
class ImageGenService(AIService):
def __init__(self, image_size, **kwargs):

28
src/dailyai/services/base_transport_service.py
View File

@@ -20,11 +20,12 @@ from dailyai.pipeline.frames import (
PipelineStartedFrame,
SpriteFrame,
StartFrame,
TranscriptionQueueFrame,
TextFrame,
UserStartedSpeakingFrame,
UserStoppedSpeakingFrame,
)
from dailyai.pipeline.pipeline import Pipeline
from dailyai.services.ai_services import TTSService
torch.set_num_threads(1)
@@ -125,7 +126,7 @@ class BaseTransportService:
self._logger: logging.Logger = logging.getLogger()
async def run(self):
async def run(self, pipeline:Pipeline | None=None, override_pipeline_source_queue=True):
self._prerun()
async_output_queue_marshal_task = asyncio.create_task(self._marshal_frames())
@@ -148,6 +149,13 @@ class BaseTransportService:
self._vad_thread = threading.Thread(target=self._vad, daemon=True)
self._vad_thread.start()
pipeline_task = None
if pipeline:
pipeline.set_sink(self.send_queue)
if override_pipeline_source_queue:
pipeline.set_source(self.receive_queue)
pipeline_task = asyncio.create_task(pipeline.run_pipeline())
try:
while time.time() < self._expiration and not self._stop_threads.is_set():
await asyncio.sleep(1)
@@ -160,9 +168,12 @@ class BaseTransportService:
self._stop_threads.set()
if pipeline_task:
pipeline_task.cancel()
await self.send_queue.put(EndFrame())
await async_output_queue_marshal_task
await self.send_queue.join()
self._frame_consumer_thread.join()
if self._speaker_enabled:
@@ -171,11 +182,6 @@ class BaseTransportService:
if self._vad_enabled:
self._vad_thread.join()
async def run_uninterruptible_pipeline(self, pipeline: Pipeline):
pipeline.set_sink(self.send_queue)
pipeline.set_source(self.receive_queue)
await pipeline.run_pipeline()
async def run_interruptible_pipeline(
self,
pipeline: Pipeline,
@@ -232,6 +238,11 @@ class BaseTransportService:
await asyncio.gather(pipeline_task, post_process_task)
async def say(self, text:str, tts:TTSService):
"""Say a phrase. Use with caution; this bypasses any running pipelines."""
async for frame in tts.process_frame(TextFrame(text)):
await self.send_queue.put(frame)
def _post_run(self):
# Note that this function must be idempotent! It can be called multiple times
# if, for example, a keyboard interrupt occurs.
@@ -399,6 +410,7 @@ class BaseTransportService:
for frame in frames:
if isinstance(frame, EndFrame):
self._logger.info("Stopping frame consumer thread")
self._stop_threads.set()
self._threadsafe_send_queue.task_done()
if self._loop:
asyncio.run_coroutine_threadsafe(

3
src/dailyai/services/daily_transport_service.py
View File

@@ -219,7 +219,8 @@ class DailyTransportService(BaseTransportService, EventHandler):
pass
def call_joined(self, join_data, client_error):
self._logger.info(f"Call_joined: {join_data}, {client_error}")
#self._logger.info(f"Call_joined: {join_data}, {client_error}")
pass
def dialout(self, number):
self.client.start_dialout({"phoneNumber": number})

26
src/dailyai/tests/test_ai_services.py
View File

@@ -12,7 +12,7 @@ class SimpleAIService(AIService):
class TestBaseAIService(unittest.IsolatedAsyncioTestCase):
async def test_async_input(self):
async def test_simple_processing(self):
service = SimpleAIService()
input_frames = [
@@ -20,28 +20,10 @@ class TestBaseAIService(unittest.IsolatedAsyncioTestCase):
EndFrame()
]
async def iterate_frames() -> AsyncGenerator[Frame, None]:
for frame in input_frames:
yield frame
output_frames = []
async for frame in service.run(iterate_frames()):
output_frames.append(frame)
self.assertEqual(input_frames, output_frames)
async def test_nonasync_input(self):
service = SimpleAIService()
input_frames = [TextFrame("hello"), EndFrame()]
def iterate_frames() -> Generator[Frame, None, None]:
for frame in input_frames:
yield frame
output_frames = []
async for frame in service.run(iterate_frames()):
output_frames.append(frame)
for input_frame in input_frames:
async for output_frame in service.process_frame(input_frame):
output_frames.append(output_frame)
self.assertEqual(input_frames, output_frames)