Aau/#40/refacto events

Makefile

@@ -7,4 +7,4 @@ lint:
 test:
 	pytest
 
-check: lint test
+check: lint test

pyproject.toml

@@ -25,6 +25,9 @@ strict_optional = true
 warn_unused_ignores = true
 warn_return_any = true
 warn_unused_configs = true
+exclude = [
+  "^tests/",
+]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]

src/core/bus.py

@@ -0,0 +1,94 @@
+import asyncio
+import logging
+from collections import defaultdict
+from typing import Any, AsyncGenerator, Coroutine, cast
+
+from .events import EventData
+from .module import Module
+
+logger = logging.getLogger("ray.serve")
+
+
+class EventGraph:
+    """
+    Asynchronous event routing system for HuRI modules.
+
+    The EventGraph is responsible for:
+        - Registering module subscribers
+        - Routing events between modules
+        - Executing module pipelines asynchronously
+        - Handling coroutine and async generator outputs
+
+    Modules subscribe to events through their `input_type`.
+    When an event is published, all subscribed modules are executed
+    concurrently.
+
+    Supports:
+        - Coroutine-based modules
+        - Async generator streaming modules
+        - Recursive event propagation
+
+    :subscribers:
+        Dictionary mapping event topics to subscribed modules.
+    """
+
+    def __init__(self):
+        self.subscribers = defaultdict(list)
+
+    def register(self, module: Module):
+        self.subscribers[module.input_type].append(module)
+
+    async def publish(self, event_topic, data):
+        subs = self.subscribers[event_topic]
+        if event_topic not in ("audio_in",):  # skip mic-frame spam
+            logger.info(
+                "[GRAPH] publish topic=%r subscribers=%s",
+                event_topic,
+                [type(m).__name__ for m in subs],
+            )
+        for module in subs:
+            asyncio.create_task(self._run(module, data))
+
+    async def _run(self, module: Module, data):
+        try:
+            result = module.process(data)
+
+            if hasattr(result, "__aiter__"):
+                try:
+                    generator = cast(AsyncGenerator[EventData | None, None], result)
+                    async for item in generator:
+                        if item is None:
+                            continue
+                        logger.info(
+                            "[GRAPH] %s -> %r: %s",
+                            type(module).__name__,
+                            module.output_type,
+                            item.summarize(),
+                        )
+                        await self.publish(module.output_type, item)
+                except Exception:
+                    logger.exception(
+                        "[GRAPH] async generator failed in %s", type(module).__name__
+                    )
+
+            else:
+                coroutine = cast(Coroutine[Any, Any, EventData | None], result)
+                try:
+                    value = await coroutine
+                    if value is not None:
+                        logger.info(
+                            "[GRAPH] %s -> %r: %s",
+                            type(module).__name__,
+                            module.output_type,
+                            value.summarize(),
+                        )
+                        await self.publish(module.output_type, value)
+                except Exception:
+                    logger.exception(
+                        "[GRAPH] coroutine failed in %s", type(module).__name__
+                    )
+
+        except Exception:
+            logger.exception(
+                "[GRAPH] process() call failed in %s", type(module).__name__
+            )

src/core/client.py

@@ -4,15 +4,14 @@
 import struct
 from collections import defaultdict
 from dataclasses import asdict
-from typing import Any, Dict, Generic, List, Type, TypeVar
+from typing import Any, Dict, Generic, List, Mapping, Type, TypeVar
 
-import numpy as np
 import websockets
 
 from src.core.dataclasses.config import ClientConfig
 from src.core.events import EventData
 
-T = TypeVar("T", bound=EventData | bytes)
+T = TypeVar("T", bound=EventData)
 
 
 class ClientSender(Generic[T]):
@@ -45,16 +44,20 @@ async def _send_bytes(self, ws: websockets.ClientConnection, data: bytes):
 
         await ws.send(packet)
 
-    async def _send_event_data(self, ws: websockets.ClientConnection, data: EventData):
-        packet = json.dumps({"topic": self.topic, "data": asdict(data)})
+    async def _send_event_data(
+        self, ws: websockets.ClientConnection, data: Mapping[str, Any]
+    ):
+        packet = json.dumps({"topic": self.topic, "data": data})
 
         await ws.send(packet)
 
     async def send(self, ws: websockets.ClientConnection, data: T):
