mirror of
https://github.com/microsoft/autogen.git
synced 2025-07-19 15:01:52 +00:00
e1a823fb6d
* initial impl of topics and subscriptions * Update python/src/agnext/core/_agent_runtime.py Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com> * add topic in context * migrate * migrate code for topics * migrate team one * edit notebooks * formatting * fix imports * Build proto * Fix circular import --------- Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import pytest
|
|
from agnext.application import SingleThreadedAgentRuntime
|
|
|
|
from agnext.components._type_subscription import TypeSubscription
|
|
from agnext.core import AgentRuntime, AgentId
|
|
|
|
from agnext.components import ClosureAgent
|
|
|
|
|
|
import asyncio
|
|
|
|
from agnext.core import MessageContext
|
|
from agnext.core import TopicId
|
|
|
|
@dataclass
|
|
class Message:
|
|
content: str
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_register_receives_publish() -> None:
|
|
runtime = SingleThreadedAgentRuntime()
|
|
|
|
queue = asyncio.Queue[tuple[str, str]]()
|
|
|
|
async def log_message(_runtime: AgentRuntime, id: AgentId, message: Message, ctx: MessageContext) -> None:
|
|
key = id.key
|
|
await queue.put((key, message.content))
|
|
|
|
await runtime.register("name", lambda: ClosureAgent("my_agent", log_message))
|
|
await runtime.add_subscription(TypeSubscription("default", "name"))
|
|
topic_id = TopicId("default", "default")
|
|
run_context = runtime.start()
|
|
|
|
await runtime.publish_message(Message("first message"), topic_id=topic_id)
|
|
await runtime.publish_message(Message("second message"), topic_id=topic_id)
|
|
await runtime.publish_message(Message("third message"), topic_id=topic_id)
|
|
|
|
|
|
await run_context.stop_when_idle()
|
|
|
|
assert queue.qsize() == 3
|
|
assert queue.get_nowait() == ("default", "first message")
|
|
assert queue.get_nowait() == ("default", "second message")
|
|
assert queue.get_nowait() == ("default", "third message")
|
|
assert queue.empty()
|