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

import re
from datetime import datetime
from enum import Enum

from pydantic import BaseModel, Field, field_validator, model_validator

# OCC option symbol: ROOT (<=6 chars) + YYMMDD + C/P + strike*1000 as 8 digits,
# e.g. AAPL260724C00295000. LLMs sometimes fabricate near-miss symbols; reject them
# here so an invented contract can never reach order placement.
_OCC_SYMBOL_RE = re.compile(r"^[A-Z.]{1,6}\d{6}[CP]\d{8}$")


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)")
            self.contract_symbol = self.contract_symbol.strip().upper()
            if not _OCC_SYMBOL_RE.match(self.contract_symbol):
                raise ValueError(
                    f"contract_symbol '{self.contract_symbol}' is not a valid OCC symbol — "
                    "copy it EXACTLY from get_option_chain output, never construct one"
                )
            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()
            # We only BUY options. A put marked direction=short is the same bearish bet
            # expressed sloppily (LLMs do this) — normalize to a long put. A short call has
            # no long-only equivalent, so reject it outright.
            if self.direction == Direction.SHORT:
                if self.option_right == "put":
                    self.direction = Direction.LONG
                else:
                    raise ValueError(
                        "cannot short a call (selling options unsupported) — express a "
                        "bearish view with a long put instead"
                    )
        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)