-        if isinstance(data, bytes):
-            await self._send_bytes(ws, data)
+        wire = data.to_wire()
+
+        if isinstance(wire, bytes):
+            await self._send_bytes(ws, wire)
         else:
-            await self._send_event_data(ws, data)
+            await self._send_event_data(ws, wire)
 
 
 class ClientHook(Generic[T]):
@@ -120,40 +123,12 @@ async def _receive_loop(self, ws: websockets.ClientConnection):
                     topic = msg[2 : 2 + topic_len].decode()
                     data = msg[2 + topic_len :]
 
-                    if topic == "audio" and len(data) >= 13:
-                        sample_rate, end, pts = struct.unpack(">IBd", data[:13])
-                        # Samples are native-endian float32
-                        # (Sender uses ndarray.tobytes()).
-                        samples = np.frombuffer(data[13:], dtype=np.float32)
-                        data = {
-                            "sample_rate": sample_rate,
-                            "end": end,
-                            "pts": pts,
-                            "data": samples,
-                        }
-                    elif topic == "motion" and len(data) >= 16:
-                        pts, fps, n_frames = struct.unpack(">dII", data[:16])
-                        print(f"<< motion: pts={pts:.3f}s frames={n_frames} @ {fps}fps")
-                        data = {
-                            "poses": np.ndarray(0),
-                            "expressions": np.ndarray(0),
-                            "trans": np.ndarray(0),
-                            "fps": fps,
-                            "pts": pts,
-                        }
-                    else:
-                        print(f"<< {topic}: bytes ({len(data)}B)")
                 else:
                     event = json.loads(msg)
                     topic = event["topic"]
                     data = event["data"]
 
                 for hook in self.hooks[topic]:
