Merge pull request #1442 from pipecat-ai/aleix/on-context-updated-as-task

LLMAssistantContextAggregator: create a task to run on_context_updated
This commit is contained in:
Aleix Conchillo Flaqué
2025-03-25 15:39:36 -07:00
committed by GitHub
9 changed files with 128 additions and 39 deletions

View File

@@ -17,6 +17,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Fixed
- Fixed a `GoogleAssistantContextAggregator` issue where function calls
placeholders where not being updated when then function call result was
different from a string.
- Fixed an issue that would cause `LLMAssistantContextAggregator` to block
processing more frames while processing a function call result.
- Fixed an issue where the `RTVIObserver` would report two bot started and
stopped speaking events for each bot turn.

View File

@@ -384,7 +384,7 @@ class FunctionCallResultFrame(DataFrame):
function_name: str
tool_call_id: str
arguments: str
arguments: Any
result: Any
properties: Optional[FunctionCallResultProperties] = None
@@ -633,8 +633,8 @@ class FunctionCallInProgressFrame(SystemFrame):
function_name: str
tool_call_id: str
arguments: str
cancel_on_interruption: bool
arguments: Any
cancel_on_interruption: bool = False
@dataclass

View File

@@ -6,7 +6,7 @@
import asyncio
from abc import abstractmethod
from typing import Dict, List
from typing import Dict, List, Set
from loguru import logger
@@ -380,6 +380,7 @@ class LLMAssistantContextAggregator(LLMContextResponseAggregator):
self._started = 0
self._function_calls_in_progress: Dict[str, FunctionCallInProgressFrame] = {}
self._context_updated_tasks: Set[asyncio.Task] = set()
async def handle_aggregation(self, aggregation: str):
self._context.add_message({"role": "assistant", "content": aggregation})
@@ -486,10 +487,14 @@ class LLMAssistantContextAggregator(LLMContextResponseAggregator):
if run_llm:
await self.push_context_frame(FrameDirection.UPSTREAM)
# Emit the on_context_updated callback once the function call
# result is added to the context
# Call the `on_context_updated` callback once the function call result
# is added to the context. Also, run this in a separate task to make
# sure we don't block the pipeline.
if properties and properties.on_context_updated:
await properties.on_context_updated()
task_name = f"{frame.function_name}:{frame.tool_call_id}:on_context_updated"
task = self.create_task(properties.on_context_updated(), task_name)
self._context_updated_tasks.add(task)
task.add_done_callback(self._context_updated_task_finished)
async def _handle_function_call_cancel(self, frame: FunctionCallCancelFrame):
logger.debug(
@@ -535,6 +540,13 @@ class LLMAssistantContextAggregator(LLMContextResponseAggregator):
else:
self._aggregation += frame.text
def _context_updated_task_finished(self, task: asyncio.Task):
self._context_updated_tasks.discard(task)
# The task is finished so this should exit immediately. We need to do
# this because otherwise the task manager would report a dangling task
# if we don't remove it.
asyncio.run_coroutine_threadsafe(self.wait_for_task(task), self.get_event_loop())
class LLMUserResponseAggregator(LLMUserContextAggregator):
def __init__(self, messages: List[dict] = [], **kwargs):

View File

@@ -147,10 +147,13 @@ class FrameProcessor(BaseObject):
await self.stop_ttfb_metrics()
await self.stop_processing_metrics()
def create_task(self, coroutine: Coroutine) -> asyncio.Task:
def create_task(self, coroutine: Coroutine, name: Optional[str] = None) -> asyncio.Task:
if not self._task_manager:
raise Exception(f"{self} TaskManager is still not initialized.")
name = f"{self}::{coroutine.cr_code.co_name}"
if name:
name = f"{self}::{name}"
else:
name = f"{self}::{coroutine.cr_code.co_name}"
return self._task_manager.create_task(coroutine, name)
async def cancel_task(self, task: asyncio.Task, timeout: Optional[float] = None):

View File

@@ -369,7 +369,7 @@ class LLMService(AIService):
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
# do this because otherwise the task manager would report a dangling
# task if we don't remove it.
asyncio.run_coroutine_threadsafe(self.wait_for_task(task), self.get_event_loop())

View File

@@ -725,7 +725,7 @@ class AnthropicAssistantContextAggregator(LLMAssistantContextAggregator):
)
async def _update_function_call_result(
self, function_name: str, tool_call_id: str, result: str
self, function_name: str, tool_call_id: str, result: Any
):
for message in self._context.messages:
if message["role"] == "user":

