Add streaming agent response callback to magentic

Author: TaoChenOSU

python/semantic_kernel/agents/orchestration/magentic.py

@@ -31,6 +31,7 @@
 from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
 from semantic_kernel.contents.chat_history import ChatHistory
 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.functions.kernel_arguments import KernelArguments
 from semantic_kernel.kernel import Kernel
@@ -701,29 +702,17 @@ async def _handle_request_message(self, message: MagenticRequestMessage, ctx: Me
             return
 
         logger.debug(f"{self.id}: Received request message.")
-        if self._agent_thread is None:
-            # Add a user message to steer the agent to respond more closely to the instructions.
-            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]
-            self._agent_thread = response_item.thread
-        else:
-            # Add a user message to steer the agent to respond more closely to the instructions.
-            new_message = ChatMessageContent(
-                role=AuthorRole.USER,
-                content=f"Transferred to {self._agent.name}, adopt the persona immediately.",
-            )
-            response_item = await self._agent.get_response(messages=new_message, thread=self._agent_thread)
 
-        logger.debug(f"{self.id} responded with {response_item.message.content}.")
-        await self._call_agent_response_callback(response_item.message)
+        persona_adoption_message = ChatMessageContent(
+            role=AuthorRole.USER,
+            content=f"Transferred to {self._agent.name}, adopt the persona immediately.",
+        )
+        response = await self._invoke_agent(additional_messages=persona_adoption_message)
+
+        logger.debug(f"{self.id} responded with {response}.")
 
         await self.publish_message(
-            MagenticResponseMessage(body=response_item.message),
+            MagenticResponseMessage(body=response),
             TopicId(self._internal_topic_type, self.id.key),
             cancellation_token=ctx.cancellation_token,
         )
@@ -756,6 +745,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,
     ) -> None:
         """Initialize the Magentic One orchestration.
 
@@ -768,6 +759,8 @@ def __init__(
             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
                 by the agents.
+            streaming_agent_response_callback (Callable | None): A function that is called when a streaming response
+                is produced by the agents.
         """
         self._manager = manager
 
@@ -782,6 +775,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,
         )
 
     @override
@@ -826,6 +820,7 @@ async def _register_members(self, runtime: CoreRuntime, internal_topic_type: str
                     agent,
                     internal_topic_type,
                     self._agent_response_callback,
+                    self._streaming_agent_response_callback,
                 ),
             )
             for agent in self._members