function calling now run in tasks

This commit is contained in:
Aleix Conchillo Flaqué
2025-03-18 14:57:27 -07:00
parent fc06306efd
commit a98000fd1d
12 changed files with 537 additions and 567 deletions

View File

@@ -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):