LLMAssistantAggregator: cache function call requested images

This commit is contained in:
Aleix Conchillo Flaqué
2026-01-13 11:56:43 -08:00
parent d3c57e2da0
commit e268c73c41

View File

@@ -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:
return
image_appended = False
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(
format=frame.format,
size=frame.size,
image=frame.image,
text=frame.text,
)
await self._trigger_assistant_turn_stopped()
await self.push_context_frame(FrameDirection.UPSTREAM)
if image_appended:
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)