LLMAssistantAggregator: cache function call requested images

src/pipecat/processors/aggregators/llm_response_universal.py

@@ -641,6 +641,7 @@ class LLMAssistantAggregator(LLMContextAggregator):
 
         self._started = 0
         self._function_calls_in_progress: Dict[str, Optional[FunctionCallInProgressFrame]] = {}
+        self._function_calls_image_results: Dict[str, UserImageRawFrame] = {}
         self._context_updated_tasks: Set[asyncio.Task] = set()
 
         self._assistant_turn_start_timestamp = ""
@@ -820,6 +821,15 @@ class LLMAssistantAggregator(LLMContextAggregator):
 
         run_llm = False
 
+        # Append any images that were generated by function calls.
+        if frame.tool_call_id in self._function_calls_image_results:
+            image_frame = self._function_calls_image_results[frame.tool_call_id]
+
+            del self._function_calls_image_results[frame.tool_call_id]
+
+            # If an image frame has been added to the context, let's run inference.
+            run_llm = await self._maybe_append_image_to_context(image_frame)
+
         # Run inference if the function call result requires it.
         if frame.result:
             if properties and properties.run_llm is not None:
@@ -856,31 +866,24 @@ class LLMAssistantAggregator(LLMContextAggregator):
             self._update_function_call_result(frame.function_name, frame.tool_call_id, "CANCELLED")
             del self._function_calls_in_progress[frame.tool_call_id]
 
-    def _update_function_call_result(self, function_name: str, tool_call_id: str, result: Any):
-        for message in self._context.get_messages():
-            if (
-                not isinstance(message, LLMSpecificMessage)
-                and message["role"] == "tool"
-                and message["tool_call_id"]
-                and message["tool_call_id"] == tool_call_id
-            ):
-                message["content"] = result
-
     async def _handle_user_image_frame(self, frame: UserImageRawFrame):
-        if not frame.append_to_context:
+        image_appended = False
-            return
 
-        logger.debug(f"{self} Appending UserImageRawFrame to LLM context (size: {frame.size})")
+        # Check if this image is a result of a function call if so, let's cache.
+        # TODO(aleix): The function call might have already been executed
+        # because FunctionCallResultFrame was just faster, in that case we just
+        # push the context frame now.
+        if (
+            frame.request
+            and frame.request.tool_call_id
+            and frame.request.tool_call_id in self._function_calls_in_progress
+        ):
+            self._function_calls_image_results[frame.request.tool_call_id] = frame
+        else:
+            image_appended = await self._maybe_append_image_to_context(frame)
 
-        await self._context.add_image_frame_message(
+        if image_appended:
-            format=frame.format,
+            await self.push_context_frame(FrameDirection.UPSTREAM)
-            size=frame.size,
-            image=frame.image,
-            text=frame.text,
-        )
-
-        await self._trigger_assistant_turn_stopped()
-        await self.push_context_frame(FrameDirection.UPSTREAM)
 
     async def _handle_assistant_image_frame(self, frame: AssistantImageRawFrame):
         logger.debug(f"{self} Appending AssistantImageRawFrame to LLM context (size: {frame.size})")
@@ -970,6 +973,31 @@ class LLMAssistantAggregator(LLMContextAggregator):
 
         await self._call_event_handler("on_assistant_thought", message)
 
+    async def _maybe_append_image_to_context(self, frame: UserImageRawFrame) -> bool:
+        if not frame.append_to_context:
+            return False
+
+        logger.debug(f"{self} Appending UserImageRawFrame to LLM context (size: {frame.size})")
+
+        await self._context.add_image_frame_message(
+            format=frame.format,
+            size=frame.size,
+            image=frame.image,
+            text=frame.text,
+        )
+
+        return True
+
+    def _update_function_call_result(self, function_name: str, tool_call_id: str, result: Any):
+        for message in self._context.get_messages():
+            if (
+                not isinstance(message, LLMSpecificMessage)
+                and message["role"] == "tool"
+                and message["tool_call_id"]
+                and message["tool_call_id"] == tool_call_id
+            ):
+                message["content"] = result
+
     def _context_updated_task_finished(self, task: asyncio.Task):
         self._context_updated_tasks.discard(task)