chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,205 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from semantic_kernel.agents.agent import Agent
|
||||
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
|
||||
from semantic_kernel.agents.runtime.core.message_context import MessageContext
|
||||
from semantic_kernel.agents.runtime.core.routed_agent import message_handler
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
|
||||
from semantic_kernel.kernel_pydantic import KernelBaseModel
|
||||
from semantic_kernel.utils.feature_stage_decorator import experimental
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
from typing import override # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # pragma: no cover
|
||||
|
||||
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@experimental
|
||||
class SequentialRequestMessage(KernelBaseModel):
|
||||
"""A request message type for sequential agents."""
|
||||
|
||||
body: DefaultTypeAlias
|
||||
|
||||
|
||||
@experimental
|
||||
class SequentialResultMessage(KernelBaseModel):
|
||||
"""A result message type for sequential agents."""
|
||||
|
||||
body: ChatMessageContent
|
||||
|
||||
|
||||
@experimental
|
||||
class SequentialAgentActor(AgentActorBase):
|
||||
"""A agent actor for sequential agents that process tasks."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: Agent,
|
||||
internal_topic_type: str,
|
||||
next_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,
|
||||
) -> None:
|
||||
"""Initialize the agent actor."""
|
||||
self._next_agent_type = next_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,
|
||||
)
|
||||
|
||||
@message_handler
|
||||
async def _handle_message(self, message: SequentialRequestMessage, ctx: MessageContext) -> None:
|
||||
"""Handle a message."""
|
||||
logger.debug(f"Sequential actor (Actor ID: {self.id}; Agent name: {self._agent.name}) started processing...")
|
||||
|
||||
response = await self._invoke_agent(additional_messages=message.body)
|
||||
|
||||
logger.debug(f"Sequential actor (Actor ID: {self.id}; Agent name: {self._agent.name}) finished processing.")
|
||||
|
||||
target_actor_id = await self.runtime.get(self._next_agent_type)
|
||||
await self.send_message(
|
||||
SequentialRequestMessage(body=response),
|
||||
target_actor_id,
|
||||
cancellation_token=ctx.cancellation_token,
|
||||
)
|
||||
|
||||
|
||||
@experimental
|
||||
class CollectionActor(ActorBase):
|
||||
"""A agent container for collection results from the last agent in the sequence."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
description: str,
|
||||
exception_callback: Callable[[BaseException], None],
|
||||
result_callback: Callable[[DefaultTypeAlias], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Initialize the collection actor."""
|
||||
self._result_callback = result_callback
|
||||
|
||||
super().__init__(description, exception_callback)
|
||||
|
||||
@message_handler
|
||||
async def _handle_message(self, message: SequentialRequestMessage, _: MessageContext) -> None:
|
||||
"""Handle the last message."""
|
||||
await self._result_callback(message.body)
|
||||
|
||||
|
||||
@experimental
|
||||
class SequentialOrchestration(OrchestrationBase[TIn, TOut]):
|
||||
"""A sequential multi-agent pattern orchestration."""
|
||||
|
||||
@override
|
||||
async def _start(
|
||||
self,
|
||||
task: DefaultTypeAlias,
|
||||
runtime: CoreRuntime,
|
||||
internal_topic_type: str,
|
||||
cancellation_token: CancellationToken,
|
||||
) -> None:
|
||||
"""Start the sequential pattern."""
|
||||
target_actor_id = await runtime.get(self._get_agent_actor_type(self._members[0], internal_topic_type))
|
||||
await runtime.send_message(
|
||||
SequentialRequestMessage(body=task),
|
||||
target_actor_id,
|
||||
cancellation_token=cancellation_token,
|
||||
)
|
||||
|
||||
@override
|
||||
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, exception_callback)
|
||||
await self._register_collection_actor(runtime, internal_topic_type, exception_callback, result_callback)
|
||||
|
||||
async def _register_members(
|
||||
self,
|
||||
runtime: CoreRuntime,
|
||||
internal_topic_type: str,
|
||||
exception_callback: Callable[[BaseException], None],
|
||||
) -> None:
|
||||
"""Register the members.
|
||||
|
||||
The members will be registered in the reverse order so that the actor type of the next worker
|
||||
is available when the current worker is registered. This is important for the sequential
|
||||
orchestration, where actors need to know its next actor type to send the message to.
|
||||
|
||||
Args:
|
||||
runtime (CoreRuntime): The agent runtime.
|
||||
internal_topic_type (str): The internal topic type for the orchestration that this actor is part of.
|
||||
exception_callback (Callable[[BaseException], None]): A callback function to handle exceptions.
|
||||
|
||||
Returns:
|
||||
str: The first actor type in the sequence.
|
||||
"""
|
||||
next_actor_type = self._get_collection_actor_type(internal_topic_type)
|
||||
for agent in reversed(self._members):
|
||||
await SequentialAgentActor.register(
|
||||
runtime,
|
||||
self._get_agent_actor_type(agent, internal_topic_type),
|
||||
lambda agent=agent, next_actor_type=next_actor_type: SequentialAgentActor( # type: ignore[misc]
|
||||
agent,
|
||||
internal_topic_type,
|
||||
next_agent_type=next_actor_type,
|
||||
exception_callback=exception_callback,
|
||||
agent_response_callback=self._agent_response_callback,
|
||||
streaming_agent_response_callback=self._streaming_agent_response_callback,
|
||||
),
|
||||
)
|
||||
logger.debug(f"Registered agent actor of type {self._get_agent_actor_type(agent, internal_topic_type)}")
|
||||
next_actor_type = self._get_agent_actor_type(agent, internal_topic_type)
|
||||
|
||||
async def _register_collection_actor(
|
||||
self,
|
||||
runtime: CoreRuntime,
|
||||
internal_topic_type: str,
|
||||
exception_callback: Callable[[BaseException], None],
|
||||
result_callback: Callable[[DefaultTypeAlias], Awaitable[None]],
|
||||
) -> None:
|
||||
"""Register the collection actor."""
|
||||
await CollectionActor.register(
|
||||
runtime,
|
||||
self._get_collection_actor_type(internal_topic_type),
|
||||
lambda: CollectionActor(
|
||||
description="An internal agent that is responsible for collection results",
|
||||
exception_callback=exception_callback,
|
||||
result_callback=result_callback,
|
||||
),
|
||||
)
|
||||
|
||||
def _get_agent_actor_type(self, agent: Agent, internal_topic_type: str) -> str:
|
||||
"""Get the actor type for an agent.
|
||||
|
||||
The type is appended with the internal topic type to ensure uniqueness in the runtime
|
||||
that may be shared by multiple orchestrations.
|
||||
"""
|
||||
return f"{agent.name}_{internal_topic_type}"
|
||||
|
||||
def _get_collection_actor_type(self, internal_topic_type: str) -> str:
|
||||
"""Get the collection actor type.
|
||||
|
||||
The type is appended with the internal topic type to ensure uniqueness in the runtime
|
||||
that may be shared by multiple orchestrations.
|
||||
"""
|
||||
return f"{CollectionActor.__name__}_{internal_topic_type}"
|
||||
Reference in New Issue
Block a user