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

LLMAssistantContextAggregator: create a task to run on_context_updated

CHANGELOG.md

@@ -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.
 

src/pipecat/frames/frames.py

@@ -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
+    arguments: Any
-    cancel_on_interruption: bool
+    cancel_on_interruption: bool = False
 
 
 @dataclass

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

src/pipecat/services/anthropic.py

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

src/pipecat/services/google/google.py

@@ -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):
+            await self._update_function_call_result(
-                return
+                frame.function_name, frame.tool_call_id, frame.result
-
+            )
-            response = {"response": 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(
+        response = {"response": "CANCELLED"}
-            frame.function_name, frame.tool_call_id, "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,

src/pipecat/services/openai.py

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

tests/test_context_aggregators.py

@@ -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 (
+from pipecat.processors.aggregators.llm_response import LLMUserContextAggregator
-    LLMAssistantContextAggregator,
-    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