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) -> dedup across agents (highest conviction) ->
per-agent trade cap -> price/liquidity filter + skip already-held -> size & clamp weight ->
per-day cap -> buying-power trim.
"""

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)
    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
            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
    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)
        if len(picks) > risk.max_trades_per_agent:
            notes.append(f"{pid}: capped to {risk.max_trades_per_agent} trades")
        kept.extend((pid, p) for p in picks[: risk.max_trades_per_agent])

    # 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
    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,
        )
        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,
            )
        )

    # 6. per-day cap — rank all by conviction, truncate
    orders.sort(key=lambda o: o.conviction, reverse=True)
    if len(orders) > risk.max_trades_per_day:
        notes.append(f"day cap: {len(orders)} -> {risk.max_trades_per_day}")
        orders = orders[: risk.max_trades_per_day]

    # 7. buying-power trim — drop lowest-conviction until it fits (notional = qty*price*mult)
    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:
        dropped = orders.pop()
        total -= _notional(dropped)
        notes.append(f"drop {dropped.ticker}: buying power (need ${total:,.0f} > ${buying_power:,.0f})")

    return orders, notes