""" Data models for the DataSource interface. Inspired by TradingView's Datafeed API but with flexible column schemas for AI-native trading platform needs. """ from enum import StrEnum from typing import Any, Dict, List, Literal, Optional from pydantic import BaseModel, Field class Resolution(StrEnum): """Standard time resolutions for bar data""" # Seconds S1 = "1S" S5 = "5S" S15 = "15S" S30 = "30S" # Minutes M1 = "1" M5 = "5" M15 = "15" M30 = "30" # Hours H1 = "60" H2 = "120" H4 = "240" H6 = "360" H12 = "720" # Days D1 = "1D" # Weeks W1 = "1W" # Months MO1 = "1M" class ColumnInfo(BaseModel): """ Metadata for a single data column. Provides rich, LLM-readable descriptions so AI agents can understand and reason about available data fields. """ model_config = {"extra": "forbid"} name: str = Field(description="Column name (e.g., 'close', 'volume', 'funding_rate')") type: Literal["float", "int", "bool", "string", "decimal"] = Field(description="Data type") description: str = Field(description="Human and LLM-readable description of what this column represents") unit: Optional[str] = Field(default=None, description="Unit of measurement (e.g., 'USD', 'BTC', '%', 'contracts')") nullable: bool = Field(default=False, description="Whether this column can contain null values") class SymbolInfo(BaseModel): """ Complete metadata for a tradeable symbol. Includes both TradingView-compatible fields and flexible schema definition for arbitrary data columns. """ model_config = {"extra": "forbid"} # Core identification symbol: str = Field(description="Unique symbol identifier (primary key for data fetching)") ticker: Optional[str] = Field(default=None, description="TradingView ticker (if different from symbol)") name: str = Field(description="Display name") description: str = Field(description="LLM-readable description of the instrument") type: str = Field(description="Instrument type: 'crypto', 'stock', 'forex', 'futures', 'derived', etc.") exchange: str = Field(description="Exchange or data source identifier") # Trading session info timezone: str = Field(default="Etc/UTC", description="IANA timezone identifier") session: str = Field(default="24x7", description="Trading session spec (e.g., '0930-1600' or '24x7')") # Resolution support supported_resolutions: List[str] = Field(description="List of supported time resolutions") has_intraday: bool = Field(default=True, description="Whether intraday resolutions are supported") has_daily: bool = Field(default=True, description="Whether daily resolution is supported") has_weekly_and_monthly: bool = Field(default=False, description="Whether weekly/monthly resolutions are supported") # Flexible schema definition columns: List[ColumnInfo] = Field(description="Available data columns for this symbol") time_column: str = Field(default="time", description="Name of the timestamp column") # Convenience flags has_ohlcv: bool = Field(default=False, description="Whether standard OHLCV columns are present") # Price display (for OHLCV data) pricescale: int = Field(default=100, description="Price scale factor (e.g., 100 for 2 decimals)") minmov: int = Field(default=1, description="Minimum price movement in pricescale units") # Additional metadata base_currency: Optional[str] = Field(default=None, description="Base currency (for crypto/forex)") quote_currency: Optional[str] = Field(default=None, description="Quote currency (for crypto/forex)") class Bar(BaseModel): """ A single bar/row of time-series data with flexible columns. All bars must have a timestamp. Additional columns are stored in the data dict and described by the associated ColumnInfo metadata. """ model_config = {"extra": "forbid"} time: int = Field(description="Unix timestamp in seconds") data: Dict[str, Any] = Field(description="Column name -> value mapping") # Convenience accessors for common OHLCV columns @property def open(self) -> Optional[float]: return self.data.get("open") @property def high(self) -> Optional[float]: return self.data.get("high") @property def low(self) -> Optional[float]: return self.data.get("low") @property def close(self) -> Optional[float]: return self.data.get("close") @property def volume(self) -> Optional[float]: return self.data.get("volume") class HistoryResult(BaseModel): """ Result from a historical data query. Includes the bars, schema information, and pagination metadata. """ model_config = {"extra": "forbid"} symbol: str = Field(description="Symbol identifier") resolution: str = Field(description="Time resolution of the bars") bars: List[Bar] = Field(description="The actual data bars") columns: List[ColumnInfo] = Field(description="Schema describing the bar data columns") nextTime: Optional[int] = Field(default=None, description="Unix timestamp for pagination (if more data available)") class SearchResult(BaseModel): """ A single result from symbol search. """ model_config = {"extra": "forbid"} symbol: str = Field(description="Display symbol (e.g., 'BINANCE:ETH/BTC')") ticker: Optional[str] = Field(default=None, description="Backend ticker for data fetching (e.g., 'ETH/BTC')") full_name: str = Field(description="Full display name including exchange") description: str = Field(description="Human-readable description") exchange: str = Field(description="Exchange identifier") type: str = Field(description="Instrument type") class DatafeedConfig(BaseModel): """ Configuration and capabilities of a DataSource. Similar to TradingView's onReady configuration object. """ model_config = {"extra": "forbid"} # Supported features supported_resolutions: List[str] = Field(description="All resolutions this datafeed supports") supports_search: bool = Field(default=True, description="Whether symbol search is available") supports_time: bool = Field(default=True, description="Whether time-based queries are supported") supports_marks: bool = Field(default=False, description="Whether marks/events are supported") # Data characteristics exchanges: List[str] = Field(default_factory=list, description="Available exchanges") symbols_types: List[str] = Field(default_factory=list, description="Available instrument types") # Metadata name: str = Field(description="Datafeed name") description: str = Field(description="LLM-readable description of this data source")