Add streaming agent response callback to handoff

Author: TaoChenOSU

python/semantic_kernel/agents/orchestration/agent_actor_base.py

@@ -98,55 +98,55 @@ async def _call_streaming_agent_response_callback(
             else:
                 self._streaming_agent_response_callback(message_chunk, is_final)
 
-    async def _invoke_agent(self, additional_messages: DefaultTypeAlias | None = None) -> ChatMessageContent:
+    async def _invoke_agent(self, additional_messages: DefaultTypeAlias | None = None, **kwargs) -> ChatMessageContent:
         """Invoke the agent with the current chat history or thread and optionally additional messages.
 
         Args:
             additional_messages (DefaultTypeAlias | None): Additional messages to be sent to the agent.
+            **kwargs: Additional keyword arguments to be passed to the agent's invoke method:
+                - kernel: The kernel to use for the agent invocation.
 
         Returns:
             DefaultTypeAlias: The response from the agent.
         """
         streaming_message_buffer: list[StreamingChatMessageContent] = []
+        messages = self._create_messages(additional_messages)
 
-        if self._agent_thread is None:
-            messages = self._chat_history.messages[:]
-            if additional_messages:
-                messages.extend(additional_messages if isinstance(additional_messages, list) else [additional_messages])
-
-            async for response_item in self._agent.invoke_stream(messages=messages):
-                # Buffer message chunks and stream them with correct is_final flag.
-                streaming_message_buffer.append(response_item.message)
-                if len(streaming_message_buffer) > 1:
-                    await self._call_streaming_agent_response_callback(streaming_message_buffer[-2], is_final=False)
-                if self._agent_thread is None:
-                    self._agent_thread = response_item.thread
-            if streaming_message_buffer:
-                await self._call_streaming_agent_response_callback(streaming_message_buffer[-1], is_final=True)
-        else:
-            messages = (
-                []
-                if additional_messages is None
-                else additional_messages
-                if isinstance(additional_messages, list)
-                else [additional_messages]
-            )
-
-            async for response_item in self._agent.invoke_stream(
-                messages=messages,
-                thread=self._agent_thread,
-            ):
-                # Buffer message chunks and stream them with correct is_final flag.
-                streaming_message_buffer.append(response_item.message)
-                if len(streaming_message_buffer) > 1:
-                    await self._call_streaming_agent_response_callback(streaming_message_buffer[-2], is_final=False)
-            if streaming_message_buffer:
-                await self._call_streaming_agent_response_callback(streaming_message_buffer[-1], is_final=True)
+        async for response_item in self._agent.invoke_stream(messages=messages, thread=self._agent_thread, **kwargs):
+            # Buffer message chunks and stream them with correct is_final flag.
+            streaming_message_buffer.append(response_item.message)
+            if len(streaming_message_buffer) > 1:
+                await self._call_streaming_agent_response_callback(streaming_message_buffer[-2], is_final=False)
+            if self._agent_thread is None:
+                self._agent_thread = response_item.thread
+
+        if streaming_message_buffer:
+            # Call the callback for the last message chunk with is_final=True.
+            await self._call_streaming_agent_response_callback(streaming_message_buffer[-1], is_final=True)
 
         if not streaming_message_buffer:
             raise RuntimeError(f'Agent "{self._agent.name}" did not return any response.')
 
+        # Build the full response from the streaming messages
         full_response = sum(streaming_message_buffer[1:], streaming_message_buffer[0])
         await self._call_agent_response_callback(full_response)
 
         return full_response
+
+    def _create_messages(self, additional_messages: DefaultTypeAlias | None = None) -> list[ChatMessageContent]:
+        """Create a list of messages to be sent to the agent along with a potential thread.
+
+        Args:
+            additional_messages (DefaultTypeAlias | None): Additional messages to be sent to the agent.
+
+        Returns:
+            list[ChatMessageContent]: A list of messages to be sent to the agent.
+        """
+        base_messages = self._chat_history.messages[:] if self._agent_thread is None else []
+
+        if additional_messages is None:
+            return base_messages
+
+        if isinstance(additional_messages, list):
+            return base_messages + additional_messages
+        return [*base_messages, additional_messages]

python/semantic_kernel/agents/orchestration/handoffs.py

@@ -18,6 +18,7 @@
 from semantic_kernel.agents.runtime.core.topic import TopicId
 from semantic_kernel.agents.runtime.in_process.type_subscription import TypeSubscription
 from semantic_kernel.contents.chat_message_content import ChatMessageContent
+from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
 from semantic_kernel.contents.utils.author_role import AuthorRole
 from semantic_kernel.filters.auto_function_invocation.auto_function_invocation_context import (
     AutoFunctionInvocationContext,
@@ -162,6 +163,8 @@ def __init__(
         handoff_connections: AgentHandoffs,
         result_callback: Callable[[DefaultTypeAlias], Awaitable[None]] | None = None,
         agent_response_callback: Callable[[DefaultTypeAlias], Awaitable[None] | None] | None = None,
+        streaming_agent_response_callback: Callable[[StreamingChatMessageContent, bool], Awaitable[None] | None]
+        | None = None,
         human_response_function: Callable[[], Awaitable[ChatMessageContent] | ChatMessageContent] | None = None,
     ) -> None:
         """Initialize the handoff agent actor."""
@@ -179,6 +182,7 @@ def __init__(
             agent=agent,
             internal_topic_type=internal_topic_type,
             agent_response_callback=agent_response_callback,
+            streaming_agent_response_callback=streaming_agent_response_callback,
         )
 
     def _add_handoff_functions(self) -> None:
@@ -275,62 +279,44 @@ async def _handle_request_message(self, message: HandoffRequestMessage, cts: Mes
             return
         logger.debug(f"{self.id}: Received handoff request message.")
 
-        if self._agent_thread is None:
-            self._chat_history.add_message(
-                ChatMessageContent(
-                    role=AuthorRole.USER,
-                    content=f"Transferred to {self._agent.name}, adopt the persona immediately.",
-                )
-            )
-            response_item = await self._agent.get_response(
-                messages=self._chat_history.messages,  # type: ignore[arg-type]
-                kernel=self._kernel,
-            )
-        else:
-            response_item = await self._agent.get_response(
-                messages=ChatMessageContent(
-                    role=AuthorRole.USER,
-                    content=f"Transferred to {self._agent.name}, adopt the persona immediately.",
-                ),
-                thread=self._agent_thread,
-                kernel=self._kernel,
-            )
-
-        if self._agent_thread is None:
-            self._agent_thread = response_item.thread
+        persona_adoption_message = ChatMessageContent(
+            role=AuthorRole.USER,
+            content=f"Transferred to {self._agent.name}, adopt the persona immediately.",
+        )
+        response = await self._invoke_agent_with_potentially_no_response(
+            additional_messages=persona_adoption_message,
+            kernel=self._kernel,
+        )
 
         while not self._task_completed:
-            if response_item.message.role == AuthorRole.ASSISTANT:
-                # The response can potentially be a TOOL message from the Handoff plugin
-                # since we have added a filter which will terminate the conversation when
-                # a function from the handoff plugin is called. And we don't want to publish
-                # that message. So we only publish if the response is an ASSISTANT message.
-                logger.debug(f"{self.id} responded with: {response_item.message.content}")
-                await self._call_agent_response_callback(response_item.message)
-
-                await self.publish_message(
-                    HandoffResponseMessage(body=response_item.message),
-                    TopicId(self._internal_topic_type, self.id.key),
-                    cancellation_token=cts.cancellation_token,
-                )
-
             if self._handoff_agent_name:
                 await self.publish_message(
                     HandoffRequestMessage(agent_name=self._handoff_agent_name),
                     TopicId(self._internal_topic_type, self.id.key),
                 )
                 self._handoff_agent_name = None
                 break
+
+            if response is None:
+                raise RuntimeError(
+                    f'Agent "{self._agent.name}" did not return any response nor did not set a handoff agent name.'
+                )
+
+            await self.publish_message(
+                HandoffResponseMessage(body=response),
+                TopicId(self._internal_topic_type, self.id.key),
+                cancellation_token=cts.cancellation_token,
+            )
+
             if self._human_response_function:
                 human_response = await self._call_human_response_function()
                 await self.publish_message(
                     HandoffResponseMessage(body=human_response),
                     TopicId(self._internal_topic_type, self.id.key),
                     cancellation_token=cts.cancellation_token,
                 )
-                response_item = await self._agent.get_response(
-                    messages=human_response,
-                    thread=self._agent_thread,
+                response = await self._invoke_agent_with_potentially_no_response(
+                    additional_messages=human_response,
                     kernel=self._kernel,
                 )
             else:
@@ -346,6 +332,43 @@ async def _call_human_response_function(self) -> ChatMessageContent:
             return await self._human_response_function()
         return self._human_response_function()  # type: ignore[return-value]
 
+    async def _invoke_agent_with_potentially_no_response(
+        self, additional_messages: DefaultTypeAlias | None = None, **kwargs
+    ) -> ChatMessageContent | None:
+        """Invoke the agent with the current chat history or thread and optionally additional messages.
+
+        This method differs from `_invoke_agent` in that it handles the case where no response is returned
+        from the agent gracefully, returning `None` instead of raising an error.
+
+        The reason for this is that agents in a handoff group chat may not always produce a response when
+        a handoff function is invoked, where the `_handoff_function_filter` will terminate the auto function
+        invocation loop before a response is produced. In such cases, this method will return `None`
+        instead of raising an error.
+        """
+        streaming_message_buffer: list[StreamingChatMessageContent] = []
+        messages = self._create_messages(additional_messages)
+
+        async for response_item in self._agent.invoke_stream(messages=messages, thread=self._agent_thread, **kwargs):
+            # Buffer message chunks and stream them with correct is_final flag.
+            streaming_message_buffer.append(response_item.message)
+            if len(streaming_message_buffer) > 1:
+                await self._call_streaming_agent_response_callback(streaming_message_buffer[-2], is_final=False)
+            if self._agent_thread is None:
+                self._agent_thread = response_item.thread
+
+        if streaming_message_buffer:
+            # Call the callback for the last message chunk with is_final=True.
+            await self._call_streaming_agent_response_callback(streaming_message_buffer[-1], is_final=True)
+
+        if not streaming_message_buffer:
+            return None
+
+        # Build the full response from the streaming messages
+        full_response = sum(streaming_message_buffer[1:], streaming_message_buffer[0])
+        await self._call_agent_response_callback(full_response)
+
+        return full_response
+
 
 # endregion HandoffAgentActor
 
@@ -365,6 +388,8 @@ def __init__(
         input_transform: Callable[[TIn], Awaitable[DefaultTypeAlias] | DefaultTypeAlias] | None = None,
         output_transform: Callable[[DefaultTypeAlias], Awaitable[TOut] | TOut] | None = None,
         agent_response_callback: Callable[[DefaultTypeAlias], Awaitable[None] | None] | None = None,
+        streaming_agent_response_callback: Callable[[StreamingChatMessageContent, bool], Awaitable[None] | None]
+        | None = None,
         human_response_function: Callable[[], Awaitable[ChatMessageContent] | ChatMessageContent] | None = None,
     ) -> None:
         """Initialize the handoff orchestration.
@@ -377,8 +402,10 @@ def __init__(
             description (str | None): The description of the orchestration.
             input_transform (Callable | None): A function that transforms the external input message.
             output_transform (Callable | None): A function that transforms the internal output message.
-            agent_response_callback (Callable | None): A function that is called when a response is produced
+            agent_response_callback (Callable | None): A function that is called when a full response is produced
                 by the agents.
+            streaming_agent_response_callback (Callable | None): A function that is called when a streaming response
+                is produced by the agents.
             human_response_function (Callable | None): A function that is called when a human response is
                 needed.
         """
@@ -392,6 +419,7 @@ def __init__(
             input_transform=input_transform,
             output_transform=output_transform,
             agent_response_callback=agent_response_callback,
+            streaming_agent_response_callback=streaming_agent_response_callback,
         )
 
         self._validate_handoffs()
@@ -461,6 +489,7 @@ async def _register_helper(agent: Agent) -> None:
                     handoff_connections,
                     result_callback=result_callback,
                     agent_response_callback=self._agent_response_callback,
+                    streaming_agent_response_callback=self._streaming_agent_response_callback,
                     human_response_function=self._human_response_function,
                 ),
             )