ai/backend.old/src/datasource/schema.py

"""
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")