LLMAssistantAggregator: cache function call requested images
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user