From f41cf75d8a8892ff3c31a2db6abbeb87c1ab61b6 Mon Sep 17 00:00:00 2001 From: Aryan Singh Date: Wed, 1 Jul 2026 16:29:29 -0500 Subject: [PATCH] feat(agents): pydantic output schemas (Pick/AgentOutput/RunManifest) + tests Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_016DvdRZkyCy9NJpHYDdPqmS --- agents/__init__.py | 0 agents/schemas.py | 134 ++++++++++++++++++++++++++++++++++++++++++ tests/test_schemas.py | 73 +++++++++++++++++++++++ 3 files changed, 207 insertions(+) create mode 100644 agents/__init__.py create mode 100644 agents/schemas.py create mode 100644 tests/test_schemas.py diff --git a/agents/__init__.py b/agents/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/agents/schemas.py b/agents/schemas.py new file mode 100644 index 0000000..b66bd49 --- /dev/null +++ b/agents/schemas.py @@ -0,0 +1,134 @@ +"""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) diff --git a/tests/test_schemas.py b/tests/test_schemas.py new file mode 100644 index 0000000..e80e530 --- /dev/null +++ b/tests/test_schemas.py @@ -0,0 +1,73 @@ +"""Phase-0 gate: the output contract validates and rejects garbage.""" + +from __future__ import annotations + +from datetime import datetime, timezone + +import pytest +from pydantic import ValidationError + +from agents.schemas import AgentOutput, AgentStatus, Direction, Pick, RunManifest + + +def _pick(**kw) -> Pick: + base = dict(ticker="nvda", conviction=0.8, thesis="momentum", requested_weight_pct=0.05) + base.update(kw) + return Pick(**base) + + +def test_pick_valid_and_ticker_uppercased(): + p = _pick() + assert p.ticker == "NVDA" + assert p.direction == Direction.LONG + + +def test_conviction_out_of_range_rejected(): + with pytest.raises(ValidationError): + _pick(conviction=1.5) + + +def test_weight_must_be_positive_fraction(): + with pytest.raises(ValidationError): + _pick(requested_weight_pct=0.0) + with pytest.raises(ValidationError): + _pick(requested_weight_pct=1.5) + + +def test_stop_take_pct_must_be_fraction(): + with pytest.raises(ValidationError): + _pick(stop_loss_pct=7.0) # 700% — clearly a mistake + assert _pick(stop_loss_pct=0.07).stop_loss_pct == 0.07 + + +def test_long_absolute_levels_ordering_enforced(): + with pytest.raises(ValidationError): + _pick(entry_type="limit", limit_price=100.0, stop_price=105.0) # stop above entry + with pytest.raises(ValidationError): + _pick(entry_type="limit", limit_price=100.0, take_price=95.0) # take below entry + ok = _pick(entry_type="limit", limit_price=100.0, stop_price=93.0, take_price=120.0) + assert ok.limit_price == 100.0 + + +def test_agent_output_roundtrip(): + out = AgentOutput( + agent_name="Nova", + persona_id="nova", + started_at=datetime.now(timezone.utc), + status=AgentStatus.COMPLETE, + picks=[_pick(), _pick(ticker="amd", conviction=0.6)], + ) + restored = AgentOutput.model_validate_json(out.model_dump_json()) + assert restored.agent_name == "Nova" + assert [p.ticker for p in restored.picks] == ["NVDA", "AMD"] + + +def test_run_manifest_minimal(): + m = RunManifest( + run_id="r1", + trade_date="2026-07-01", + model="gpt-oss-20b", + started_at=datetime.now(timezone.utc), + ) + assert m.status == "running" + assert m.agents == []