Merge pull request #35 from AKKI0511/implement-real-time-websocket-data-streaming

Add streaming gateway and WebSocket infrastructure

Author: AKKI0511

config/streaming.yaml

@@ -0,0 +1,15 @@
+streaming:
+  providers:
+    - name: "alpaca"
+      websocket_url: "wss://stream.data.alpaca.markets/v2/iex"
+      auth_method: "api_key"
+      subscriptions: ["trades", "quotes"]
+      rate_limit:
+        default_rate: 100
+        burst_allowance: 50
+      circuit_breaker:
+        failure_threshold: 5
+        timeout: 30
+  buffer_size: 10000
+  reconnect_attempts: 5
+  health_check_interval: 30

poetry.lock

@@ -25,7 +25,12 @@ 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)",
+    "structlog (>=25.4.0,<26.0.0)",
+    "pybreaker (>=1.4.0,<2.0.0)",
+    "asyncio-throttle (>=1.0.2,<2.0.0)",
+    "prometheus-client (>=0.22.1,<0.23.0)"
 ]
 
 [tool.poetry]

pyproject.toml

@@ -0,0 +1,13 @@
+"""Streaming infrastructure package."""
+
+from .gateway import StreamingGateway
+from .auth_manager import AuthManager
+from .rate_limiter import AdaptiveRateLimiter
+from .connection_pool import ConnectionPool
+
+__all__ = [
+    "StreamingGateway",
+    "AuthManager",
+    "AdaptiveRateLimiter",
+    "ConnectionPool",
+]

quanttradeai/streaming/__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"]

quanttradeai/streaming/adapters/__init__.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}

quanttradeai/streaming/adapters/alpaca_adapter.py

@@ -0,0 +1,61 @@
+"""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
+
+from ..auth_manager import AuthManager
+from ..rate_limiter import AdaptiveRateLimiter
+
+
+@dataclass
+class DataProviderAdapter(ABC):
+    """Abstract adapter defining basic WebSocket operations."""
+
+    websocket_url: str
+    name: str
+    connection: Optional[websockets.WebSocketClientProtocol] = field(
+        default=None, init=False
+    )
+    circuit_breaker: Any = field(default=None, init=False)
+    rate_limiter: Optional[AdaptiveRateLimiter] = field(default=None, init=False)
+    auth_manager: Optional[AuthManager] = field(default=None, init=False)
+
+    async def connect(self) -> None:
+        """Establish a WebSocket connection to the provider."""
+
+        headers = None
+        if self.auth_manager:
+            headers = await self.auth_manager.get_auth_headers()
+        self.connection = await websockets.connect(
+            self.websocket_url, extra_headers=headers
+        )
+
+    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)
+        if self.rate_limiter:
+            await self.rate_limiter.acquire()
+        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."""

quanttradeai/streaming/adapters/base_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}

quanttradeai/streaming/adapters/ib_adapter.py

@@ -0,0 +1,37 @@
+"""Authentication and authorization utilities for data providers."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from typing import Dict, Optional
+
+
+@dataclass
+class AuthManager:
+    """Simple authentication manager with token refresh support."""
+
+    provider: str
+    _token: Optional[str] = None
+    _expires_at: datetime = datetime.fromtimestamp(0)
+
+    def _load_credentials(self) -> Dict[str, str]:
+        """Load credentials from environment variables."""
+        key = os.getenv(f"{self.provider.upper()}_API_KEY", "")
+        secret = os.getenv(f"{self.provider.upper()}_API_SECRET", "")
+        return {"key": key, "secret": secret}
+
+    def _token_needs_refresh(self) -> bool:
+        return datetime.utcnow() >= self._expires_at - timedelta(minutes=5)
+
+    async def _refresh_token(self) -> None:
+        creds = self._load_credentials()
+        # In production, call provider-specific auth endpoint.
+        self._token = creds.get("key")
+        self._expires_at = datetime.utcnow() + timedelta(hours=1)
+
+    async def get_auth_headers(self) -> Dict[str, str]:
+        if self._token is None or self._token_needs_refresh():
+            await self._refresh_token()
+        return {"Authorization": f"Bearer {self._token}"}

quanttradeai/streaming/auth_manager.py

@@ -0,0 +1,44 @@
+"""Connection pooling and health management."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import Any, Set, Callable, Awaitable
+
+
+@dataclass
+class ConnectionPool:
+    """Maintain a pool of reusable WebSocket connections."""
+
+    max_connections: int = 10
+    _pool: asyncio.Queue = field(init=False)
+    _active: Set[Any] = field(init=False, default_factory=set)
+
+    def __post_init__(self) -> None:
+        self._pool = asyncio.Queue(maxsize=self.max_connections)
+
+    async def acquire_connection(self, factory: Callable[[], Awaitable[Any]]) -> Any:
+        if not self._pool.empty():
+            conn = await self._pool.get()
+        elif len(self._active) < self.max_connections:
+            conn = await factory()
+        else:
+            conn = await self._pool.get()
+        self._active.add(conn)
+        return conn
+
+    async def release(self, conn: Any) -> None:
+        if conn in self._active:
+            self._active.remove(conn)
+            await self._pool.put(conn)
+
+    async def _ping_all_connections(self) -> None:
+        for conn in list(self._active):
+            if hasattr(conn, "ping"):
+                await conn.ping()
+
+    async def health_check_loop(self, interval: int = 30) -> None:
+        while True:
+            await self._ping_all_connections()
+            await asyncio.sleep(interval)