hyperforge 1.0.0.post19__py3-none-any.whl

hyperforge/broker/__init__.py

@@ -0,0 +1,52 @@
+from abc import ABC, abstractmethod
+from typing import AsyncIterator
+
+from hyperforge.pubsub import AgentMessage, StartInteraction
+
+
+class AgentTimeoutError(Exception):
+    pass
+
+
+class Broker(ABC):
+    # Activation
+    @abstractmethod
+    async def publish_activation(
+        self, msg: StartInteraction, trace: dict[str, str]
+    ) -> None: ...
+
+    @abstractmethod
+    def subscribe_activations(
+        self,
+    ) -> AsyncIterator[tuple[StartInteraction, dict[str, str]]]:
+        """Yields (StartInteraction, trace_headers) pairs. Called only by the server."""
+        ...
+
+    # Answer stream
+    @abstractmethod
+    async def publish(self, topic: str, message: AgentMessage) -> None: ...
+
+    @abstractmethod
+    def subscribe(
+        self, topic: str, from_cursor: str = "0"
+    ) -> AsyncIterator[tuple[str, AgentMessage]]:
+        """Yields (cursor, message) pairs. Raises AgentTimeoutError on keepalive timeout."""
+        ...
+
+    # Reply channel — used for feedback and OAuth callbacks
+    @abstractmethod
+    async def send_reply(self, key: str, payload: str) -> None: ...
+
+    @abstractmethod
+    async def receive_reply(self, key: str, timeout_ms: int) -> str | None: ...
+
+    # Lifecycle
+    @property
+    @abstractmethod
+    def keepalive_seconds(self) -> float: ...
+
+    @abstractmethod
+    async def initialize(self) -> None: ...
+
+    @abstractmethod
+    async def finalize(self) -> None: ...

hyperforge/broker/local.py

@@ -0,0 +1,116 @@
+import asyncio
+from typing import AsyncIterator
+
+from lru import LRU
+from pydantic import TypeAdapter
+
+from hyperforge.broker import AgentTimeoutError, Broker
+from hyperforge.pubsub import AgentMessage, StartInteraction
+
+
+class LocalBroker(Broker):
+    """In-process broker implementation. No Redis required.
+    API and server must share the same instance."""
+
+    def __init__(self, keepalive_ms: int = 20000, max_streams: int = 500):
+        self._keepalive_ms = keepalive_ms
+        self._activation_queue: asyncio.Queue[
+            tuple[StartInteraction, dict[str, str]]
+        ] = asyncio.Queue()
+        # topic -> (messages list, condition). Bounded by LRU to avoid unbounded growth.
+        self._streams: LRU = LRU(max_streams)
+        # key -> queue (single item). Deleted after receive_reply.
+        self._reply_channels: dict[str, asyncio.Queue[str]] = {}
+        self._counter = 0
+
+    def _next_id(self) -> str:
+        self._counter += 1
+        return str(self._counter)
+
+    def _get_or_create_stream(
+        self, topic: str
+    ) -> tuple[list[tuple[str, str]], asyncio.Condition]:
+        if topic not in self._streams:
+            self._streams[topic] = ([], asyncio.Condition())
+        return self._streams[topic]
+
+    @property
+    def keepalive_seconds(self) -> float:
+        return self._keepalive_ms / 1000
+
+    async def publish_activation(
+        self, msg: StartInteraction, trace: dict[str, str]
+    ) -> None:
+        await self._activation_queue.put((msg, trace))
+
+    async def subscribe_activations(
+        self,
+    ) -> AsyncIterator[tuple[StartInteraction, dict[str, str]]]:
+        while True:
+            msg, trace = await self._activation_queue.get()
+            yield msg, trace
+
+    async def publish(self, topic: str, message: AgentMessage) -> None:
+        messages, condition = self._get_or_create_stream(topic)
+        cursor = self._next_id()
+        async with condition:
+            messages.append((cursor, message.model_dump_json()))
+            condition.notify_all()
+
+    async def subscribe(
+        self, topic: str, from_cursor: str = "0"
+    ) -> AsyncIterator[tuple[str, AgentMessage]]:
+        messages, condition = self._get_or_create_stream(topic)
+        adapter: TypeAdapter[AgentMessage] = TypeAdapter(AgentMessage)
+
+        # Find the starting index from the cursor
+        last_index = 0
+        if from_cursor != "0":
+            for i, (cursor, _) in enumerate(messages):
+                if cursor == from_cursor:
+                    last_index = i + 1
+                    break
+
+        while True:
+            async with condition:
+                if last_index < len(messages):
+                    batch = messages[last_index:]
+                    last_index = len(messages)
+                else:
+                    try:
+                        await asyncio.wait_for(
+                            condition.wait(),
+                            timeout=self._keepalive_ms / 1000,
+                        )
+                    except asyncio.TimeoutError:
+                        raise AgentTimeoutError(topic)
+                    batch = messages[last_index:]
+                    last_index = len(messages)
+
+            for cursor, raw in batch:
+                yield cursor, adapter.validate_json(raw)
+
+    async def send_reply(self, key: str, payload: str) -> None:
+        if key not in self._reply_channels:
+            self._reply_channels[key] = asyncio.Queue(maxsize=1)
+        await self._reply_channels[key].put(payload)
+
+    async def receive_reply(self, key: str, timeout_ms: int) -> str | None:
+        if key not in self._reply_channels:
+            self._reply_channels[key] = asyncio.Queue(maxsize=1)
+        try:
+            return await asyncio.wait_for(
+                self._reply_channels[key].get(),
+                timeout=timeout_ms / 1000,
+            )
+        except asyncio.TimeoutError:
+            return None
+        finally:
+            self._reply_channels.pop(key, None)
+
+    async def initialize(self) -> None:
+        pass
+
+    async def finalize(self) -> None:
+        self._streams.clear()
+        self._reply_channels.clear()

