Python: Improve orchestration exception handling (#12716)

### Motivation and Context

<!-- Thank you for your contribution to the semantic-kernel repo!
Please help reviewers and future users, providing the following
information:
  1. Why is this change required?
  2. What problem does it solve?
  3. What scenario does it contribute to?
  4. If it fixes an open issue, please link to the issue here.
-->
Currently, exceptions that occur inside the orchestration actors are not
properly surfaced to the caller as it happens.

Fixes: #12719

### Description

<!-- Describe your changes, the overall approach, the underlying design.
These notes will help understanding how your code works. Thanks! -->
This PR addresses the issue by introducing a new exception_callback that
is hidden from the user but will raise exceptions that occurs inside the
orchestration actors to the caller.

### Contribution Checklist

<!-- Before submitting this PR, please make sure: -->

- [x] The code builds clean without any errors or warnings
- [x] The PR follows the [SK Contribution
Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md)
and the [pre-submission formatting
script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts)
raises no violations
- [x] All unit tests pass, and I have added new tests where possible
- [x] I didn't break anyone 😄

Author: TaoChenOSU

python/semantic_kernel/agents/orchestration/agent_actor_base.py

@@ -2,8 +2,10 @@
 
 
 import inspect
+import logging
 import sys
 from collections.abc import Awaitable, Callable
+from functools import wraps
 from typing import Any
 
 from semantic_kernel.agents.agent import Agent, AgentThread
@@ -18,11 +20,27 @@
 else:
     from typing_extensions import override  # pragma: no cover
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 @experimental
 class ActorBase(RoutedAgent):
     """A base class for actors running in the AgentRuntime."""
 
+    def __init__(
+        self,
+        description: str,
+        exception_callback: Callable[[BaseException], None],
+    ):
+        """Initialize the actor with a description and an exception callback.
+
+        Args:
+            description (str): A description of the actor.
+            exception_callback (Callable[[BaseException], None]): A callback function to handle exceptions.
+        """
+        super().__init__(description=description)
+        self._exception_callback = exception_callback
+
     @override
     async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any | None:
         """Handle a message.
@@ -34,6 +52,46 @@ async def on_message_impl(self, message: Any, ctx: MessageContext) -> Any | None
 
         return await super().on_message_impl(message, ctx)
 
+    @staticmethod
+    def exception_handler(func: Callable[..., Any]) -> Callable[..., Any]:
+        """Decorator that wraps a function in a try-catch block and calls the exception callback on errors.
+
+        This decorator can be used on both synchronous and asynchronous functions. When an exception
+        occurs during function execution, it will call the exception_callback with the exception
+        and then re-raise the exception.
+
+        Args:
+            func: The function to be wrapped.
+
+        Returns:
+            The wrapped function.
+        """
+        log_message_template = "Exception occurred in agent {agent_id}:\n{exception}"
+
+        if inspect.iscoroutinefunction(func):
+
+            @wraps(func)
+            async def async_wrapper(self, *args, **kwargs):
+                try:
+                    return await func(self, *args, **kwargs)
+                except BaseException as e:
+                    self._exception_callback(e)
+                    logger.error(log_message_template.format(agent_id=self.id, exception=e))
+                    raise
+
+            return async_wrapper
+
+        @wraps(func)
+        def sync_wrapper(self, *args, **kwargs):
+            try:
+                return func(self, *args, **kwargs)
+            except BaseException as e:
+                self._exception_callback(e)
+                logger.error(log_message_template.format(agent_id=self.id, exception=e))
+                raise
+
+        return sync_wrapper
+
 
 @experimental
 class AgentActorBase(ActorBase):
@@ -43,6 +101,7 @@ def __init__(
         self,
         agent: Agent,
         internal_topic_type: str,
+        exception_callback: Callable[[BaseException], None],
         agent_response_callback: Callable[[DefaultTypeAlias], Awaitable[None] | None] | None = None,
         streaming_agent_response_callback: Callable[[StreamingChatMessageContent, bool], Awaitable[None] | None]
         | None = None,
@@ -52,6 +111,7 @@ def __init__(
         Args:
             agent (Agent): An agent to be run in the container.
             internal_topic_type (str): The topic type of the internal topic.
+            exception_callback (Callable): A function that is called when an exception occurs.
             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
@@ -66,7 +126,7 @@ def __init__(
         # Chat history to temporarily store messages before each invoke.
         self._message_cache: ChatHistory = ChatHistory()
 
-        ActorBase.__init__(self, description=agent.description or "Semantic Kernel Agent")
+        super().__init__(agent.description or "Semantic Kernel Actor", exception_callback)
 
     async def _call_agent_response_callback(self, message: DefaultTypeAlias) -> None:
         """Call the agent_response_callback function if it is set.
@@ -97,6 +157,7 @@ async def _call_streaming_agent_response_callback(
             else:
                 self._streaming_agent_response_callback(message_chunk, is_final)
 
+    @ActorBase.exception_handler
     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.
 

python/semantic_kernel/agents/orchestration/concurrent.py

@@ -56,6 +56,7 @@ def __init__(
         agent: Agent,
         internal_topic_type: str,
         collection_agent_type: str,
+        exception_callback: Callable[[BaseException], None],
         agent_response_callback: Callable[[DefaultTypeAlias], Awaitable[None] | None] | None = None,
         streaming_agent_response_callback: Callable[[StreamingChatMessageContent, bool], Awaitable[None] | None]
         | None = None,
@@ -66,13 +67,15 @@ def __init__(
             agent: The agent to be executed.
             internal_topic_type: The internal topic type for the actor.
             collection_agent_type: The collection agent type for the actor.
+            exception_callback: A callback function to handle exceptions.
             agent_response_callback: A callback function to handle the full response from the agent.
             streaming_agent_response_callback: A callback function to handle streaming responses from the agent.
         """
         self._collection_agent_type = collection_agent_type
         super().__init__(
             agent=agent,
             internal_topic_type=internal_topic_type,
+            exception_callback=exception_callback,
             agent_response_callback=agent_response_callback,
             streaming_agent_response_callback=streaming_agent_response_callback,
         )
@@ -102,6 +105,7 @@ def __init__(
         self,
         description: str,
         expected_answer_count: int,
+        exception_callback: Callable[[BaseException], None],
         result_callback: Callable[[DefaultTypeAlias], Awaitable[None]] | None = None,
     ) -> None:
         """Initialize the collection agent container."""
@@ -110,7 +114,7 @@ def __init__(
         self._results: list[ChatMessageContent] = []
         self._lock = asyncio.Lock()
 
-        super().__init__(description=description)
+        super().__init__(description, exception_callback)
 
     @message_handler
     async def _handle_message(self, message: ConcurrentResponseMessage, _: MessageContext) -> None:
@@ -147,17 +151,20 @@ async def _prepare(
         self,
         runtime: CoreRuntime,
         internal_topic_type: str,
-        result_callback: Callable[[DefaultTypeAlias], Awaitable[None]] | None = None,
+        exception_callback: Callable[[BaseException], None],
+        result_callback: Callable[[DefaultTypeAlias], Awaitable[None]],
     ) -> None:
         """Register the actors and orchestrations with the runtime and add the required subscriptions."""
         await asyncio.gather(*[
             self._register_members(
                 runtime,
                 internal_topic_type,
+                exception_callback,
             ),
             self._register_collection_actor(
                 runtime,
                 internal_topic_type,
+                exception_callback,
                 result_callback=result_callback,
             ),
             self._add_subscriptions(
@@ -170,6 +177,7 @@ async def _register_members(
         self,
         runtime: CoreRuntime,
         internal_topic_type: str,
+        exception_callback: Callable[[BaseException], None],
     ) -> None:
         """Register the members."""
 
@@ -180,7 +188,8 @@ async def _internal_helper(agent: Agent) -> None:
                 lambda agent=agent: ConcurrentAgentActor(  # type: ignore[misc]
                     agent,
                     internal_topic_type,
-                    collection_agent_type=self._get_collection_actor_type(internal_topic_type),
+                    self._get_collection_actor_type(internal_topic_type),
+                    exception_callback,
                     agent_response_callback=self._agent_response_callback,
                     streaming_agent_response_callback=self._streaming_agent_response_callback,
                 ),
@@ -192,6 +201,7 @@ async def _register_collection_actor(
         self,
         runtime: CoreRuntime,
         internal_topic_type: str,
+        exception_callback: Callable[[BaseException], None],
         result_callback: Callable[[DefaultTypeAlias], Awaitable[None]] | None = None,
     ) -> None:
         await CollectionActor.register(
@@ -200,6 +210,7 @@ async def _register_collection_actor(
             lambda: CollectionActor(
                 description="An internal agent that is responsible for collection results",
                 expected_answer_count=len(self._members),
+                exception_callback=exception_callback,
                 result_callback=result_callback,
             ),
         )

python/semantic_kernel/agents/orchestration/group_chat.py

@@ -257,6 +257,7 @@ def __init__(
         manager: GroupChatManager,
         internal_topic_type: str,
         participant_descriptions: dict[str, str],
+        exception_callback: Callable[[BaseException], None],
         result_callback: Callable[[DefaultTypeAlias], Awaitable[None]] | None = None,
     ):
         """Initialize the group chat manager actor.
@@ -265,8 +266,7 @@ def __init__(
             manager (GroupChatManager): The group chat manager that manages the flow of the group chat.
             internal_topic_type (str): The topic type of the internal topic.
             participant_descriptions (dict[str, str]): The descriptions of the participants in the group chat.
-            agent_response_callback (Callable | None): A function that is called when a response is produced
-                by the agents.
+            exception_callback (Callable[[BaseException], None]): A function that is called when an exception occurs.
             result_callback (Callable | None): A function that is called when the group chat manager produces a result.
         """
         self._manager = manager
@@ -275,7 +275,7 @@ def __init__(
         self._participant_descriptions = participant_descriptions
         self._result_callback = result_callback
 
-        super().__init__(description="An actor for the group chat manager.")
+        super().__init__("An actor for the group chat manager.", exception_callback)
 
     @message_handler
     async def _handle_start_message(self, message: GroupChatStartMessage, ctx: MessageContext) -> None:
@@ -304,6 +304,7 @@ async def _handle_response_message(self, message: GroupChatResponseMessage, ctx:
 
         await self._determine_state_and_take_action(ctx.cancellation_token)
 
+    @ActorBase.exception_handler
     async def _determine_state_and_take_action(self, cancellation_token: CancellationToken) -> None:
         """Determine the state of the group chat and take action accordingly."""
         # User input state
@@ -448,14 +449,20 @@ async def _prepare(
         self,
         runtime: CoreRuntime,
         internal_topic_type: str,
+        exception_callback: Callable[[BaseException], None],
         result_callback: Callable[[DefaultTypeAlias], Awaitable[None]],
     ) -> None:
         """Register the actors and orchestrations with the runtime and add the required subscriptions."""
-        await self._register_members(runtime, internal_topic_type)
-        await self._register_manager(runtime, internal_topic_type, result_callback=result_callback)
+        await self._register_members(runtime, internal_topic_type, exception_callback)
+        await self._register_manager(runtime, internal_topic_type, exception_callback, result_callback)
         await self._add_subscriptions(runtime, internal_topic_type)
 
-    async def _register_members(self, runtime: CoreRuntime, internal_topic_type: str) -> None:
+    async def _register_members(
+        self,
+        runtime: CoreRuntime,
+        internal_topic_type: str,
+        exception_callback: Callable[[BaseException], None],
+    ) -> None:
         """Register the agents."""
         await asyncio.gather(*[
             GroupChatAgentActor.register(
@@ -464,6 +471,7 @@ async def _register_members(self, runtime: CoreRuntime, internal_topic_type: str
                 lambda agent=agent: GroupChatAgentActor(  # type: ignore[misc]
                     agent,
                     internal_topic_type,
+                    exception_callback=exception_callback,
                     agent_response_callback=self._agent_response_callback,
                     streaming_agent_response_callback=self._streaming_agent_response_callback,
                 ),
@@ -475,6 +483,7 @@ async def _register_manager(
         self,
         runtime: CoreRuntime,
         internal_topic_type: str,
+        exception_callback: Callable[[BaseException], None],
         result_callback: Callable[[DefaultTypeAlias], Awaitable[None]] | None = None,
     ) -> None:
         """Register the group chat manager."""
@@ -485,6 +494,7 @@ async def _register_manager(
                 self._manager,
                 internal_topic_type=internal_topic_type,
                 participant_descriptions={agent.name: agent.description for agent in self._members},  # type: ignore[misc]
+                exception_callback=exception_callback,
                 result_callback=result_callback,
             ),
         )

python/semantic_kernel/agents/orchestration/handoffs.py

@@ -9,7 +9,7 @@
 from functools import partial
 
 from semantic_kernel.agents.agent import Agent
-from semantic_kernel.agents.orchestration.agent_actor_base import AgentActorBase
+from semantic_kernel.agents.orchestration.agent_actor_base import ActorBase, AgentActorBase
 from semantic_kernel.agents.orchestration.orchestration_base import DefaultTypeAlias, OrchestrationBase, TIn, TOut
 from semantic_kernel.agents.runtime.core.cancellation_token import CancellationToken
 from semantic_kernel.agents.runtime.core.core_runtime import CoreRuntime
@@ -161,6 +161,7 @@ def __init__(
         agent: Agent,
         internal_topic_type: str,
         handoff_connections: AgentHandoffs,
+        exception_callback: Callable[[BaseException], None],
         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]
@@ -181,6 +182,7 @@ def __init__(
         super().__init__(
             agent=agent,
             internal_topic_type=internal_topic_type,
+            exception_callback=exception_callback,
             agent_response_callback=agent_response_callback,
             streaming_agent_response_callback=streaming_agent_response_callback,
         )
@@ -316,6 +318,7 @@ async def _call_human_response_function(self) -> ChatMessageContent:
             return await self._human_response_function()
         return self._human_response_function()  # type: ignore[return-value]
 
+    @ActorBase.exception_handler
     async def _invoke_agent_with_potentially_no_response(
         self, additional_messages: DefaultTypeAlias | None = None, **kwargs
     ) -> ChatMessageContent | None:
@@ -453,16 +456,18 @@ async def _prepare(
         self,
         runtime: CoreRuntime,
         internal_topic_type: str,
+        exception_callback: Callable[[BaseException], None],
         result_callback: Callable[[DefaultTypeAlias], Awaitable[None]],
     ) -> None:
         """Register the actors and orchestrations with the runtime and add the required subscriptions."""
-        await self._register_members(runtime, internal_topic_type, result_callback)
+        await self._register_members(runtime, internal_topic_type, exception_callback, result_callback)
         await self._add_subscriptions(runtime, internal_topic_type)
 
     async def _register_members(
         self,
         runtime: CoreRuntime,
         internal_topic_type: str,
+        exception_callback: Callable[[BaseException], None],
         result_callback: Callable[[DefaultTypeAlias], Awaitable[None]] | None = None,
     ) -> None:
         """Register the members with the runtime."""
@@ -476,6 +481,7 @@ async def _register_helper(agent: Agent) -> None:
                     agent,
                     internal_topic_type,
                     handoff_connections,
+                    exception_callback,
                     result_callback=result_callback,
                     agent_response_callback=self._agent_response_callback,
                     streaming_agent_response_callback=self._streaming_agent_response_callback,