sample 05 works
This commit is contained in:
@@ -14,11 +14,6 @@ from dailyai.services.ai_services import AIService, PipeService
|
|||||||
|
|
||||||
from typing import Any, AsyncGenerator, Callable, List, Tuple
|
from typing import Any, AsyncGenerator, Callable, List, Tuple
|
||||||
|
|
||||||
@dataclass
|
|
||||||
class SinkAndQueuePair:
|
|
||||||
sink: PipeService
|
|
||||||
queue: asyncio.Queue
|
|
||||||
|
|
||||||
|
|
||||||
class QueueTee(PipeService):
|
class QueueTee(PipeService):
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -26,11 +21,10 @@ class QueueTee(PipeService):
|
|||||||
):
|
):
|
||||||
super().__init__(*args, **kwargs)
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
self.sinks: List[SinkAndQueuePair] = []
|
self.sinks: List[PipeService] = []
|
||||||
for sink in sinks:
|
for sink in sinks:
|
||||||
pair = SinkAndQueuePair(sink, asyncio.Queue())
|
sink.source_queue = asyncio.Queue()
|
||||||
self.sinks.append(pair)
|
self.sinks.append(sink)
|
||||||
sink.source_queue = pair.queue
|
|
||||||
|
|
||||||
async def process_queue(self):
|
async def process_queue(self):
|
||||||
if not self.source_queue:
|
if not self.source_queue:
|
||||||
@@ -38,10 +32,10 @@ class QueueTee(PipeService):
|
|||||||
|
|
||||||
while True:
|
while True:
|
||||||
frame: QueueFrame = await self.source_queue.get()
|
frame: QueueFrame = await self.source_queue.get()
|
||||||
print("got frame")
|
|
||||||
for sink in self.sinks:
|
for sink in self.sinks:
|
||||||
print("putting frame in sink")
|
if sink.source_queue:
|
||||||
await sink.queue.put(frame)
|
await sink.source_queue.put(frame)
|
||||||
|
|
||||||
if isinstance(frame, EndStreamQueueFrame):
|
if isinstance(frame, EndStreamQueueFrame):
|
||||||
break
|
break
|
||||||
|
|
||||||
@@ -63,10 +57,6 @@ class QueueFrameAggregator(PipeService):
|
|||||||
async def process_frame(
|
async def process_frame(
|
||||||
self, frame: QueueFrame
|
self, frame: QueueFrame
|
||||||
) -> AsyncGenerator[QueueFrame, None]:
|
) -> AsyncGenerator[QueueFrame, None]:
|
||||||
if isinstance(frame, ControlQueueFrame):
|
|
||||||
yield frame
|
|
||||||
return
|
|
||||||
|
|
||||||
output_frame: QueueFrame | None = None
|
output_frame: QueueFrame | None = None
|
||||||
(self.aggregation, output_frame) = self.aggregator(
|
(self.aggregation, output_frame) = self.aggregator(
|
||||||
self.aggregation, frame
|
self.aggregation, frame
|
||||||
@@ -92,8 +82,6 @@ class QueueMergeGateOnFirst(PipeService):
|
|||||||
*[source_queue.get() for source_queue in self.source_queues]
|
*[source_queue.get() for source_queue in self.source_queues]
|
||||||
)
|
)
|
||||||
for idx, frame in enumerate(frames):
|
for idx, frame in enumerate(frames):
|
||||||
print("frame", idx, frame)
|
|
||||||
|
|
||||||
# if the frame we got from a source is an EndStreamQueueFrame, remove that source
|
# if the frame we got from a source is an EndStreamQueueFrame, remove that source
|
||||||
if isinstance(frame, EndStreamQueueFrame):
|
if isinstance(frame, EndStreamQueueFrame):
|
||||||
self.source_queues.pop(idx)
|
self.source_queues.pop(idx)
|
||||||
|
|||||||
@@ -47,7 +47,6 @@ class PipeService(AbstractPipeService):
|
|||||||
frame: QueueFrame = await self.source_queue.get()
|
frame: QueueFrame = await self.source_queue.get()
|
||||||
async for output_frame in self.process_frame(frame):
|
async for output_frame in self.process_frame(frame):
|
||||||
if isinstance(frame, EndStreamQueueFrame):
|
if isinstance(frame, EndStreamQueueFrame):
|
||||||
print("end of stream", type(self))
|
|
||||||
async for final_frame in self.finalize():
|
async for final_frame in self.finalize():
|
||||||
await self.sink_queue.put(final_frame)
|
await self.sink_queue.put(final_frame)
|
||||||
await self.sink_queue.put(output_frame)
|
await self.sink_queue.put(output_frame)
|
||||||
|
|||||||
@@ -32,9 +32,9 @@ async def main(room_url):
|
|||||||
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
/ TTS \
|
/ TTS \
|
||||||
Month prompt -> LLM -> Fork -> -> Gate -> Transport
|
Month prompt -> LLM -> Fork -> -> Gate -> Transport
|
||||||
\ ImageGen /
|
\ Aggregate -> ImageGen /
|
||||||
"""
|
"""
|
||||||
|
|
||||||
month_description_queue: asyncio.Queue[QueueFrame] = asyncio.Queue()
|
month_description_queue: asyncio.Queue[QueueFrame] = asyncio.Queue()
|
||||||
@@ -60,6 +60,9 @@ async def main(room_url):
|
|||||||
def aggregator(
|
def aggregator(
|
||||||
accumulation, frame: QueueFrame
|
accumulation, frame: QueueFrame
|
||||||
) -> tuple[Any, QueueFrame | None]:
|
) -> tuple[Any, QueueFrame | None]:
|
||||||
|
if not accumulation:
|
||||||
|
accumulation = ""
|
||||||
|
|
||||||
if isinstance(frame, TextQueueFrame):
|
if isinstance(frame, TextQueueFrame):
|
||||||
accumulation += frame.text
|
accumulation += frame.text
|
||||||
return (accumulation, None)
|
return (accumulation, None)
|
||||||
@@ -70,7 +73,7 @@ async def main(room_url):
|
|||||||
|
|
||||||
# This queue service takes chunks from LLM output and merges them into one text frame
|
# This queue service takes chunks from LLM output and merges them into one text frame
|
||||||
# that will be used to prompt the image service.
|
# that will be used to prompt the image service.
|
||||||
llm_aggregator_for_image = QueueFrameAggregator(source_queue=llm.sink_queue, aggregator=aggregator, finalizer=lambda x: None)
|
llm_aggregator_for_image = QueueFrameAggregator(aggregator=aggregator, finalizer=lambda x: None)
|
||||||
|
|
||||||
# Set the source queue for the image service to the sink of the aggregator service
|
# Set the source queue for the image service to the sink of the aggregator service
|
||||||
dalle.source_queue = llm_aggregator_for_image.sink_queue
|
dalle.source_queue = llm_aggregator_for_image.sink_queue
|
||||||
@@ -89,7 +92,7 @@ async def main(room_url):
|
|||||||
tts_image_gate.sink_queue = transport.send_queue
|
tts_image_gate.sink_queue = transport.send_queue
|
||||||
|
|
||||||
# Queue up all the months in the LLM service source queue
|
# Queue up all the months in the LLM service source queue
|
||||||
months = ["January", "February"]
|
months = ["January"] #, "February"]
|
||||||
for month in months:
|
for month in months:
|
||||||
messages = [
|
messages = [
|
||||||
{
|
{
|
||||||
@@ -102,7 +105,7 @@ async def main(room_url):
|
|||||||
|
|
||||||
await month_description_queue.put(EndStreamQueueFrame())
|
await month_description_queue.put(EndStreamQueueFrame())
|
||||||
|
|
||||||
await asyncio.gather(transport.run(), *[service.process_queue() for service in [llm, tts, dalle, tee, tts_image_gate]])
|
await asyncio.gather(transport.run(), *[service.process_queue() for service in [llm, tts, dalle, tee, tts_image_gate, llm_aggregator_for_image]])
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
(url, token) = configure()
|
(url, token) = configure()
|
||||||
|
|||||||
Reference in New Issue
Block a user