hyperforge/broker/redis.py

@@ -0,0 +1,161 @@
+import asyncio
+import json
+from typing import AsyncIterator, cast
+
+import opentelemetry.propagate
+from pydantic import TypeAdapter
+from redis.asyncio import Redis, ResponseError
+
+from hyperforge import logger
+from hyperforge.broker import AgentTimeoutError, Broker
+from hyperforge.pubsub import AgentMessage, StartInteraction
+from hyperforge.redis_utils import ManualStreamKeysRedisCluster
+
+
+class RedisBroker(Broker):
+    def __init__(self, client: Redis, activate_subject: str, keepalive_ms: int):
+        self._client = client
+        self._activate_subject = activate_subject
+        self._keepalive_ms = int(keepalive_ms)
+
+    @property
+    def keepalive_seconds(self) -> float:
+        return self._keepalive_ms / 1000
+
+    @classmethod
+    def from_url(
+        cls,
+        url: str,
+        activate_subject: str,
+        keepalive_ms: int,
+        cluster_mode: bool = False,
+    ) -> "RedisBroker":
+        if cluster_mode:
+            client = cast(
+                Redis,
+                ManualStreamKeysRedisCluster.from_url(
+                    url=url,
+                    decode_responses=True,
+                    dynamic_startup_nodes=False,
+                ),
+            )
+        else:
+            client = Redis.from_url(
+                url,
+                decode_responses=True,
+            )
+        return cls(client, activate_subject, keepalive_ms)
+
+    async def publish_activation(
+        self, msg: StartInteraction, trace: dict[str, str]
+    ) -> None:
+        await self._client.xadd(
+            self._activate_subject,
+            {"msg": msg.model_dump_json(), "trace": json.dumps(trace)},
+            maxlen=100,
+        )
+
+    async def subscribe_activations(
+        self,
+    ) -> AsyncIterator[tuple[StartInteraction, dict[str, str]]]:
+        try:
+            await self._client.xgroup_create(
+                name=self._activate_subject,
+                groupname="arag_server",
+                mkstream=True,
+            )
+        except ResponseError as e:
+            if "BUSYGROUP" not in str(e):
+                raise
+
+        while True:
+            try:
+                response = await self._client.xreadgroup(
+                    groupname="arag_server",
+                    consumername="any_server",
+                    streams={self._activate_subject: ">"},
+                    block=1000,
+                    count=1,
+                    noack=True,
+                )
+                if not response:
+                    continue
+                [_stream, messages] = response[0]
+                if messages:
+                    [_msgid, fields] = messages[0]
+                    msg = StartInteraction.model_validate_json(fields["msg"])
+                    trace = json.loads(fields.get("trace", "{}"))
+                    yield msg, trace
+            except (asyncio.CancelledError, KeyboardInterrupt):
+                logger.info("Activation subscription cancelled, exiting...")
+                break
+            except Exception:
+                logger.exception("Error while subscribing to activations, retrying...")
+                await asyncio.sleep(1)
+
+    async def publish(self, topic: str, message: AgentMessage) -> None:
+        async with self._client.pipeline() as pipe:
+            await (
+                pipe.xadd(topic, {"msg": message.model_dump_json()}, maxlen=100)
+                .expire(topic, 300)
+                .execute()
+            )
+
+    async def subscribe(
+        self, topic: str, from_cursor: str = "0"
+    ) -> AsyncIterator[tuple[str, AgentMessage]]:
+        cursor = from_cursor
+        retries = 0
+        while True:
+            try:
+                response = await self._client.xread(
+                    {topic: cursor},
+                    block=max(50, int(self._keepalive_ms)),
+                )
+            except (asyncio.CancelledError, KeyboardInterrupt):
+                logger.info(f"Subscription to topic '{topic}' cancelled, exiting...")
+                break
+            except Exception:
+                logger.exception(
+                    f"Error while subscribing to topic '{topic}', retrying..."
+                )
+                retries += 1
+                if retries > 5:
+                    logger.error(
+                        f"Too many errors while subscribing to topic '{topic}', giving up..."
+                    )
+                    raise AgentTimeoutError(topic)
+                await asyncio.sleep(1)
+                continue
+            if not response:
+                raise AgentTimeoutError(topic)
+            [_stream, messages] = response[0]
+            for msgid, fields in messages:
+                cursor = msgid
+                obj: AgentMessage = TypeAdapter(AgentMessage).validate_json(
+                    fields["msg"]
+                )
+                yield cursor, obj
+
+    async def send_reply(self, key: str, payload: str) -> None:
+        trace_headers: dict[str, str] = {}
+        opentelemetry.propagate.inject(trace_headers)
+        await self._client.xadd(
+            key, {"msg": payload, "trace": json.dumps(trace_headers)}, maxlen=100
+        )
+
+    async def receive_reply(self, key: str, timeout_ms: int) -> str | None:
+        response = await self._client.xread(
+            {key: "$"},
+            block=timeout_ms,
+            count=1,
+        )
+        if not response:
+            return None
+        return response[0][1][0][1]["msg"]
+
+    async def initialize(self) -> None:
+        pass
+
+    async def finalize(self) -> None:
+        await self._client.aclose()  # type: ignore[attr-defined]