LLMAssistantContextAggregator: create a task to run on_context_updated

CHANGELOG.md

@@ -13,6 +13,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- 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.
 

src/pipecat/processors/aggregators/llm_response.py

@@ -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
+        # Call the `on_context_updated` callback once the function call result
-        # result is added to the context
+        # 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):

src/pipecat/processors/frame_processor.py

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

src/pipecat/services/ai_services.py

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