trading-agent/execution/guardrails.py

"""Pure, unit-tested risk guardrails. Turns the agents' picks into a safe order plan.

No network, no Alpaca objects — takes plain data (equity, prices, held symbols) so the
whole pipeline is deterministic and testable. Handles equities and single-leg long options
(sized by premium × 100). The LLM's requested sizing / stops are inputs that get clamped
here; nothing is trusted as-is.

Pipeline: collect (long-only unless enabled; options quality gate) -> dedup across agents
(highest conviction) -> per-agent trade cap (option slot reserved) -> price/liquidity filter
+ skip already-held -> size & clamp weight (options premium-capped) -> per-day cap ->
buying-power trim (reserved option dropped last).

The "reserved option": each night one options trade is mandated per agent, so the
highest-conviction option that passes the quality gate is protected from the caps and the
buying-power trim — otherwise conviction-ranked truncation would routinely cut it.
"""

from __future__ import annotations

from dataclasses import dataclass
from math import floor

from agents.schemas import AgentOutput, Direction, Instrument, Pick
from config.risk import RiskConfig


@dataclass
class PlannedOrder:
    ticker: str
    persona_id: str
    qty: int
    est_price: float  # share price (equity) or premium per share (option)
    weight_pct: float
    conviction: float
    stop_pct: float | None
    take_pct: float | None
    entry_type: str
    direction: str = "long"
    instrument: str = "equity"
    contract_symbol: str | None = None
    multiplier: int = 1  # 100 for options
    option_right: str | None = None


def _exec_key(pick: Pick) -> str:
    """The symbol we price/dedup/hold on: contract for options, ticker for equities."""
    return pick.contract_symbol if pick.instrument == Instrument.OPTION else pick.ticker


def _multiplier(pick: Pick) -> int:
    return 100 if pick.instrument == Instrument.OPTION else 1


