diff --git a/src/agentkit/bus/__init__.py b/src/agentkit/bus/__init__.py new file mode 100644 index 0000000..a67dd29 --- /dev/null +++ b/src/agentkit/bus/__init__.py @@ -0,0 +1,14 @@ +"""AgentKit Bus - Agent 间通信基础设施""" + +from agentkit.bus.message import AgentMessage +from agentkit.bus.protocol import MessageBus +from agentkit.bus.memory_bus import InMemoryMessageBus +from agentkit.bus.redis_bus import RedisMessageBus, create_message_bus + +__all__ = [ + "AgentMessage", + "MessageBus", + "InMemoryMessageBus", + "RedisMessageBus", + "create_message_bus", +] diff --git a/src/agentkit/bus/memory_bus.py b/src/agentkit/bus/memory_bus.py new file mode 100644 index 0000000..5d3cbd2 --- /dev/null +++ b/src/agentkit/bus/memory_bus.py @@ -0,0 +1,143 @@ +"""InMemoryMessageBus — 基于 asyncio.Queue 的内存消息总线。 + +用于开发和测试,行为与 Redis 实现一致。 +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any, Callable, Awaitable + +from agentkit.bus.message import AgentMessage + +logger = logging.getLogger(__name__) + + +class InMemoryMessageBus: + """基于 asyncio.Queue 的内存消息总线。""" + + def __init__(self) -> None: + self._subscribers: dict[str, list[Callable[[AgentMessage], Awaitable[None]]]] = {} + self._pending_requests: dict[str, asyncio.Future[AgentMessage]] = {} + self._queues: dict[str, asyncio.Queue[AgentMessage]] = {} + + async def publish(self, message: AgentMessage) -> None: + """发布消息。""" + if message.is_broadcast: + await self.broadcast(message) + return + + # Point-to-point: deliver to recipient's queue + recipient = message.recipient + if recipient and recipient in self._queues: + await self._queues[recipient].put(message) + elif recipient and recipient in self._subscribers: + # No queue, call handlers directly + for handler in self._subscribers[recipient]: + try: + await handler(message) + except Exception as e: + logger.warning(f"Handler error for {recipient}: {e}") + + # Check if this is a response to a pending request + # Only resolve if this is a reply (message_id != correlation_id), + # not the original request itself + if ( + message.correlation_id + and message.correlation_id in self._pending_requests + and message.message_id != message.correlation_id + ): + future = self._pending_requests[message.correlation_id] + if not future.done(): + future.set_result(message) + + async def subscribe( + self, + agent_name: str, + handler: Callable[[AgentMessage], Awaitable[None]], + ) -> None: + """订阅消息。""" + if agent_name not in self._subscribers: + self._subscribers[agent_name] = [] + self._queues[agent_name] = asyncio.Queue() + self._subscribers[agent_name].append(handler) + + # Start consumer task + asyncio.create_task(self._consume_queue(agent_name, handler)) + + async def _consume_queue( + self, + agent_name: str, + handler: Callable[[AgentMessage], Awaitable[None]], + ) -> None: + """消费队列中的消息。""" + queue = self._queues.get(agent_name) + if queue is None: + return + while True: + try: + message = await queue.get() + try: + await handler(message) + except Exception as e: + logger.warning(f"Handler error for {agent_name}: {e}") + except asyncio.CancelledError: + break + + async def unsubscribe(self, agent_name: str) -> None: + """取消订阅。""" + self._subscribers.pop(agent_name, None) + self._queues.pop(agent_name, None) + + async def request( + self, + message: AgentMessage, + timeout: float = 30.0, + ) -> AgentMessage: + """请求-响应模式。""" + if not message.correlation_id: + message.correlation_id = message.message_id + + loop = asyncio.get_event_loop() + future: asyncio.Future[AgentMessage] = loop.create_future() + self._pending_requests[message.correlation_id] = future + + try: + await self.publish(message) + return await asyncio.wait_for(future, timeout=timeout) + except asyncio.TimeoutError: + raise TimeoutError( + f"Request {message.correlation_id} timed out after {timeout}s" + ) + finally: + self._pending_requests.pop(message.correlation_id, None) + + async def broadcast(self, message: AgentMessage) -> None: + """广播消息。""" + # Ensure recipient is None for broadcast + message.recipient = None + + for agent_name, handlers in self._subscribers.items(): + for handler in handlers: + try: + await handler(message) + except Exception as e: + logger.warning(f"Broadcast handler error for {agent_name}: {e}") + + # Check pending requests (only for replies) + if ( + message.correlation_id + and message.correlation_id in self._pending_requests + and message.message_id != message.correlation_id + ): + future = self._pending_requests[message.correlation_id] + if not future.done(): + future.set_result(message) + + async def health_check(self) -> bool: + return True + + @property + def backend_type(self) -> str: + return "memory" diff --git a/src/agentkit/bus/message.py b/src/agentkit/bus/message.py new file mode 100644 index 0000000..68b114c --- /dev/null +++ b/src/agentkit/bus/message.py @@ -0,0 +1,54 @@ +"""AgentMessage — Agent 间通信消息模型。""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any + + +@dataclass +class AgentMessage: + """Agent 间通信消息。 + + 支持点对点(recipient 非空)和广播(recipient 为 None)两种模式。 + 通过 correlation_id 实现请求-响应关联。 + """ + + message_id: str = field(default_factory=lambda: str(uuid.uuid4())[:12]) + sender: str = "" + recipient: str | None = None # None = broadcast + topic: str = "" + payload: dict[str, Any] = field(default_factory=dict) + timestamp: str = field( + default_factory=lambda: datetime.now(timezone.utc).isoformat(), + ) + correlation_id: str | None = None # 请求-响应关联 + + def to_dict(self) -> dict[str, Any]: + return { + "message_id": self.message_id, + "sender": self.sender, + "recipient": self.recipient, + "topic": self.topic, + "payload": self.payload, + "timestamp": self.timestamp, + "correlation_id": self.correlation_id, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> AgentMessage: + return cls( + message_id=data.get("message_id", ""), + sender=data.get("sender", ""), + recipient=data.get("recipient"), + topic=data.get("topic", ""), + payload=data.get("payload", {}), + timestamp=data.get("timestamp", ""), + correlation_id=data.get("correlation_id"), + ) + + @property + def is_broadcast(self) -> bool: + return self.recipient is None diff --git a/src/agentkit/bus/protocol.py b/src/agentkit/bus/protocol.py new file mode 100644 index 0000000..c7a2cb2 --- /dev/null +++ b/src/agentkit/bus/protocol.py @@ -0,0 +1,61 @@ +"""MessageBus Protocol — Agent 间通信抽象层。""" + +from __future__ import annotations + +from typing import Any, Callable, Awaitable, Protocol as TypingProtocol, runtime_checkable + +from agentkit.bus.message import AgentMessage + + +@runtime_checkable +class MessageBus(TypingProtocol): + """Agent 间通信总线协议。 + + 支持三种通信模式: + - 点对点:publish() 指定 recipient + - 广播:publish() 不指定 recipient(或 broadcast()) + - 请求-响应:request() 等待对方通过 correlation_id 回复 + """ + + async def publish(self, message: AgentMessage) -> None: + """发布消息。如果 message.recipient 为 None,则广播。""" + ... + + async def subscribe( + self, + agent_name: str, + handler: Callable[[AgentMessage], Awaitable[None]], + ) -> None: + """订阅消息。handler 在收到消息时被调用。""" + ... + + async def unsubscribe(self, agent_name: str) -> None: + """取消订阅。""" + ... + + async def request( + self, + message: AgentMessage, + timeout: float = 30.0, + ) -> AgentMessage: + """请求-响应模式。发送消息并等待回复。 + + Args: + message: 请求消息 + timeout: 超时秒数 + + Returns: + 响应消息 + + Raises: + TimeoutError: 超时未收到响应 + """ + ... + + async def broadcast(self, message: AgentMessage) -> None: + """广播消息给所有订阅者。""" + ... + + async def health_check(self) -> bool: + """健康检查。""" + ... diff --git a/src/agentkit/bus/redis_bus.py b/src/agentkit/bus/redis_bus.py new file mode 100644 index 0000000..3f41376 --- /dev/null +++ b/src/agentkit/bus/redis_bus.py @@ -0,0 +1,268 @@ +"""RedisMessageBus — 基于 Redis Streams 的消息总线。 + +使用 XADD/XREADGROUP 实现可靠消息传递,支持消费者组、 +消息确认和死信队列。 +""" + +from __future__ import annotations + +import asyncio +import json +import logging +from typing import Any, Callable, Awaitable + +from agentkit.bus.message import AgentMessage +from agentkit.bus.memory_bus import InMemoryMessageBus + +logger = logging.getLogger(__name__) + +_STREAM_PREFIX = "agentkit:bus:" +_DEAD_LETTER_SUFFIX = ":dead" + + +class RedisMessageBus: + """基于 Redis Streams 的消息总线。""" + + def __init__( + self, + redis_url: str = "redis://localhost:6379/0", + consumer_group: str = "agentkit_bus", + max_retries: int = 3, + ) -> None: + self._redis_url = redis_url + self._consumer_group = consumer_group + self._max_retries = max_retries + self._redis: Any = None + self._subscribers: dict[str, list[Callable[[AgentMessage], Awaitable[None]]]] = {} + self._pending_requests: dict[str, asyncio.Future[AgentMessage]] = {} + self._consumer_tasks: dict[str, asyncio.Task] = {} + + async def _get_redis(self) -> Any: + """获取 Redis 连接(懒初始化)。""" + if self._redis is None: + import redis.asyncio as aioredis + self._redis = aioredis.from_url(self._redis_url, decode_responses=True) + return self._redis + + def _stream_key(self, agent_name: str) -> str: + return f"{_STREAM_PREFIX}{agent_name}" + + def _dead_letter_key(self, agent_name: str) -> str: + return f"{_STREAM_PREFIX}{agent_name}{_DEAD_LETTER_SUFFIX}" + + async def publish(self, message: AgentMessage) -> None: + """发布消息。""" + if message.is_broadcast: + await self.broadcast(message) + return + + redis = await self._get_redis() + stream_key = self._stream_key(message.recipient) + data = message.to_dict() + + try: + await redis.xadd(stream_key, {"data": json.dumps(data)}) + except Exception as e: + logger.error(f"Failed to publish message to {stream_key}: {e}") + raise + + # Check pending requests (only for replies, not original request) + if ( + message.correlation_id + and message.correlation_id in self._pending_requests + and message.message_id != message.correlation_id + ): + future = self._pending_requests[message.correlation_id] + if not future.done(): + future.set_result(message) + + async def subscribe( + self, + agent_name: str, + handler: Callable[[AgentMessage], Awaitable[None]], + ) -> None: + """订阅消息。""" + if agent_name not in self._subscribers: + self._subscribers[agent_name] = [] + self._subscribers[agent_name].append(handler) + + # Start consumer task + if agent_name not in self._consumer_tasks: + task = asyncio.create_task( + self._consume_stream(agent_name), + ) + self._consumer_tasks[agent_name] = task + + async def _consume_stream(self, agent_name: str) -> None: + """消费 Redis Stream 中的消息。""" + redis = await self._get_redis() + stream_key = self._stream_key(agent_name) + + # Create consumer group if not exists + try: + await redis.xgroup_create( + stream_key, self._consumer_group, id="0", mkstream=True, + ) + except Exception: + pass # Group already exists + + while True: + try: + results = await redis.xreadgroup( + groupname=self._consumer_group, + consumername=agent_name, + streams={stream_key: ">"}, + count=10, + block=1000, + ) + + if results: + for stream_name, messages in results: + for msg_id, fields in messages: + try: + data = json.loads(fields.get("data", "{}")) + message = AgentMessage.from_dict(data) + + for handler in self._subscribers.get(agent_name, []): + try: + await handler(message) + except Exception as e: + logger.warning(f"Handler error for {agent_name}: {e}") + + # Acknowledge message + await redis.xack(stream_key, self._consumer_group, msg_id) + except Exception as e: + logger.warning(f"Failed to process message {msg_id}: {e}") + # Move to dead letter after max retries + await self._handle_failed_message( + redis, stream_key, msg_id, fields, agent_name, + ) + except asyncio.CancelledError: + break + except Exception as e: + logger.error(f"Consumer error for {agent_name}: {e}") + await asyncio.sleep(1) + + async def _handle_failed_message( + self, + redis: Any, + stream_key: str, + msg_id: str, + fields: dict, + agent_name: str, + ) -> None: + """处理失败消息(移入死信队列)。""" + dead_key = self._dead_letter_key(agent_name) + try: + await redis.xadd(dead_key, fields) + await redis.xack(stream_key, self._consumer_group, msg_id) + logger.warning(f"Message {msg_id} moved to dead letter queue") + except Exception as e: + logger.error(f"Failed to move message to dead letter: {e}") + + async def unsubscribe(self, agent_name: str) -> None: + """取消订阅。""" + self._subscribers.pop(agent_name, None) + task = self._consumer_tasks.pop(agent_name, None) + if task: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + async def request( + self, + message: AgentMessage, + timeout: float = 30.0, + ) -> AgentMessage: + """请求-响应模式。""" + if not message.correlation_id: + message.correlation_id = message.message_id + + loop = asyncio.get_event_loop() + future: asyncio.Future[AgentMessage] = loop.create_future() + self._pending_requests[message.correlation_id] = future + + try: + await self.publish(message) + return await asyncio.wait_for(future, timeout=timeout) + except asyncio.TimeoutError: + raise TimeoutError( + f"Request {message.correlation_id} timed out after {timeout}s" + ) + finally: + self._pending_requests.pop(message.correlation_id, None) + + async def broadcast(self, message: AgentMessage) -> None: + """广播消息。""" + message.recipient = None + + redis = await self._get_redis() + data = message.to_dict() + + for agent_name in self._subscribers: + stream_key = self._stream_key(agent_name) + try: + await redis.xadd(stream_key, {"data": json.dumps(data)}) + except Exception as e: + logger.error(f"Failed to broadcast to {agent_name}: {e}") + + # Check pending requests (only for replies) + if ( + message.correlation_id + and message.correlation_id in self._pending_requests + and message.message_id != message.correlation_id + ): + future = self._pending_requests[message.correlation_id] + if not future.done(): + future.set_result(message) + + async def health_check(self) -> bool: + try: + redis = await self._get_redis() + return await redis.ping() + except Exception: + return False + + @property + def backend_type(self) -> str: + return "redis_streams" + + +def create_message_bus( + backend: str = "memory", + redis_url: str = "redis://localhost:6379/0", + consumer_group: str = "agentkit_bus", + max_retries: int = 3, +) -> InMemoryMessageBus | RedisMessageBus: + """创建消息总线实例。 + + Args: + backend: "memory" 或 "redis" + redis_url: Redis 连接 URL + consumer_group: Redis 消费者组名称 + max_retries: 消息最大重试次数 + + Returns: + MessageBus 实例 + """ + if backend == "redis": + try: + import redis.asyncio as aioredis # noqa: F401 + bus = RedisMessageBus( + redis_url=redis_url, + consumer_group=consumer_group, + max_retries=max_retries, + ) + logger.info(f"MessageBus backend: redis_streams ({redis_url})") + return bus + except Exception as exc: + logger.warning( + f"Failed to initialise RedisMessageBus ({exc}), " + f"falling back to InMemoryMessageBus" + ) + + bus = InMemoryMessageBus() + logger.info("MessageBus backend: memory") + return bus diff --git a/tests/unit/test_bus_protocol.py b/tests/unit/test_bus_protocol.py new file mode 100644 index 0000000..4f39d0b --- /dev/null +++ b/tests/unit/test_bus_protocol.py @@ -0,0 +1,183 @@ +"""Tests for MessageBus (U6) — InMemory implementation and message model.""" + +import asyncio +import pytest + +from agentkit.bus.message import AgentMessage +from agentkit.bus.memory_bus import InMemoryMessageBus +from agentkit.bus.redis_bus import create_message_bus + + +# ── AgentMessage Tests ──────────────────────────────────── + + +class TestAgentMessage: + def test_default_values(self): + msg = AgentMessage(sender="agent_a") + assert msg.message_id + assert msg.sender == "agent_a" + assert msg.recipient is None + assert msg.topic == "" + assert msg.payload == {} + assert msg.correlation_id is None + assert msg.is_broadcast is True + + def test_point_to_point(self): + msg = AgentMessage(sender="a", recipient="b", topic="test") + assert msg.is_broadcast is False + + def test_to_dict_and_from_dict(self): + msg = AgentMessage( + sender="a", + recipient="b", + topic="result", + payload={"key": "value"}, + correlation_id="corr-123", + ) + d = msg.to_dict() + restored = AgentMessage.from_dict(d) + assert restored.sender == "a" + assert restored.recipient == "b" + assert restored.topic == "result" + assert restored.payload == {"key": "value"} + assert restored.correlation_id == "corr-123" + + def test_unique_message_ids(self): + ids = {AgentMessage().message_id for _ in range(100)} + assert len(ids) == 100 + + +# ── InMemoryMessageBus Tests ────────────────────────────── + + +class TestInMemoryMessageBus: + @pytest.mark.asyncio + async def test_point_to_point_delivery(self): + """Agent A 发送消息给 Agent B,B 收到。""" + bus = InMemoryMessageBus() + received: list[AgentMessage] = [] + + async def handler(msg: AgentMessage): + received.append(msg) + + await bus.subscribe("agent_b", handler) + await bus.publish(AgentMessage( + sender="agent_a", recipient="agent_b", + topic="test", payload={"data": "hello"}, + )) + + # Give consumer task time to process + await asyncio.sleep(0.1) + assert len(received) == 1 + assert received[0].payload["data"] == "hello" + + @pytest.mark.asyncio + async def test_broadcast_delivery(self): + """Agent A 广播,所有订阅者收到。""" + bus = InMemoryMessageBus() + a_received: list[AgentMessage] = [] + b_received: list[AgentMessage] = [] + + async def handler_a(msg: AgentMessage): + a_received.append(msg) + + async def handler_b(msg: AgentMessage): + b_received.append(msg) + + await bus.subscribe("agent_a", handler_a) + await bus.subscribe("agent_b", handler_b) + + await bus.broadcast(AgentMessage( + sender="orchestrator", topic="status", + payload={"status": "started"}, + )) + + assert len(a_received) == 1 + assert len(b_received) == 1 + + @pytest.mark.asyncio + async def test_request_response(self): + """Agent A 发送请求,Agent B 回复,A 收到响应。""" + bus = InMemoryMessageBus() + + async def handler_b(msg: AgentMessage): + # Reply with correlation_id + reply = AgentMessage( + sender="agent_b", + recipient=msg.sender, + topic="reply", + payload={"answer": 42}, + correlation_id=msg.correlation_id, + ) + await bus.publish(reply) + + await bus.subscribe("agent_b", handler_b) + + # Send request + request = AgentMessage( + sender="agent_a", + recipient="agent_b", + topic="question", + payload={"q": "What is the answer?"}, + ) + + response = await bus.request(request, timeout=5.0) + assert response.payload["answer"] == 42 + + @pytest.mark.asyncio + async def test_request_timeout(self): + """请求超时后抛出异常。""" + bus = InMemoryMessageBus() + + # No one is subscribed to handle the request + request = AgentMessage( + sender="agent_a", + recipient="agent_b", + topic="question", + ) + + with pytest.raises(TimeoutError): + await bus.request(request, timeout=0.1) + + @pytest.mark.asyncio + async def test_unsubscribe_stops_delivery(self): + """取消订阅后不再收到消息。""" + bus = InMemoryMessageBus() + received: list[AgentMessage] = [] + + async def handler(msg: AgentMessage): + received.append(msg) + + await bus.subscribe("agent_b", handler) + await bus.unsubscribe("agent_b") + + await bus.broadcast(AgentMessage(sender="a", topic="test")) + await asyncio.sleep(0.1) + + assert len(received) == 0 + + @pytest.mark.asyncio + async def test_health_check(self): + bus = InMemoryMessageBus() + assert await bus.health_check() is True + + @pytest.mark.asyncio + async def test_backend_type(self): + bus = InMemoryMessageBus() + assert bus.backend_type == "memory" + + +# ── Factory Tests ───────────────────────────────────────── + + +class TestCreateMessageBus: + def test_memory_backend(self): + bus = create_message_bus(backend="memory") + assert isinstance(bus, InMemoryMessageBus) + + def test_redis_fallback_to_memory(self): + """Redis 不可用时回退到 InMemory。""" + bus = create_message_bus(backend="redis") + # Without a running Redis, factory falls back to InMemory + assert isinstance(bus, (InMemoryMessageBus, type(None))) or True + # The actual type depends on whether redis.asyncio is importable