View File

@@ -601,23 +601,18 @@ class GoogleAssistantContextAggregator(OpenAIAssistantContextAggregator):
async def handle_function_call_result(self, frame: FunctionCallResultFrame):
if frame.result:
if not isinstance(frame.result, str):
return
response = {"response": frame.result}
await self._update_function_call_result(
frame.function_name, frame.tool_call_id, frame.result
)
else:
response = {"response": "COMPLETED"}
await self._update_function_call_result(
frame.function_name, frame.tool_call_id, response
)
else:
await self._update_function_call_result(
frame.function_name, frame.tool_call_id, "COMPLETED"
)
async def handle_function_call_cancel(self, frame: FunctionCallCancelFrame):
await self._update_function_call_result(
frame.function_name, frame.tool_call_id, "CANCELLED"
)
response = {"response": "CANCELLED"}
await self._update_function_call_result(frame.function_name, frame.tool_call_id, response)
async def _update_function_call_result(
self, function_name: str, tool_call_id: str, result: Any
@@ -626,11 +621,12 @@ class GoogleAssistantContextAggregator(OpenAIAssistantContextAggregator):
if message.role == "user":
for part in message.parts:
if part.function_response and part.function_response.id == tool_call_id:
part.function_response.response = {"response": result}
part.function_response.response = result
async def handle_user_image_frame(self, frame: UserImageRawFrame):
response = {"response": "COMPLETED"}
await self._update_function_call_result(
frame.request.function_name, frame.request.tool_call_id, "COMPLETED"
frame.request.function_name, frame.request.tool_call_id, response
)
self._context.add_image_frame_message(
format=frame.format,

View File

@@ -613,7 +613,7 @@ class OpenAIAssistantContextAggregator(LLMAssistantContextAggregator):
)
async def _update_function_call_result(
self, function_name: str, tool_call_id: str, result: str
self, function_name: str, tool_call_id: str, result: Any
):
for message in self._context.messages:
if (

View File

@@ -4,13 +4,18 @@
# SPDX-License-Identifier: BSD 2-Clause License
#
import json
import unittest
from typing import Any
import google.ai.generativelanguage as glm
from pipecat.frames.frames import (
EmulateUserStartedSpeakingFrame,
EmulateUserStoppedSpeakingFrame,
FunctionCallInProgressFrame,
FunctionCallResultFrame,
FunctionCallResultProperties,
InterimTranscriptionFrame,
LLMFullResponseEndFrame,
LLMFullResponseStartFrame,
@@ -21,10 +26,7 @@ from pipecat.frames.frames import (
UserStartedSpeakingFrame,
UserStoppedSpeakingFrame,
)
from pipecat.processors.aggregators.llm_response import (
LLMAssistantContextAggregator,
LLMUserContextAggregator,
)
from pipecat.processors.aggregators.llm_response import LLMUserContextAggregator
from pipecat.processors.aggregators.openai_llm_context import (
OpenAILLMContext,
OpenAILLMContextFrame,
@@ -423,6 +425,9 @@ class BaseTestAssistantContextAggreagator:
):
assert context.messages[index]["content"] == content
def check_function_call_result(self, context: OpenAILLMContext, index: int, content: str):
assert json.loads(context.messages[index]["content"]) == content
async def test_empty(self):
assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
assert self.AGGREGATOR_CLASS is not None, "AGGREGATOR_CLASS must be set in a subclass"
@@ -556,9 +561,76 @@ class BaseTestAssistantContextAggreagator:
self.check_message_multi_content(context, 0, 0, "Hello Pipecat.")
self.check_message_multi_content(context, 0, 1, "How are you?")
async def test_function_call(self):
assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
assert self.AGGREGATOR_CLASS is not None, "AGGREGATOR_CLASS must be set in a subclass"
context = self.CONTEXT_CLASS()
aggregator = self.AGGREGATOR_CLASS(context)
frames_to_send = [
FunctionCallInProgressFrame(
function_name="get_weather",
tool_call_id="1",
arguments={"location": "Los Angeles"},
cancel_on_interruption=False,
),
SleepFrame(),
FunctionCallResultFrame(
function_name="get_weather",
tool_call_id="1",
arguments={"location": "Los Angeles"},
result={"conditions": "Sunny"},
),
]
expected_down_frames = []
await run_test(
aggregator,
frames_to_send=frames_to_send,
expected_down_frames=expected_down_frames,
)
self.check_function_call_result(context, -1, {"conditions": "Sunny"})
async def test_function_call_on_context_updated(self):
assert self.CONTEXT_CLASS is not None, "CONTEXT_CLASS must be set in a subclass"
assert self.AGGREGATOR_CLASS is not None, "AGGREGATOR_CLASS must be set in a subclass"
context_updated = False
async def on_context_updated():
nonlocal context_updated
context_updated = True
context = self.CONTEXT_CLASS()
aggregator = self.AGGREGATOR_CLASS(context)
frames_to_send = [
FunctionCallInProgressFrame(
function_name="get_weather",
tool_call_id="1",
arguments={"location": "Los Angeles"},
cancel_on_interruption=False,
),
SleepFrame(),
FunctionCallResultFrame(
function_name="get_weather",
tool_call_id="1",
arguments={"location": "Los Angeles"},
result={"conditions": "Sunny"},
properties=FunctionCallResultProperties(on_context_updated=on_context_updated),
),
SleepFrame(),
]
expected_down_frames = []
await run_test(
aggregator,
frames_to_send=frames_to_send,
expected_down_frames=expected_down_frames,
)
self.check_function_call_result(context, -1, {"conditions": "Sunny"})
assert context_updated
#
# LLMUserContextAggregator, LLMAssistantContextAggregator
# LLMUserContextAggregator
#
@@ -567,14 +639,6 @@ class TestLLMUserContextAggregator(BaseTestUserContextAggregator, unittest.Isola
AGGREGATOR_CLASS = LLMUserContextAggregator
class TestLLMAssistantContextAggregator(
BaseTestAssistantContextAggreagator, unittest.IsolatedAsyncioTestCase
):
CONTEXT_CLASS = OpenAILLMContext
AGGREGATOR_CLASS = LLMAssistantContextAggregator
EXPECTED_CONTEXT_FRAMES = [OpenAILLMContextFrame, OpenAILLMContextAssistantTimestampFrame]
#
# OpenAI
#
@@ -626,6 +690,9 @@ class TestAnthropicAssistantContextAggregator(
messages = context.messages[content_index]
assert messages["content"][index]["text"] == content
def check_function_call_result(self, context: OpenAILLMContext, index: int, content: Any):
assert context.messages[index]["content"][0]["content"] == json.dumps(content)
#
# Google
@@ -665,3 +732,7 @@ class TestGoogleAssistantContextAggregator(
):
obj = glm.Content.to_dict(context.messages[index])
assert obj["parts"][0]["text"] == content
def check_function_call_result(self, context: OpenAILLMContext, index: int, content: Any):
obj = glm.Content.to_dict(context.messages[index])
assert obj["parts"][0]["function_response"]["response"] == content