def plan_orders(
    agent_outputs: list[AgentOutput],
    *,
    equity: float,
    buying_power: float,
    prices: dict[str, float],
    held_symbols: set[str],
    risk: RiskConfig,
    persona_caps: dict[str, float],
) -> tuple[list[PlannedOrder], list[str]]:
    notes: list[str] = []

    # 1. collect candidates (long-only unless shorts allowed; options are long-only here)
    #    + options quality gate: the mandate is a HIGH-PROBABILITY options trade, so option
    #    picks below the conviction floor never reach sizing.
    candidates: list[tuple[str, Pick]] = []
    for out in agent_outputs:
        for pick in out.picks:
            if pick.direction != Direction.LONG and not risk.allow_shorts:
                notes.append(f"drop {pick.ticker} ({out.persona_id}): shorts disabled")
                continue
            if pick.instrument == Instrument.OPTION and pick.direction != Direction.LONG:
                notes.append(f"drop {pick.ticker} ({out.persona_id}): only long options supported")
                continue
            if pick.instrument == Instrument.OPTION and pick.conviction < risk.min_option_conviction:
                notes.append(
                    f"drop {pick.ticker} option ({out.persona_id}): conviction "
                    f"{pick.conviction:.2f} < floor {risk.min_option_conviction}"
                )
                continue
            candidates.append((out.persona_id, pick))

    # 2. dedup across agents — highest conviction wins the exec key
    best: dict[str, tuple[str, Pick]] = {}
    for pid, pick in candidates:
        k = _exec_key(pick)
        cur = best.get(k)
        if cur is None or pick.conviction > cur[1].conviction:
            if cur is not None:
                notes.append(f"dedup {pick.ticker}: {pid} outranks {cur[0]}")
            best[k] = (pid, pick)

    # 3. per-agent trade cap — keep each persona's top-conviction picks, but the persona's
    #    best option pick survives the cap (options-trade mandate)
    by_persona: dict[str, list[Pick]] = {}
    for pid, pick in best.values():
        by_persona.setdefault(pid, []).append(pick)
    kept: list[tuple[str, Pick]] = []
    for pid, picks in by_persona.items():
        picks.sort(key=lambda p: p.conviction, reverse=True)
        keep = picks[: risk.max_trades_per_agent]
        if len(picks) > risk.max_trades_per_agent:
            notes.append(f"{pid}: capped to {risk.max_trades_per_agent} trades")
            options = [p for p in picks if p.instrument == Instrument.OPTION]
            if options and not any(p.instrument == Instrument.OPTION for p in keep):
                notes.append(f"{pid}: reserved option slot for {options[0].ticker}")
                keep[-1] = options[0]
        kept.extend((pid, p) for p in keep)

    # 4. price / liquidity filter + skip already-held (min_price only applies to equities)
    priced: list[tuple[str, Pick, float]] = []
    for pid, pick in kept:
        key = _exec_key(pick)
        px = prices.get(key)
        if not px or px <= 0:
            notes.append(f"drop {pick.ticker}: no price")
        elif pick.instrument == Instrument.EQUITY and px < risk.min_price:
            notes.append(f"drop {pick.ticker}: price {px:.2f} < min {risk.min_price}")
        elif pick.instrument == Instrument.EQUITY and risk.max_price and px > risk.max_price:
            notes.append(f"drop {pick.ticker}: price {px:.2f} > max {risk.max_price}")
        elif key in held_symbols:
            notes.append(f"skip {pick.ticker}: already held")
        else:
            priced.append((pid, pick, px))

    # 5. clamp weight, size to whole shares/contracts (options additionally premium-capped:
    #    weight = premium notional, i.e. the max loss — keep any single bet survivable)
    orders: list[PlannedOrder] = []
    for pid, pick, px in priced:
        cap = min(
            persona_caps.get(pid, risk.max_position_pct),
            risk.max_position_pct,
            risk.max_symbol_exposure_pct,
        )
        if pick.instrument == Instrument.OPTION:
            cap = min(cap, risk.max_option_weight_pct)
        weight = min(pick.requested_weight_pct, cap)
        mult = _multiplier(pick)
        qty = floor(weight * equity / (px * mult))
        if qty < 1:
            notes.append(f"drop {pick.ticker}: qty<1 at weight {weight:.3f}")
            continue
        orders.append(
            PlannedOrder(
                ticker=pick.ticker, persona_id=pid, qty=qty, est_price=px,
                weight_pct=weight, conviction=pick.conviction,
                stop_pct=pick.stop_loss_pct, take_pct=pick.take_profit_pct,
                entry_type=pick.entry_type.value, direction=pick.direction.value,
                instrument=pick.instrument.value, contract_symbol=pick.contract_symbol,
                multiplier=mult, option_right=pick.option_right,
            )
        )

    # the mandated trade: highest-conviction option that made it through sizing is protected
    # from the day cap and the buying-power trim below
    reserved = max(
        (o for o in orders if o.instrument == "option"), key=lambda o: o.conviction, default=None
    )

    # 6. per-day cap — rank all by conviction, truncate (reserved option swapped back in)
    orders.sort(key=lambda o: o.conviction, reverse=True)
    if len(orders) > risk.max_trades_per_day:
        keep, cut = orders[: risk.max_trades_per_day], orders[risk.max_trades_per_day :]
        if reserved is not None and reserved in cut:
            swapped = keep.pop()
            cut.remove(reserved)
            keep.append(reserved)
            notes.append(f"day cap: reserved option {reserved.ticker} replaces {swapped.ticker}")
        notes.append(f"day cap: {len(orders)} -> {risk.max_trades_per_day}")
        orders = keep

    # 7. buying-power trim — drop lowest-conviction until it fits (notional = qty*price*mult);
    #    the reserved option goes last (it's also the cheapest kind of trade to keep)
    def _notional(o: PlannedOrder) -> float:
        return o.qty * o.est_price * o.multiplier

    total = sum(_notional(o) for o in orders)
    while orders and total > buying_power:
        droppable = [o for o in orders if o is not reserved] or orders
        dropped = min(droppable, key=lambda o: o.conviction)
        orders.remove(dropped)
        total -= _notional(dropped)
        notes.append(f"drop {dropped.ticker}: buying power (need ${total:,.0f} > ${buying_power:,.0f})")

    return orders, notes