feat(agents): pydantic output schemas (Pick/AgentOutput/RunManifest) + tests
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016DvdRZkyCy9NJpHYDdPqmS
This commit is contained in:
1 parent
031a051687
commit
f41cf75d8a
3 files changed
+207
No files matched your search
Whitespace-only changes.
@@ -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)
|
||||||
@@ -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 == []
|
||||||
Reference in new issue
Block a user