trading-agent/agents/schemas.py

"""Pydantic output contract shared by the nightly agents and the morning executor.

The agents emit ``AgentOutput`` (containing ``Pick`` items) to JSON; the executor reads
it back. Validators enforce basic sanity so garbage never reaches order placement.
"""

from __future__ import annotations

from datetime import datetime
from enum import Enum

from pydantic import BaseModel, Field, field_validator, model_validator


class Direction(str, Enum):
    LONG = "long"
    SHORT = "short"


class Horizon(str, Enum):
    INTRADAY = "intraday"
    SWING = "swing"
    WEEKS = "weeks"
    MONTHS = "months"


class EntryType(str, Enum):
    MARKET = "market"
    LIMIT = "limit"


class Instrument(str, Enum):
    EQUITY = "equity"
    OPTION = "option"


class AgentStatus(str, Enum):
    COMPLETE = "complete"
    PARTIAL = "partial"
    TIMEOUT = "timeout"
    ERROR = "error"


class Pick(BaseModel):
    """One trade idea from an agent. Absolute prices are advisory (formed overnight on
    stale data); the executor derives real stop/take from pct off the actual fill."""

    ticker: str
    direction: Direction = Direction.LONG
    conviction: float = Field(ge=0.0, le=1.0)
    thesis: str
    catalysts: list[str] = Field(default_factory=list)
    time_horizon: Horizon = Horizon.SWING
    requested_weight_pct: float = Field(gt=0.0, le=1.0, description="fraction of equity")
    entry_type: EntryType = EntryType.MARKET
    limit_price: float | None = None
    stop_loss_pct: float | None = None
    stop_price: float | None = None
    take_profit_pct: float | None = None
    take_price: float | None = None
    risk_notes: str = ""
    data_refs: list[str] = Field(default_factory=list)
    # --- options (instrument=option): long calls/puts. contract_symbol from get_option_chain ---
    instrument: Instrument = Instrument.EQUITY
    option_right: str | None = None  # "call" | "put"
    contract_symbol: str | None = None  # OCC symbol, e.g. AAPL260724C00295000
    expiry: str | None = None
    strike: float | None = None

    @field_validator("ticker")
    @classmethod
    def _normalize_ticker(cls, v: str) -> str:
        v = v.strip().upper()
        if not v:
            raise ValueError("ticker must be non-empty")
        return v

    @field_validator("stop_loss_pct", "take_profit_pct")
    @classmethod
    def _pct_fraction(cls, v: float | None) -> float | None:
        if v is not None and not (0.0 < v < 1.0):
            raise ValueError("stop/take pct must be a fraction between 0 and 1")
        return v

    @model_validator(mode="after")
    def _check_option_fields(self) -> "Pick":
        if self.instrument == Instrument.OPTION:
            if not self.contract_symbol:
                raise ValueError("option pick requires contract_symbol (from get_option_chain)")
            if (self.option_right or "").lower() not in ("call", "put"):
                raise ValueError("option pick requires option_right 'call' or 'put'")
            self.option_right = self.option_right.lower()
        return self

    @model_validator(mode="after")
    def _check_absolute_levels(self) -> "Pick":
        # Only enforce ordering when an absolute entry reference exists (limit orders).
        ref = self.limit_price
        if ref is not None:
            if self.direction == Direction.LONG:
                if self.stop_price is not None and self.stop_price >= ref:
                    raise ValueError("long stop_price must be below entry")
                if self.take_price is not None and self.take_price <= ref:
                    raise ValueError("long take_price must be above entry")
            else:  # SHORT
                if self.stop_price is not None and self.stop_price <= ref:
                    raise ValueError("short stop_price must be above entry")
                if self.take_price is not None and self.take_price >= ref:
                    raise ValueError("short take_price must be below entry")
        return self


class AgentOutput(BaseModel):
    agent_name: str
    persona_id: str
    started_at: datetime
    ended_at: datetime | None = None
    status: AgentStatus = AgentStatus.PARTIAL
    steps_used: int = 0
    wall_seconds: float = 0.0
    market_summary: str = ""
    picks: list[Pick] = Field(default_factory=list)
    reasoning_trace: list[str] = Field(default_factory=list)
    errors: list[str] = Field(default_factory=list)


class RunManifest(BaseModel):
    run_id: str
    trade_date: str  # ET session date, YYYY-MM-DD
    model: str
    started_at: datetime
    ended_at: datetime | None = None
    status: str = "running"
    agents: list[str] = Field(default_factory=list)