Files
ai/backend/src/datasource/websocket_protocol.py

171 lines
5.0 KiB
Python

"""
WebSocket protocol messages for TradingView-compatible datafeed API.
These messages define the wire format for client-server communication
over WebSocket for symbol search, historical data, and real-time subscriptions.
"""
from typing import Any, Dict, List, Literal, Optional, Union
from pydantic import BaseModel, Field
# ============================================================================
# Client -> Server Messages
# ============================================================================
class SearchSymbolsRequest(BaseModel):
"""Request to search for symbols matching a query"""
type: Literal["search_symbols"] = "search_symbols"
request_id: str = Field(description="Client-generated request ID for matching responses")
query: str = Field(description="Search query string")
symbol_type: Optional[str] = Field(default=None, description="Filter by instrument type")
exchange: Optional[str] = Field(default=None, description="Filter by exchange")
limit: int = Field(default=30, description="Maximum number of results")
class ResolveSymbolRequest(BaseModel):
"""Request full metadata for a specific symbol"""
type: Literal["resolve_symbol"] = "resolve_symbol"
request_id: str
symbol: str = Field(description="Symbol identifier to resolve")
class GetBarsRequest(BaseModel):
"""Request historical bar data"""
type: Literal["get_bars"] = "get_bars"
request_id: str
symbol: str
resolution: str = Field(description="Time resolution (e.g., '1', '5', '60', '1D')")
from_time: int = Field(description="Start time (Unix timestamp in seconds)")
to_time: int = Field(description="End time (Unix timestamp in seconds)")
countback: Optional[int] = Field(default=None, description="Maximum number of bars to return")
class SubscribeBarsRequest(BaseModel):
"""Subscribe to real-time bar updates"""
type: Literal["subscribe_bars"] = "subscribe_bars"
request_id: str
symbol: str
resolution: str
subscription_id: str = Field(description="Client-generated subscription ID")
class UnsubscribeBarsRequest(BaseModel):
"""Unsubscribe from real-time updates"""
type: Literal["unsubscribe_bars"] = "unsubscribe_bars"
request_id: str
subscription_id: str
class GetConfigRequest(BaseModel):
"""Request datafeed configuration"""
type: Literal["get_config"] = "get_config"
request_id: str
# Union of all client request types
ClientRequest = Union[
SearchSymbolsRequest,
ResolveSymbolRequest,
GetBarsRequest,
SubscribeBarsRequest,
UnsubscribeBarsRequest,
GetConfigRequest,
]
# ============================================================================
# Server -> Client Messages
# ============================================================================
class SearchSymbolsResponse(BaseModel):
"""Response with search results"""
type: Literal["search_symbols_response"] = "search_symbols_response"
request_id: str
results: List[Dict[str, Any]] = Field(description="List of SearchResult objects")
class ResolveSymbolResponse(BaseModel):
"""Response with symbol metadata"""
type: Literal["resolve_symbol_response"] = "resolve_symbol_response"
request_id: str
symbol_info: Dict[str, Any] = Field(description="SymbolInfo object")
class GetBarsResponse(BaseModel):
"""Response with historical bars"""
type: Literal["get_bars_response"] = "get_bars_response"
request_id: str
history: Dict[str, Any] = Field(description="HistoryResult object with bars and metadata")
class SubscribeBarsResponse(BaseModel):
"""Acknowledgment of subscription"""
type: Literal["subscribe_bars_response"] = "subscribe_bars_response"
request_id: str
subscription_id: str
success: bool
message: Optional[str] = None
class UnsubscribeBarsResponse(BaseModel):
"""Acknowledgment of unsubscribe"""
type: Literal["unsubscribe_bars_response"] = "unsubscribe_bars_response"
request_id: str
subscription_id: str
success: bool
class GetConfigResponse(BaseModel):
"""Response with datafeed configuration"""
type: Literal["get_config_response"] = "get_config_response"
request_id: str
config: Dict[str, Any] = Field(description="DatafeedConfig object")
class BarUpdateMessage(BaseModel):
"""Real-time bar update (server-initiated, no request_id)"""
type: Literal["bar_update"] = "bar_update"
subscription_id: str
symbol: str
resolution: str
bar: Dict[str, Any] = Field(description="Bar data including time and all columns")
class ErrorResponse(BaseModel):
"""Error response for any failed request"""
type: Literal["error"] = "error"
request_id: str
error_code: str = Field(description="Machine-readable error code")
error_message: str = Field(description="Human-readable error description")
# Union of all server response types
ServerResponse = Union[
SearchSymbolsResponse,
ResolveSymbolResponse,
GetBarsResponse,
SubscribeBarsResponse,
UnsubscribeBarsResponse,
GetConfigResponse,
BarUpdateMessage,
ErrorResponse,
]