Improve HTTPChannel

Author: RussellLuo

coagent/runtimes/http_runtime.py

@@ -44,7 +44,7 @@ class HTTPChannel(BaseChannel):
     """An HTTP-based channel.
 
     _publish: POST /publish
-    _publish_stream: POST /publish_multi
+    _publish_stream: POST /publish stream=True
     subscribe: POST /subscribe
     new_reply_topic: POST /reply-topics
     """
@@ -68,17 +68,17 @@ async def _publish(
         self,
         addr: Address,
         msg: RawMessage,
-        request: bool = False,
         stream: bool = False,
+        request: bool = False,
         reply: str = "",
         timeout: float = 5.0,
         probe: bool = True,
     ) -> RawMessage | None:
         data = dict(
             addr=addr.encode(mode="json"),
             msg=msg.encode(mode="json"),
-            request=request,
             stream=stream,
+            request=request,
             reply=reply,
             timeout=timeout,
             probe=probe,
@@ -108,13 +108,14 @@ async def _publish_stream(
         data = dict(
             addr=addr.encode(mode="json"),
             msg=msg.encode(mode="json"),
+            stream=True,
             probe=probe,
         )
         headers = {"Authorization": self._auth} if self._auth else None
 
         queue: QueueSubscriptionIterator = QueueSubscriptionIterator()
         sub: HTTPChannelSubscription = HTTPChannelSubscription(
-            f"{self._server}/publish_multi", data, headers, queue.receive
+            f"{self._server}/publish", data, headers, queue.receive
         )
         await sub.subscribe()
 
@@ -268,37 +269,22 @@ async def publish(
         self,
         addr: Address,
         msg: RawMessage,
-        request: bool = False,
         stream: bool = False,
+        request: bool = False,
         reply: str = "",
         timeout: float = 5.0,
         probe: bool = True,
     ) -> RawMessage | None:
         return await self._channel.publish(
             addr,
             msg,
-            request=request,
             stream=stream,
+            request=request,
             reply=reply,
             timeout=timeout,
             probe=probe,
         )
 
-    async def publish_multi(
-        self,
-        addr: Address,
-        msg: RawMessage,
-        probe: bool = True,
-    ) -> AsyncIterator[RawMessage]:
-        msgs = await self._channel.publish(
-            addr,
-            msg,
-            stream=True,
-            probe=probe,
-        )
-        async for msg in msgs:
-            yield msg
-
     async def subscribe(
         self,
         addr: Address,

coagent/runtimes/local_runtime.py

@@ -69,8 +69,8 @@ async def _publish(
         self,
         addr: Address,
         msg: RawMessage,
-        request: bool = False,
         stream: bool = False,
+        request: bool = False,
         reply: str = "",
         timeout: float = 0.5,
         probe: bool = True,

coagent/runtimes/nats_runtime.py

@@ -77,8 +77,8 @@ async def _publish(
         self,
         addr: Address,
         msg: RawMessage,
-        request: bool = False,
         stream: bool = False,
+        request: bool = False,
         reply: str = "",
         timeout: float = 0.5,
         probe: bool = True,

examples/ping-pong/http_runtime_server.py

@@ -27,15 +27,40 @@ async def shutdown():
 
 async def publish(request):
     data: dict = await request.json()
+
+    addr: Address = Address.model_validate(data["addr"])
+    msg: RawMessage = RawMessage.model_validate(data["msg"])
+    stream: bool = data.get("stream", False)
+    request: bool = data.get("request", False)
+    reply: str = data.get("reply", "")
+    timeout: float = data.get("timeout", 0.5)
+    probe: bool = data.get("probe", True)
+
+    # Streaming
+    if stream:
+        msgs: AsyncIterator[RawMessage] = await backend.publish(
+            addr=addr, msg=msg, stream=stream, probe=probe
+        )
+
+        async def event_stream() -> AsyncIterator[str]:
+            try:
+                async for raw in msgs:
+                    yield dict(data=raw.encode_json())
+            except BaseError as exc:
+                yield dict(event="error", data=exc.encode_json())
+
+        return EventSourceResponse(event_stream())
+
+    # Non-streaming
     try:
         resp: RawMessage | None = await backend.publish(
-            addr=Address.model_validate(data["addr"]),
-            msg=RawMessage.model_validate(data["msg"]),
-            request=data.get("request", False),
-            stream=data.get("stream", False),
-            reply=data.get("reply", ""),
-            timeout=data.get("timeout", 0.5),
-            probe=data.get("probe", True),
+            addr=addr,
+            msg=msg,
+            stream=stream,
+            request=request,
+            reply=reply,
+            timeout=timeout,
+            probe=probe,
         )
     except BaseError as exc:
         return JSONResponse(exc.encode(mode="json"), status_code=404)
@@ -46,24 +71,6 @@ async def publish(request):
         return JSONResponse(resp.encode(mode="json"))
 
 
-async def publish_multi(request):
-    data: dict = await request.json()
-    msgs = backend.publish_multi(
-        addr=Address.model_validate(data["addr"]),
-        msg=RawMessage.model_validate(data["msg"]),
-        probe=data.get("probe", True),
-    )
-
-    async def event_stream() -> AsyncIterator[str]:
-        try:
-            async for raw in msgs:
-                yield dict(data=raw.encode_json())
-        except BaseError as exc:
-            yield dict(event="error", data=exc.encode_json())
-
-    return EventSourceResponse(event_stream())
-
-
 async def subscribe(request):
     data: dict = await request.json()
     msgs: AsyncIterator[RawMessage] = backend.subscribe(
@@ -85,7 +92,6 @@ async def new_reply_topic(request):
 
 routes = [
     Route("/publish", publish, methods=["POST"]),
-    Route("/publish_multi", publish_multi, methods=["POST"]),
     Route("/subscribe", subscribe, methods=["POST"]),
     Route("/reply-topics", new_reply_topic, methods=["POST"]),
 ]