78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
import asyncio
|
|
import copy
|
|
import functools
|
|
from typing import AsyncGenerator, Awaitable, Callable
|
|
from dailyai.queue_aggregators import LLMAssistantContextAggregator, LLMContextAggregator, LLMUserContextAggregator
|
|
from dailyai.queue_frame import EndStreamQueueFrame, QueueFrame, TranscriptionQueueFrame
|
|
|
|
|
|
class InterruptibleConversationWrapper:
|
|
|
|
def __init__(
|
|
self,
|
|
frame_generator: Callable[[], AsyncGenerator[QueueFrame, None]],
|
|
runner: Callable[
|
|
[str, LLMContextAggregator, LLMContextAggregator], Awaitable[None]
|
|
],
|
|
interrupt: Callable[[], None],
|
|
my_participant_id: str | None,
|
|
llm_messages: list[dict[str, str]],
|
|
llm_context_aggregator_in=LLMUserContextAggregator,
|
|
llm_context_aggregator_out=LLMAssistantContextAggregator,
|
|
delay_before_speech_seconds: float = 1.0,
|
|
):
|
|
self._frame_generator: Callable[[], AsyncGenerator[QueueFrame, None]] = frame_generator
|
|
self._runner: Callable[
|
|
[str, LLMContextAggregator, LLMContextAggregator], Awaitable[None]
|
|
] = runner
|
|
self._interrupt: Callable[[], None] = interrupt
|
|
self._my_participant_id = my_participant_id
|
|
self._messages: list[dict[str, str]] = llm_messages
|
|
self._delay_before_speech_seconds = delay_before_speech_seconds
|
|
self._llm_context_aggregator_in = llm_context_aggregator_in
|
|
self._llm_context_aggregator_out = llm_context_aggregator_out
|
|
|
|
self._current_phrase = ""
|
|
|
|
def update_messages(self, new_messages: list[dict[str, str]], task: asyncio.Task | None):
|
|
if task:
|
|
if not task.cancelled():
|
|
self._current_phrase = ""
|
|
self._messages = new_messages
|
|
|
|
async def speak_after_delay(self, user_speech, messages):
|
|
await asyncio.sleep(self._delay_before_speech_seconds)
|
|
tma_in = self._llm_context_aggregator_in(
|
|
messages, self._my_participant_id, complete_sentences=False
|
|
)
|
|
tma_out = self._llm_context_aggregator_out(
|
|
messages, self._my_participant_id
|
|
)
|
|
|
|
await self._runner(user_speech, tma_in, tma_out)
|
|
|
|
async def run_conversation(self):
|
|
current_response_task = None
|
|
|
|
async for frame in self._frame_generator():
|
|
if isinstance(frame, EndStreamQueueFrame):
|
|
break
|
|
elif not isinstance(frame, TranscriptionQueueFrame):
|
|
continue
|
|
|
|
if frame.participantId == self._my_participant_id:
|
|
continue
|
|
|
|
if current_response_task:
|
|
current_response_task.cancel()
|
|
self._interrupt()
|
|
|
|
self._current_phrase += " " + frame.text
|
|
current_llm_messages = copy.deepcopy(self._messages)
|
|
current_response_task = asyncio.create_task(
|
|
self.speak_after_delay(self._current_phrase, current_llm_messages)
|
|
)
|
|
current_response_task.add_done_callback(
|
|
functools.partial(self.update_messages, current_llm_messages)
|
|
)
|