From 7fd48c776465c6c648303751cfeeebb0efc6122e Mon Sep 17 00:00:00 2001 From: agentson Date: Sun, 8 Feb 2026 02:06:16 +0900 Subject: [PATCH] feat: add strategy/playbook Pydantic models (issue #79) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Define data contracts for the proactive strategy system: - StockCondition: AND-combined condition fields (RSI, volume, price) - StockScenario: condition-action rules with stop loss/take profit - StockPlaybook: per-stock scenario collection - GlobalRule: portfolio-level rules (e.g. REDUCE_ALL on loss limit) - DayPlaybook: complete daily playbook per market with validation - CrossMarketContext: cross-market awareness (KR↔US) - ScenarioAction, MarketOutlook, PlaybookStatus enums 33 tests covering validation, serialization, edge cases. Co-Authored-By: Claude Opus 4.6 --- src/strategy/__init__.py | 0 src/strategy/models.py | 164 +++++++++++++++ tests/test_strategy_models.py | 366 ++++++++++++++++++++++++++++++++++ 3 files changed, 530 insertions(+) create mode 100644 src/strategy/__init__.py create mode 100644 src/strategy/models.py create mode 100644 tests/test_strategy_models.py diff --git a/src/strategy/__init__.py b/src/strategy/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/strategy/models.py b/src/strategy/models.py new file mode 100644 index 0000000..f860714 --- /dev/null +++ b/src/strategy/models.py @@ -0,0 +1,164 @@ +"""Pydantic models for pre-market scenario planning. + +Defines the data contracts for the proactive strategy system: +- AI generates DayPlaybook before market open (structured JSON scenarios) +- Local ScenarioEngine matches conditions during market hours (no API calls) +""" + +from __future__ import annotations + +from datetime import UTC, date, datetime +from enum import Enum + +from pydantic import BaseModel, Field, field_validator + + +class ScenarioAction(str, Enum): + """Actions that can be taken by scenarios.""" + + BUY = "BUY" + SELL = "SELL" + HOLD = "HOLD" + REDUCE_ALL = "REDUCE_ALL" + + +class MarketOutlook(str, Enum): + """AI's assessment of market direction.""" + + BULLISH = "bullish" + NEUTRAL_TO_BULLISH = "neutral_to_bullish" + NEUTRAL = "neutral" + NEUTRAL_TO_BEARISH = "neutral_to_bearish" + BEARISH = "bearish" + + +class PlaybookStatus(str, Enum): + """Lifecycle status of a playbook.""" + + PENDING = "pending" + READY = "ready" + FAILED = "failed" + EXPIRED = "expired" + + +class StockCondition(BaseModel): + """Condition fields for scenario matching (all optional, AND-combined). + + The ScenarioEngine evaluates all non-None fields as AND conditions. + A condition matches only if ALL specified fields are satisfied. + """ + + rsi_below: float | None = None + rsi_above: float | None = None + volume_ratio_above: float | None = None + volume_ratio_below: float | None = None + price_above: float | None = None + price_below: float | None = None + price_change_pct_above: float | None = None + price_change_pct_below: float | None = None + + def has_any_condition(self) -> bool: + """Check if at least one condition field is set.""" + return any( + v is not None + for v in ( + self.rsi_below, + self.rsi_above, + self.volume_ratio_above, + self.volume_ratio_below, + self.price_above, + self.price_below, + self.price_change_pct_above, + self.price_change_pct_below, + ) + ) + + +class StockScenario(BaseModel): + """A single condition-action rule for one stock.""" + + condition: StockCondition + action: ScenarioAction + confidence: int = Field(ge=0, le=100) + allocation_pct: float = Field(ge=0, le=100, default=10.0) + stop_loss_pct: float = Field(le=0, default=-2.0) + take_profit_pct: float = Field(ge=0, default=3.0) + rationale: str = "" + + +class StockPlaybook(BaseModel): + """All scenarios for a single stock (ordered by priority).""" + + stock_code: str + stock_name: str = "" + scenarios: list[StockScenario] = Field(min_length=1) + + +class GlobalRule(BaseModel): + """Portfolio-level rule (checked before stock-level scenarios).""" + + condition: str # e.g. "portfolio_pnl_pct < -2.0" + action: ScenarioAction + rationale: str = "" + + +class CrossMarketContext(BaseModel): + """Summary of another market's state for cross-market awareness.""" + + market: str # e.g. "US" or "KR" + date: str + total_pnl: float = 0.0 + win_rate: float = 0.0 + index_change_pct: float = 0.0 # e.g. KOSPI or S&P500 change + key_events: list[str] = Field(default_factory=list) + lessons: list[str] = Field(default_factory=list) + + +class DayPlaybook(BaseModel): + """Complete playbook for a single trading day in a single market. + + Generated by PreMarketPlanner (1 Gemini call per market per day). + Consumed by ScenarioEngine during market hours (0 API calls). + """ + + date: date + market: str # "KR" or "US" + market_outlook: MarketOutlook = MarketOutlook.NEUTRAL + generated_at: str = "" # ISO timestamp + gemini_model: str = "" + token_count: int = 0 + global_rules: list[GlobalRule] = Field(default_factory=list) + stock_playbooks: list[StockPlaybook] = Field(default_factory=list) + default_action: ScenarioAction = ScenarioAction.HOLD + context_summary: dict = Field(default_factory=dict) + cross_market: CrossMarketContext | None = None + + @field_validator("stock_playbooks") + @classmethod + def validate_unique_stocks(cls, v: list[StockPlaybook]) -> list[StockPlaybook]: + codes = [pb.stock_code for pb in v] + if len(codes) != len(set(codes)): + raise ValueError("Duplicate stock codes in playbook") + return v + + def get_stock_playbook(self, stock_code: str) -> StockPlaybook | None: + """Find the playbook for a specific stock.""" + for pb in self.stock_playbooks: + if pb.stock_code == stock_code: + return pb + return None + + @property + def scenario_count(self) -> int: + """Total number of scenarios across all stocks.""" + return sum(len(pb.scenarios) for pb in self.stock_playbooks) + + @property + def stock_count(self) -> int: + """Number of stocks with scenarios.""" + return len(self.stock_playbooks) + + def model_post_init(self, __context: object) -> None: + """Set generated_at if not provided.""" + if not self.generated_at: + self.generated_at = datetime.now(UTC).isoformat() diff --git a/tests/test_strategy_models.py b/tests/test_strategy_models.py new file mode 100644 index 0000000..9ea40e0 --- /dev/null +++ b/tests/test_strategy_models.py @@ -0,0 +1,366 @@ +"""Tests for strategy/playbook Pydantic models.""" + +from __future__ import annotations + +from datetime import date + +import pytest +from pydantic import ValidationError + +from src.strategy.models import ( + CrossMarketContext, + DayPlaybook, + GlobalRule, + MarketOutlook, + PlaybookStatus, + ScenarioAction, + StockCondition, + StockPlaybook, + StockScenario, +) + + +# --------------------------------------------------------------------------- +# StockCondition +# --------------------------------------------------------------------------- + + +class TestStockCondition: + def test_empty_condition(self) -> None: + cond = StockCondition() + assert not cond.has_any_condition() + + def test_single_field(self) -> None: + cond = StockCondition(rsi_below=30.0) + assert cond.has_any_condition() + + def test_multiple_fields(self) -> None: + cond = StockCondition(rsi_below=25.0, volume_ratio_above=3.0) + assert cond.has_any_condition() + + def test_all_fields(self) -> None: + cond = StockCondition( + rsi_below=30, + rsi_above=10, + volume_ratio_above=2.0, + volume_ratio_below=10.0, + price_above=1000, + price_below=50000, + price_change_pct_above=-5.0, + price_change_pct_below=5.0, + ) + assert cond.has_any_condition() + + +# --------------------------------------------------------------------------- +# StockScenario +# --------------------------------------------------------------------------- + + +class TestStockScenario: + def test_valid_scenario(self) -> None: + s = StockScenario( + condition=StockCondition(rsi_below=25.0), + action=ScenarioAction.BUY, + confidence=85, + allocation_pct=15.0, + stop_loss_pct=-2.0, + take_profit_pct=3.0, + rationale="Oversold bounce expected", + ) + assert s.action == ScenarioAction.BUY + assert s.confidence == 85 + + def test_confidence_too_high(self) -> None: + with pytest.raises(ValidationError): + StockScenario( + condition=StockCondition(), + action=ScenarioAction.BUY, + confidence=101, + ) + + def test_confidence_too_low(self) -> None: + with pytest.raises(ValidationError): + StockScenario( + condition=StockCondition(), + action=ScenarioAction.BUY, + confidence=-1, + ) + + def test_allocation_too_high(self) -> None: + with pytest.raises(ValidationError): + StockScenario( + condition=StockCondition(), + action=ScenarioAction.BUY, + confidence=80, + allocation_pct=101.0, + ) + + def test_stop_loss_must_be_negative(self) -> None: + with pytest.raises(ValidationError): + StockScenario( + condition=StockCondition(), + action=ScenarioAction.BUY, + confidence=80, + stop_loss_pct=1.0, + ) + + def test_take_profit_must_be_positive(self) -> None: + with pytest.raises(ValidationError): + StockScenario( + condition=StockCondition(), + action=ScenarioAction.BUY, + confidence=80, + take_profit_pct=-1.0, + ) + + def test_defaults(self) -> None: + s = StockScenario( + condition=StockCondition(), + action=ScenarioAction.HOLD, + confidence=50, + ) + assert s.allocation_pct == 10.0 + assert s.stop_loss_pct == -2.0 + assert s.take_profit_pct == 3.0 + assert s.rationale == "" + + +# --------------------------------------------------------------------------- +# StockPlaybook +# --------------------------------------------------------------------------- + + +class TestStockPlaybook: + def test_valid_playbook(self) -> None: + pb = StockPlaybook( + stock_code="005930", + stock_name="Samsung Electronics", + scenarios=[ + StockScenario( + condition=StockCondition(rsi_below=25.0), + action=ScenarioAction.BUY, + confidence=85, + ), + ], + ) + assert pb.stock_code == "005930" + assert len(pb.scenarios) == 1 + + def test_empty_scenarios_rejected(self) -> None: + with pytest.raises(ValidationError): + StockPlaybook( + stock_code="005930", + scenarios=[], + ) + + def test_multiple_scenarios(self) -> None: + pb = StockPlaybook( + stock_code="AAPL", + scenarios=[ + StockScenario( + condition=StockCondition(rsi_below=25.0), + action=ScenarioAction.BUY, + confidence=85, + ), + StockScenario( + condition=StockCondition(rsi_above=75.0), + action=ScenarioAction.SELL, + confidence=80, + ), + ], + ) + assert len(pb.scenarios) == 2 + + +# --------------------------------------------------------------------------- +# GlobalRule +# --------------------------------------------------------------------------- + + +class TestGlobalRule: + def test_valid_rule(self) -> None: + rule = GlobalRule( + condition="portfolio_pnl_pct < -2.0", + action=ScenarioAction.REDUCE_ALL, + rationale="Risk limit approaching", + ) + assert rule.action == ScenarioAction.REDUCE_ALL + + def test_hold_rule(self) -> None: + rule = GlobalRule( + condition="volatility_index > 30", + action=ScenarioAction.HOLD, + ) + assert rule.rationale == "" + + +# --------------------------------------------------------------------------- +# CrossMarketContext +# --------------------------------------------------------------------------- + + +class TestCrossMarketContext: + def test_valid_context(self) -> None: + ctx = CrossMarketContext( + market="US", + date="2026-02-07", + total_pnl=-1.5, + win_rate=40.0, + index_change_pct=-2.3, + key_events=["Fed rate decision"], + lessons=["Avoid tech sector on rate hike days"], + ) + assert ctx.market == "US" + assert len(ctx.key_events) == 1 + + def test_defaults(self) -> None: + ctx = CrossMarketContext(market="KR", date="2026-02-07") + assert ctx.total_pnl == 0.0 + assert ctx.key_events == [] + assert ctx.lessons == [] + + +# --------------------------------------------------------------------------- +# DayPlaybook +# --------------------------------------------------------------------------- + + +def _make_scenario(rsi_below: float = 25.0) -> StockScenario: + return StockScenario( + condition=StockCondition(rsi_below=rsi_below), + action=ScenarioAction.BUY, + confidence=85, + ) + + +def _make_playbook(**kwargs) -> DayPlaybook: + defaults = { + "date": date(2026, 2, 7), + "market": "KR", + "stock_playbooks": [ + StockPlaybook(stock_code="005930", scenarios=[_make_scenario()]), + ], + } + defaults.update(kwargs) + return DayPlaybook(**defaults) + + +class TestDayPlaybook: + def test_valid_playbook(self) -> None: + pb = _make_playbook() + assert pb.market == "KR" + assert pb.date == date(2026, 2, 7) + assert pb.default_action == ScenarioAction.HOLD + assert pb.scenario_count == 1 + assert pb.stock_count == 1 + + def test_generated_at_auto_set(self) -> None: + pb = _make_playbook() + assert pb.generated_at != "" + + def test_explicit_generated_at(self) -> None: + pb = _make_playbook(generated_at="2026-02-07T08:30:00") + assert pb.generated_at == "2026-02-07T08:30:00" + + def test_duplicate_stocks_rejected(self) -> None: + with pytest.raises(ValidationError): + DayPlaybook( + date=date(2026, 2, 7), + market="KR", + stock_playbooks=[ + StockPlaybook(stock_code="005930", scenarios=[_make_scenario()]), + StockPlaybook(stock_code="005930", scenarios=[_make_scenario(30)]), + ], + ) + + def test_empty_stock_playbooks_allowed(self) -> None: + pb = DayPlaybook( + date=date(2026, 2, 7), + market="KR", + stock_playbooks=[], + ) + assert pb.stock_count == 0 + assert pb.scenario_count == 0 + + def test_get_stock_playbook_found(self) -> None: + pb = _make_playbook() + result = pb.get_stock_playbook("005930") + assert result is not None + assert result.stock_code == "005930" + + def test_get_stock_playbook_not_found(self) -> None: + pb = _make_playbook() + result = pb.get_stock_playbook("AAPL") + assert result is None + + def test_with_global_rules(self) -> None: + pb = _make_playbook( + global_rules=[ + GlobalRule( + condition="portfolio_pnl_pct < -2.0", + action=ScenarioAction.REDUCE_ALL, + ), + ], + ) + assert len(pb.global_rules) == 1 + + def test_with_cross_market_context(self) -> None: + ctx = CrossMarketContext(market="US", date="2026-02-07", total_pnl=-1.5) + pb = _make_playbook(cross_market=ctx) + assert pb.cross_market is not None + assert pb.cross_market.market == "US" + + def test_market_outlook(self) -> None: + pb = _make_playbook(market_outlook=MarketOutlook.BEARISH) + assert pb.market_outlook == MarketOutlook.BEARISH + + def test_multiple_stocks_multiple_scenarios(self) -> None: + pb = DayPlaybook( + date=date(2026, 2, 7), + market="US", + stock_playbooks=[ + StockPlaybook( + stock_code="AAPL", + scenarios=[_make_scenario(), _make_scenario(30)], + ), + StockPlaybook( + stock_code="MSFT", + scenarios=[_make_scenario()], + ), + ], + ) + assert pb.stock_count == 2 + assert pb.scenario_count == 3 + + def test_serialization_roundtrip(self) -> None: + pb = _make_playbook( + market_outlook=MarketOutlook.BULLISH, + cross_market=CrossMarketContext(market="US", date="2026-02-07"), + ) + json_str = pb.model_dump_json() + restored = DayPlaybook.model_validate_json(json_str) + assert restored.market == pb.market + assert restored.date == pb.date + assert restored.scenario_count == pb.scenario_count + assert restored.cross_market is not None + + +# --------------------------------------------------------------------------- +# Enums +# --------------------------------------------------------------------------- + + +class TestEnums: + def test_scenario_action_values(self) -> None: + assert ScenarioAction.BUY.value == "BUY" + assert ScenarioAction.SELL.value == "SELL" + assert ScenarioAction.HOLD.value == "HOLD" + assert ScenarioAction.REDUCE_ALL.value == "REDUCE_ALL" + + def test_market_outlook_values(self) -> None: + assert len(MarketOutlook) == 5 + + def test_playbook_status_values(self) -> None: + assert PlaybookStatus.READY.value == "ready" + assert PlaybookStatus.EXPIRED.value == "expired"