mirror of
https://github.com/microsoft/autogen.git
synced 2025-08-28 10:30:08 +00:00
91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
from dataclasses import dataclass
|
|
|
|
import pytest
|
|
from agnext.application import SingleThreadedAgentRuntime
|
|
from agnext.components import TypeRoutedAgent, message_handler
|
|
from agnext.core import Agent, AgentRuntime, CancellationToken
|
|
from agnext.core.exceptions import MessageDroppedException
|
|
from agnext.core.intervention import DefaultInterventionHandler, DropMessage
|
|
|
|
|
|
@dataclass
|
|
class MessageType:
|
|
...
|
|
|
|
class LoopbackAgent(TypeRoutedAgent): # type: ignore
|
|
def __init__(self, name: str, router: AgentRuntime) -> None: # type: ignore
|
|
super().__init__(name, "A loop back agent.", router)
|
|
self.num_calls = 0
|
|
|
|
|
|
@message_handler() # type: ignore
|
|
async def on_new_message(self, message: MessageType, cancellation_token: CancellationToken) -> MessageType: # type: ignore
|
|
self.num_calls += 1
|
|
return message
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_intervention_count_messages() -> None:
|
|
|
|
class DebugInterventionHandler(DefaultInterventionHandler): # type: ignore
|
|
def __init__(self) -> None:
|
|
self.num_messages = 0
|
|
|
|
async def on_send(self, message: MessageType, *, sender: Agent | None, recipient: Agent) -> MessageType: # type: ignore
|
|
self.num_messages += 1
|
|
return message
|
|
|
|
handler = DebugInterventionHandler()
|
|
router = SingleThreadedAgentRuntime(before_send=handler)
|
|
|
|
long_running = LoopbackAgent("name", router)
|
|
response = router.send_message(MessageType(), recipient=long_running)
|
|
|
|
while not response.done():
|
|
await router.process_next()
|
|
|
|
assert handler.num_messages == 1
|
|
assert long_running.num_calls == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_intervention_drop_send() -> None:
|
|
|
|
class DropSendInterventionHandler(DefaultInterventionHandler): # type: ignore
|
|
async def on_send(self, message: MessageType, *, sender: Agent | None, recipient: Agent) -> MessageType | type[DropMessage]: # type: ignore
|
|
return DropMessage # type: ignore
|
|
|
|
handler = DropSendInterventionHandler()
|
|
router = SingleThreadedAgentRuntime(before_send=handler)
|
|
|
|
long_running = LoopbackAgent("name", router)
|
|
response = router.send_message(MessageType(), recipient=long_running)
|
|
|
|
while not response.done():
|
|
await router.process_next()
|
|
|
|
with pytest.raises(MessageDroppedException):
|
|
await response
|
|
|
|
assert long_running.num_calls == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_intervention_drop_response() -> None:
|
|
|
|
class DropResponseInterventionHandler(DefaultInterventionHandler): # type: ignore
|
|
async def on_response(self, message: MessageType, *, sender: Agent, recipient: Agent | None) -> MessageType | type[DropMessage]: # type: ignore
|
|
return DropMessage # type: ignore
|
|
|
|
handler = DropResponseInterventionHandler()
|
|
router = SingleThreadedAgentRuntime(before_send=handler)
|
|
|
|
long_running = LoopbackAgent("name", router)
|
|
response = router.send_message(MessageType(), recipient=long_running)
|
|
|
|
while not response.done():
|
|
await router.process_next()
|
|
|
|
with pytest.raises(MessageDroppedException):
|
|
await response
|
|
|
|
assert long_running.num_calls == 1
|