""" 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, ]