Add streaming gateway and WebSocket infrastructure

AKKI0511 · AKKI0511 · commit ec5ccf5350a4 · 2025-08-17T18:20:16.000-04:00
diff --git a/config/streaming.yaml b/config/streaming.yaml
@@ -0,0 +1,9 @@
+streaming:
+  providers:
+    - name: "alpaca"
+      websocket_url: "wss://stream.data.alpaca.markets/v2/iex"
+      auth_method: "api_key"
+      subscriptions: ["trades", "quotes"]
+  buffer_size: 10000
+  reconnect_attempts: 5
+  health_check_interval: 30
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -25,7 +25,8 @@ dependencies = [
     "joblib (>=1.5.1,<2.0.0)",
     "pydantic (>=2.11.7,<3.0.0)",
     "alpha-vantage (>=3.0.0,<4.0.0)",
-    "litellm (>=1.75.0,<2.0.0)"
+    "litellm (>=1.75.0,<2.0.0)",
+    "websockets (>=15.0.1,<16.0.0)"
 ]
 
 [tool.poetry]
diff --git a/quanttradeai/streaming/__init__.py b/quanttradeai/streaming/__init__.py
@@ -0,0 +1,5 @@
+"""Streaming infrastructure package."""
+
+from .gateway import StreamingGateway
+
+__all__ = ["StreamingGateway"]
diff --git a/quanttradeai/streaming/adapters/__init__.py b/quanttradeai/streaming/adapters/__init__.py
@@ -0,0 +1,7 @@
+"""Provider adapter implementations."""
+
+from .base_adapter import DataProviderAdapter
+from .alpaca_adapter import AlpacaAdapter
+from .ib_adapter import IBAdapter
+
+__all__ = ["DataProviderAdapter", "AlpacaAdapter", "IBAdapter"]
diff --git a/quanttradeai/streaming/adapters/alpaca_adapter.py b/quanttradeai/streaming/adapters/alpaca_adapter.py
@@ -0,0 +1,20 @@
+"""Alpaca data provider adapter."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+from .base_adapter import DataProviderAdapter
+
+
+@dataclass
+class AlpacaAdapter(DataProviderAdapter):
+    """Adapter for Alpaca's streaming API."""
+
+    name: str = "alpaca"
+
+    def _build_subscribe_message(
+        self, channel: str, symbols: List[str]
+    ) -> Dict[str, Any]:
+        return {"action": "subscribe", channel: symbols}
diff --git a/quanttradeai/streaming/adapters/base_adapter.py b/quanttradeai/streaming/adapters/base_adapter.py
@@ -0,0 +1,48 @@
+"""Base classes for data provider adapters."""
+
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, AsyncIterator, Dict, List, Optional
+
+import websockets
+
+
+@dataclass
+class DataProviderAdapter(ABC):
+    """Abstract adapter defining basic WebSocket operations."""
+
+    websocket_url: str
+    name: str
+    connection: Optional[websockets.WebSocketClientProtocol] = field(
+        default=None, init=False
+    )
+
+    async def connect(self) -> None:
+        """Establish a WebSocket connection to the provider."""
+
+        self.connection = await websockets.connect(self.websocket_url)
+
+    async def subscribe(self, channel: str, symbols: List[str]) -> None:
+        """Send a subscription message for ``symbols`` on ``channel``."""
+
+        if self.connection is None:
+            await self.connect()
+        message = self._build_subscribe_message(channel, symbols)
+        await self.connection.send(json.dumps(message))
+
+    async def listen(self) -> AsyncIterator[Dict[str, Any]]:
+        """Yield parsed JSON messages from the connection."""
+
+        if self.connection is None:
+            raise RuntimeError("Connection not established")
+        async for message in self.connection:
+            yield json.loads(message)
+
+    @abstractmethod
+    def _build_subscribe_message(
+        self, channel: str, symbols: List[str]
+    ) -> Dict[str, Any]:
+        """Return provider specific subscription payload."""
diff --git a/quanttradeai/streaming/adapters/ib_adapter.py b/quanttradeai/streaming/adapters/ib_adapter.py
@@ -0,0 +1,20 @@
+"""Interactive Brokers data provider adapter."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+from .base_adapter import DataProviderAdapter
+
+
+@dataclass
+class IBAdapter(DataProviderAdapter):
+    """Adapter for Interactive Brokers TWS/Gateway streaming API."""
+
+    name: str = "interactive_brokers"
+
+    def _build_subscribe_message(
+        self, channel: str, symbols: List[str]
+    ) -> Dict[str, Any]:
+        return {"action": "subscribe", "channel": channel, "symbols": symbols}
diff --git a/quanttradeai/streaming/gateway.py b/quanttradeai/streaming/gateway.py
@@ -0,0 +1,90 @@
+"""High-level streaming gateway orchestrator."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Tuple
+
+import yaml
+
+from .stream_buffer import StreamBuffer
+from .websocket_manager import WebSocketManager
+from .adapters.alpaca_adapter import AlpacaAdapter
+from .adapters.ib_adapter import IBAdapter
+from .processors import MessageProcessor
+from .monitoring.metrics import Metrics
+
+AdapterMap = {
+    "alpaca": AlpacaAdapter,
+    "interactive_brokers": IBAdapter,
+}
+
+Callback = Callable[[str, Dict], None]
+
+
+@dataclass
+class StreamingGateway:
+    """Main entry point for real-time data streaming."""
+
+    config_path: str
+    websocket_manager: WebSocketManager = field(init=False)
+    message_processor: MessageProcessor = field(init=False)
+    buffer: StreamBuffer = field(init=False)
+    metrics: Metrics = field(init=False)
+    _subscriptions: List[Tuple[str, List[str]]] = field(default_factory=list)
+    _callbacks: Dict[str, List[Callback]] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        with open(self.config_path, "r") as f:
+            cfg = yaml.safe_load(f)["streaming"]
+        self.websocket_manager = WebSocketManager(
+            reconnect_attempts=cfg.get("reconnect_attempts", 5)
+        )
+        self.message_processor = MessageProcessor()
+        self.buffer = StreamBuffer(cfg.get("buffer_size", 1000))
+        self.metrics = Metrics()
+        for provider_cfg in cfg.get("providers", []):
+            name = provider_cfg["name"]
+            url = provider_cfg["websocket_url"]
+            adapter_cls = AdapterMap.get(name)
+            if adapter_cls is None:
+                raise ValueError(f"Unknown provider: {name}")
+            adapter = adapter_cls(websocket_url=url)
+            self.websocket_manager.add_adapter(adapter)
+
+    # Subscription API -------------------------------------------------
+    def subscribe_to_trades(self, symbols: List[str], callback: Callback) -> None:
+        """Subscribe to trade updates for ``symbols``."""
+
+        self._subscriptions.append(("trades", symbols))
+        self._callbacks.setdefault("trades", []).append(callback)
+
+    def subscribe_to_quotes(self, symbols: List[str], callback: Callback) -> None:
+        """Subscribe to quote updates for ``symbols``."""
+
+        self._subscriptions.append(("quotes", symbols))
+        self._callbacks.setdefault("quotes", []).append(callback)
+
+    # Runtime ----------------------------------------------------------
+    async def _dispatch(self, provider: str, message: Dict) -> None:
+        self.metrics.increment()
+        processed = self.message_processor.process(message)
+        await self.buffer.put(processed)
+        msg_type = message.get("type", "trades")
+        for cb in self._callbacks.get(msg_type, []):
+            res = cb(processed)
+            if asyncio.iscoroutine(res):
+                await res
+
+    async def _start(self) -> None:
+        await self.websocket_manager.connect_all()
+        for channel, symbols in self._subscriptions:
+            for adapter in self.websocket_manager.adapters:
+                await adapter.subscribe(channel, symbols)
+        await self.websocket_manager.run(self._dispatch)
+
+    def start_streaming(self) -> None:
+        """Blocking call that starts the streaming event loop."""
+
+        asyncio.run(self._start())
diff --git a/quanttradeai/streaming/monitoring/__init__.py b/quanttradeai/streaming/monitoring/__init__.py
@@ -0,0 +1,6 @@
+"""Monitoring utilities for streaming system."""
+
+from .health_monitor import HealthMonitor
+from .metrics import Metrics
+
+__all__ = ["HealthMonitor", "Metrics"]
diff --git a/quanttradeai/streaming/monitoring/health_monitor.py b/quanttradeai/streaming/monitoring/health_monitor.py
@@ -0,0 +1,19 @@
+"""Health monitoring utilities."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+
+
+@dataclass
+class HealthMonitor:
+    """Periodically perform a no-op health check."""
+
+    interval: int = 30
+
+    async def run(self) -> None:
+        """Run the health check loop."""
+
+        while True:
+            await asyncio.sleep(self.interval)
diff --git a/quanttradeai/streaming/monitoring/metrics.py b/quanttradeai/streaming/monitoring/metrics.py
@@ -0,0 +1,17 @@
+"""Streaming metrics collection."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class Metrics:
+    """Simple in-memory metrics tracker."""
+
+    messages_received: int = 0
+
+    def increment(self) -> None:
+        """Increment the received message counter."""
+
+        self.messages_received += 1
diff --git a/quanttradeai/streaming/processors/__init__.py b/quanttradeai/streaming/processors/__init__.py
@@ -0,0 +1,6 @@
+"""Message processing utilities."""
+
+from .message_processor import MessageProcessor
+from .validators import validate_trade_message
+
+__all__ = ["MessageProcessor", "validate_trade_message"]
diff --git a/quanttradeai/streaming/processors/message_processor.py b/quanttradeai/streaming/processors/message_processor.py
@@ -0,0 +1,19 @@
+"""Basic message transformation logic."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict
+
+
+@dataclass
+class MessageProcessor:
+    """Simple pass-through processor for streaming messages."""
+
+    def process(self, message: Dict[str, Any]) -> Dict[str, Any]:
+        """Return the message unchanged.
+
+        This placeholder allows future transformation logic.
+        """
+
+        return message
diff --git a/quanttradeai/streaming/processors/validators.py b/quanttradeai/streaming/processors/validators.py
@@ -0,0 +1,11 @@
+"""Validation utilities for streaming messages."""
+
+from __future__ import annotations
+
+from typing import Any, Dict
+
+
+def validate_trade_message(message: Dict[str, Any]) -> bool:
+    """Validate that a trade message contains required fields."""
+
+    return "symbol" in message
diff --git a/quanttradeai/streaming/stream_buffer.py b/quanttradeai/streaming/stream_buffer.py
@@ -0,0 +1,34 @@
+"""Async buffer for streaming messages."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from asyncio import Queue
+from typing import Any
+
+
+@dataclass
+class StreamBuffer:
+    """Thread-safe buffer based on :class:`asyncio.Queue`.
+
+    Parameters
+    ----------
+    maxsize:
+        Maximum number of messages to buffer.
+    """
+
+    maxsize: int
+    queue: Queue = field(init=False)
+
+    def __post_init__(self) -> None:
+        self.queue = Queue(maxsize=self.maxsize)
+
+    async def put(self, item: Any) -> None:
+        """Put a message into the buffer."""
+
+        await self.queue.put(item)
+
+    async def get(self) -> Any:
+        """Retrieve a message from the buffer."""
+
+        return await self.queue.get()
diff --git a/quanttradeai/streaming/websocket_manager.py b/quanttradeai/streaming/websocket_manager.py
@@ -0,0 +1,51 @@
+"""Manage multiple WebSocket connections."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import Awaitable, Callable, List
+
+from .adapters.base_adapter import DataProviderAdapter
+
+Callback = Callable[[str, dict], Awaitable[None]]
+
+
+@dataclass
+class WebSocketManager:
+    """Handle connection life-cycle for multiple adapters."""
+
+    reconnect_attempts: int = 5
+    adapters: List[DataProviderAdapter] = field(default_factory=list)
+
+    def add_adapter(self, adapter: DataProviderAdapter) -> None:
+        """Register a new data provider adapter."""
+
+        self.adapters.append(adapter)
+
+    async def _connect_with_retry(self, adapter: DataProviderAdapter) -> None:
+        delay = 1.0
+        for attempt in range(self.reconnect_attempts):
+            try:
+                await adapter.connect()
+                return
+            except Exception:
+                if attempt == self.reconnect_attempts - 1:
+                    raise
+                await asyncio.sleep(delay)
+                delay *= 2
+
+    async def connect_all(self) -> None:
+        """Connect all registered adapters."""
+
+        for adapter in self.adapters:
+            await self._connect_with_retry(adapter)
+
+    async def run(self, callback: Callback) -> None:
+        """Start listening on all connections and dispatch to ``callback``."""
+
+        async def listen(adapter: DataProviderAdapter) -> None:
+            async for msg in adapter.listen():
+                await callback(adapter.name, msg)
+
+        await asyncio.gather(*(listen(ad) for ad in self.adapters))
diff --git a/tests/streaming/test_gateway.py b/tests/streaming/test_gateway.py

Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,8 @@ dependencies = [`
`25`	`25`	`"joblib (>=1.5.1,<2.0.0)",`
`26`	`26`	`"pydantic (>=2.11.7,<3.0.0)",`
`27`	`27`	`"alpha-vantage (>=3.0.0,<4.0.0)",`
`28`		`- "litellm (>=1.75.0,<2.0.0)"`
	`28`	`+ "litellm (>=1.75.0,<2.0.0)",`
	`29`	`+ "websockets (>=15.0.1,<16.0.0)"`
`29`	`30`	`]`
`30`	`31`
`31`	`32`	`[tool.poetry]`