144 lines
4.8 KiB
Python
144 lines
4.8 KiB
Python
"""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"
|