function calling now run in tasks
This commit is contained in:
@@ -8,7 +8,8 @@ import asyncio
|
||||
import io
|
||||
import wave
|
||||
from abc import abstractmethod
|
||||
from typing import Any, AsyncGenerator, Dict, List, Mapping, Optional, Sequence, Tuple, Type
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, AsyncGenerator, Dict, List, Mapping, Optional, Sequence, Set, Tuple, Type
|
||||
|
||||
from loguru import logger
|
||||
|
||||
@@ -22,6 +23,9 @@ from pipecat.frames.frames import (
|
||||
EndFrame,
|
||||
ErrorFrame,
|
||||
Frame,
|
||||
FunctionCallCancelFrame,
|
||||
FunctionCallInProgressFrame,
|
||||
FunctionCallResultFrame,
|
||||
InterimTranscriptionFrame,
|
||||
LLMFullResponseEndFrame,
|
||||
StartFrame,
|
||||
@@ -138,6 +142,13 @@ class AIService(FrameProcessor):
|
||||
await self.push_frame(f)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FunctionEntry:
|
||||
function_name: Optional[str]
|
||||
callback: Any # TODO(aleix): add proper typing.
|
||||
cancel_on_interruption: bool
|
||||
|
||||
|
||||
class LLMService(AIService):
|
||||
"""This class is a no-op but serves as a base class for LLM services."""
|
||||
|
||||
@@ -147,38 +158,74 @@ class LLMService(AIService):
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._callbacks = {}
|
||||
self._functions = {}
|
||||
self._start_callbacks = {}
|
||||
self._adapter = self.adapter_class()
|
||||
self._function_call_tasks: Set[Tuple[asyncio.Task, str, str]] = set()
|
||||
|
||||
self._register_event_handler("on_completion_timeout")
|
||||
|
||||
def get_llm_adapter(self) -> BaseLLMAdapter:
|
||||
return self._adapter
|
||||
|
||||
def create_context_aggregator(
|
||||
self, context: OpenAILLMContext, *, assistant_expect_stripped_words: bool = True
|
||||
self,
|
||||
context: OpenAILLMContext,
|
||||
*,
|
||||
user_kwargs: Mapping[str, Any] = {},
|
||||
assistant_kwargs: Mapping[str, Any] = {},
|
||||
) -> Any:
|
||||
pass
|
||||
|
||||
self._register_event_handler("on_completion_timeout")
|
||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||
await super().process_frame(frame, direction)
|
||||
|
||||
# TODO-CB: callback function type
|
||||
def register_function(self, function_name: Optional[str], callback, start_callback=None):
|
||||
if isinstance(frame, StartInterruptionFrame):
|
||||
await self._handle_interruptions(frame)
|
||||
|
||||
async def _handle_interruptions(self, frame: StartInterruptionFrame):
|
||||
for function_name, entry in self._functions.items():
|
||||
if entry.cancel_on_interruption:
|
||||
await self._cancel_function_call(function_name)
|
||||
|
||||
def register_function(
|
||||
self,
|
||||
function_name: Optional[str],
|
||||
callback: Any,
|
||||
start_callback=None,
|
||||
*,
|
||||
cancel_on_interruption: bool = False,
|
||||
):
|
||||
# Registering a function with the function_name set to None will run that callback
|
||||
# for all functions
|
||||
self._callbacks[function_name] = callback
|
||||
# QUESTION FOR CB: maybe this isn't needed anymore?
|
||||
self._functions[function_name] = FunctionEntry(
|
||||
function_name=function_name,
|
||||
callback=callback,
|
||||
cancel_on_interruption=cancel_on_interruption,
|
||||
)
|
||||
|
||||
# Start callbacks are now deprecated.
|
||||
if start_callback:
|
||||
import warnings
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("always")
|
||||
warnings.warn(
|
||||
"Parameter 'start_callback' is deprecated, just put your code on top of the actual function call instead.",
|
||||
DeprecationWarning,
|
||||
)
|
||||
|
||||
self._start_callbacks[function_name] = start_callback
|
||||
|
||||
def unregister_function(self, function_name: Optional[str]):
|
||||
del self._callbacks[function_name]
|
||||
del self._functions[function_name]
|
||||
if self._start_callbacks[function_name]:
|
||||
del self._start_callbacks[function_name]
|
||||
|
||||
def has_function(self, function_name: str):
|
||||
if None in self._callbacks.keys():
|
||||
if None in self._functions.keys():
|
||||
return True
|
||||
return function_name in self._callbacks.keys()
|
||||
return function_name in self._functions.keys()
|
||||
|
||||
async def call_function(
|
||||
self,
|
||||
@@ -188,25 +235,18 @@ class LLMService(AIService):
|
||||
function_name: str,
|
||||
arguments: str,
|
||||
run_llm: bool = True,
|
||||
) -> None:
|
||||
f = None
|
||||
if function_name in self._callbacks.keys():
|
||||
f = self._callbacks[function_name]
|
||||
elif None in self._callbacks.keys():
|
||||
f = self._callbacks[None]
|
||||
else:
|
||||
return None
|
||||
await self.call_start_function(context, function_name)
|
||||
await context.call_function(
|
||||
f,
|
||||
function_name=function_name,
|
||||
tool_call_id=tool_call_id,
|
||||
arguments=arguments,
|
||||
llm=self,
|
||||
run_llm=run_llm,
|
||||
):
|
||||
if not function_name in self._functions.keys() and not None in self._functions.keys():
|
||||
return
|
||||
|
||||
task = self.create_task(
|
||||
self._run_function_call(context, tool_call_id, function_name, arguments, run_llm)
|
||||
)
|
||||
|
||||
# QUESTION FOR CB: maybe this isn't needed anymore?
|
||||
self._function_call_tasks.add((task, tool_call_id, function_name))
|
||||
|
||||
task.add_done_callback(self._function_call_task_finished)
|
||||
|
||||
async def call_start_function(self, context: OpenAILLMContext, function_name: str):
|
||||
if function_name in self._start_callbacks.keys():
|
||||
await self._start_callbacks[function_name](function_name, self, context)
|
||||
@@ -218,6 +258,106 @@ class LLMService(AIService):
|
||||
UserImageRequestFrame(user_id=user_id, context=text_content), FrameDirection.UPSTREAM
|
||||
)
|
||||
|
||||
async def _run_function_call(
|
||||
self,
|
||||
context: OpenAILLMContext,
|
||||
tool_call_id: str,
|
||||
function_name: str,
|
||||
arguments: str,
|
||||
run_llm: bool = True,
|
||||
):
|
||||
if function_name in self._functions.keys():
|
||||
entry = self._functions[function_name]
|
||||
elif None in self._functions.keys():
|
||||
entry = self._functions[None]
|
||||
else:
|
||||
return
|
||||
|
||||
logger.info(f"Calling function {function_name} with arguments {arguments}")
|
||||
|
||||
# NOTE(aleix): This needs to be removed after we remove the deprecation.
|
||||
await self.call_start_function(context, function_name)
|
||||
|
||||
# Push a SystemFrame downstream. This frame will let our assistant context aggregator
|
||||
# know that we are in the middle of a function call. Some contexts/aggregators may
|
||||
# not need this. But some definitely do (Anthropic, for example).
|
||||
# Also push a SystemFrame upstream for use by other processors, like STTMuteFilter.
|
||||
progress_frame_downstream = FunctionCallInProgressFrame(
|
||||
function_name=function_name,
|
||||
tool_call_id=tool_call_id,
|
||||
arguments=arguments,
|
||||
cancel_on_interruption=entry.cancel_on_interruption,
|
||||
)
|
||||
progress_frame_upstream = FunctionCallInProgressFrame(
|
||||
function_name=function_name,
|
||||
tool_call_id=tool_call_id,
|
||||
arguments=arguments,
|
||||
cancel_on_interruption=entry.cancel_on_interruption,
|
||||
)
|
||||
|
||||
# Push frame both downstream and upstream
|
||||
await self.push_frame(progress_frame_downstream, FrameDirection.DOWNSTREAM)
|
||||
await self.push_frame(progress_frame_upstream, FrameDirection.UPSTREAM)
|
||||
|
||||
# Define a callback function that pushes a FunctionCallResultFrame upstream & downstream.
|
||||
async def function_call_result_callback(result, *, properties=None):
|
||||
result_frame_downstream = FunctionCallResultFrame(
|
||||
function_name=function_name,
|
||||
tool_call_id=tool_call_id,
|
||||
arguments=arguments,
|
||||
result=result,
|
||||
properties=properties,
|
||||
)
|
||||
result_frame_upstream = FunctionCallResultFrame(
|
||||
function_name=function_name,
|
||||
tool_call_id=tool_call_id,
|
||||
arguments=arguments,
|
||||
result=result,
|
||||
properties=properties,
|
||||
)
|
||||
|
||||
await self.push_frame(result_frame_downstream, FrameDirection.DOWNSTREAM)
|
||||
await self.push_frame(result_frame_upstream, FrameDirection.UPSTREAM)
|
||||
|
||||
await entry.callback(
|
||||
function_name, tool_call_id, arguments, self, context, function_call_result_callback
|
||||
)
|
||||
|
||||
async def _cancel_function_call(self, function_name: str):
|
||||
cancelled_tasks = set()
|
||||
for task, tool_call_id, name in self._function_call_tasks:
|
||||
if name == function_name:
|
||||
# We remove the callback because we are going to cancel the task
|
||||
# now, otherwise we will be removing it from the set while we
|
||||
# are iterating.
|
||||
task.remove_done_callback(self._function_call_task_finished)
|
||||
|
||||
logger.debug(f"{self} Cancelling function call [{name}:{tool_call_id}]...")
|
||||
|
||||
await self.cancel_task(task)
|
||||
|
||||
frame = FunctionCallCancelFrame(
|
||||
function_name=function_name, tool_call_id=tool_call_id
|
||||
)
|
||||
await self.push_frame(frame)
|
||||
|
||||
logger.debug(f"{self} Function call [{name}:{tool_call_id}] has been cancelled")
|
||||
|
||||
cancelled_tasks.add(task)
|
||||
|
||||
# Remove all cancelled tasks from our set.
|
||||
for task in cancelled_tasks:
|
||||
self._function_call_task_finished(task)
|
||||
|
||||
def _function_call_task_finished(self, task: asyncio.Task):
|
||||
tuple_to_remove = next((t for t in self._function_call_tasks if t[0] == task), None)
|
||||
if tuple_to_remove:
|
||||
self._function_call_tasks.discard(tuple_to_remove)
|
||||
# The task is finished so this should exit immediately. We need to
|
||||
# do this because otherwise the task manager would have a dangling
|
||||
# task if we don't remove it.
|
||||
asyncio.run_coroutine_threadsafe(self.wait_for_task(task), self.get_event_loop())
|
||||
|
||||
|
||||
class TTSService(AIService):
|
||||
def __init__(
|
||||
@@ -366,12 +506,14 @@ class TTSService(AIService):
|
||||
else:
|
||||
await self.push_frame(frame, direction)
|
||||
elif isinstance(frame, TTSSpeakFrame):
|
||||
# Store if we were processing text or not so we can set it back.
|
||||
processing_text = self._processing_text
|
||||
await self._push_tts_frames(frame.text)
|
||||
# We pause processing incoming frames because we are sending data to
|
||||
# the TTS. We pause to avoid audio overlapping.
|
||||
await self._maybe_pause_frame_processing()
|
||||
await self.flush_audio()
|
||||
self._processing_text = False
|
||||
self._processing_text = processing_text
|
||||
elif isinstance(frame, TTSUpdateSettingsFrame):
|
||||
await self._update_settings(frame.settings)
|
||||
elif isinstance(frame, BotStoppedSpeakingFrame):
|
||||
|
||||
Reference in New Issue
Block a user