-                    # Hydrate via from_wire (not a bare **data splat) so events
-                    # with nested EventData fields — e.g. RAGQuestion.transcript /
-                    # .emotion — are rebuilt as dataclasses, mirroring the server's
-                    # EventDataFactory. Build a fresh instance per hook so a topic
-                    # with several hooks doesn't re-splat an already-built event.
                     hook_data = (
                         data
                         if isinstance(data, bytes)

src/core/events.py

@@ -1,121 +1,50 @@
-import asyncio
 import logging
-from collections import defaultdict
-from dataclasses import dataclass
-from typing import Any, Mapping
-
-import numpy as np
-
-from .module import Module
+from dataclasses import asdict, dataclass
+from typing import Any, Generic, Mapping, TypeVar
 
 logger = logging.getLogger("ray.serve")
 
+WireT = TypeVar("WireT", Mapping[str, Any], bytes)
+
 
 @dataclass
-class EventData:
+class EventData(Generic[WireT]):
     """An event data must be derived from this class, and use @dataclass decorator.
     Or they can be bytes."""
 
     @classmethod
-    def from_wire(cls, data: Mapping[str, Any]) -> "EventData":
-        """Build an event from a JSON payload sent by a client.
-
-        Default: keyword-splat the payload onto the dataclass. Events whose fields
-        are nested dataclasses (or that accept a simpler external shape than the
-        in-pipeline one) override this — e.g. RAGQuestion accepts a bare
-        ``{"text": ...}`` typed question. In-process producers construct the
-        dataclass directly and never go through this path.
-        """
-        return cls(**data)
-
-
-class EventGraph:
-    """
-    Asynchronous event routing system for HuRI modules.
-
-    The EventGraph is responsible for:
-        - Registering module subscribers
-        - Routing events between modules
-        - Executing module pipelines asynchronously
-        - Handling coroutine and async generator outputs
+    def from_wire(cls, data: WireT) -> "EventData[WireT]":
+        raise NotImplementedError
 
-    Modules subscribe to events through their `input_type`.
-    When an event is published, all subscribed modules are executed
-    concurrently.
+    def to_wire(self) -> WireT:
+        raise NotImplementedError
 
-    Supports:
-        - Coroutine-based modules
-        - Async generator streaming modules
-        - Recursive event propagation
+    def summarize(self) -> str:
+        cls = type(self).__name__
+        return f"{cls}({self!r})"
 
-    :subscribers:
-        Dictionary mapping event topics to subscribed modules.
-    """
 
-    def __init__(self):
-        self.subscribers = defaultdict(list)
-
-    def register(self, module: Module):
-        self.subscribers[module.input_type].append(module)
-
-    async def publish(self, event_topic, data):
-        subs = self.subscribers[event_topic]
-        if event_topic not in ("audio_in",):  # skip mic-frame spam
-            logger.info(
-                "[GRAPH] publish topic=%r subscribers=%s",
-                event_topic,
-                [type(m).__name__ for m in subs],
-            )
-        for module in subs:
-            asyncio.create_task(self._run(module, data))
+@dataclass
+class JsonEvent(EventData[Mapping[str, Any]]):
+    @classmethod
+    def from_wire(cls, data: Mapping[str, Any]) -> "JsonEvent":
+        return cls(**data)
 
-    async def _run(self, module: Module, data):
-        try:
-            result = module.process(data)
+    def to_wire(self) -> Mapping[str, Any]:
+        return asdict(self)
 
-            if hasattr(result, "__aiter__"):
-                try:
-                    async for item in result:
-                        if item is None:
-                            continue
-                        logger.info(
-                            "[GRAPH] %s -> %r: %s",
-                            type(module).__name__,
-                            module.output_type,
-                            _summarize(item),
-                        )
-                        await self.publish(module.output_type, item)
-                except Exception:
-                    logger.exception(
-                        "[GRAPH] async generator failed in %s", type(module).__name__
-                    )
 
-            else:
-                try:
-                    value = await result
-                    if value is not None:
-                        logger.info(
-                            "[GRAPH] %s -> %r: %s",
-                            type(module).__name__,
-                            module.output_type,
-                            _summarize(value),
-                        )
-                        await self.publish(module.output_type, value)
-                except Exception:
-                    logger.exception(
-                        "[GRAPH] coroutine failed in %s", type(module).__name__
-                    )
+@dataclass
+class BytesEvent(EventData[bytes]):
+    data: bytes
 
-        except Exception:
-            logger.exception(
-                "[GRAPH] process() call failed in %s", type(module).__name__
-            )
+    @classmethod
+    def from_wire(cls, data: bytes) -> "BytesEvent":
+        return cls(data=data)
 
+    def to_wire(self) -> bytes:
+        return self.data
 
-def _summarize(item) -> str:
-    """Short repr that avoids dumping full numpy arrays into the log."""
-    cls = type(item).__name__
-    data = getattr(item, "data", None)
-    if isinstance(data, np.ndarray):
-        return f"{cls}(shape={data.shape}, dtype={data.dtype})"
-    return f"{cls}({item!r})"
+    def summarize(self) -> str:
+        cls = type(self).__name__
+        return f"{cls}(len:{len(self.data)})"

src/core/huri.py

@@ -32,7 +32,7 @@ def __init__(
         self,
         modules: Dict[str, Type[Module]],
         handles: Dict[str, handle.DeploymentHandle],
-        events: Dict[str, Type[EventData | bytes]],
+        events: Dict[str, Type[EventData]],
     ) -> None:
         self.module_factory = ModuleFactory(handles)
         self.event_factory = EventDataFactory()

src/core/module.py

@@ -2,6 +2,8 @@
 
 from ray.serve import handle
 
+from .events import EventData
+
 
 class Module:
     """
@@ -24,7 +26,9 @@ class Module:
     input_type: str
     output_type: Optional[str]
 
-    def process(self, _) -> Coroutine[Any, Any, Any] | AsyncGenerator[Any, None]:
+    def process(
+        self, _
+    ) -> Coroutine[Any, Any, EventData | None] | AsyncGenerator[EventData | None, None]:
         raise NotImplementedError
 
 

src/core/session.py

@@ -1,6 +1,6 @@
 from typing import List
 
-from .events import EventGraph
+from .bus import EventGraph
 from .module import Module