From 19b464ba23691c939ffd06f8f7704baef753e2fb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= <aleix@daily.co>
Date: Tue, 25 Mar 2025 14:41:33 -0700
Subject: [PATCH] tests: add assistant aggregator function call frame handling

---
 tests/test_context_aggregators.py | 97 ++++++++++++++++++++++++++-----
 1 file changed, 84 insertions(+), 13 deletions(-)

diff --git a/tests/test_context_aggregators.py b/tests/test_context_aggregators.py
index 185725632..baee3496f 100644
--- a/tests/test_context_aggregators.py
+++ b/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 (
-    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