wehub-resource-sync
206 lines
8.2 KiB
Python
206 lines
8.2 KiB
Python
# 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}"
|