Janady13/missionmatch / internal/platform/magstore/magstore.go

1// Package magstore is the database glue between the MAG coefficient store (aedt_learned_models, 2// aedt_model_coefficients, aedt_calibration_bins) plus the bid_outcomes label ledger and the offline trainer 3// and evaluator (cmd/train-model, cmd/eval-model). It reads the pooled labeled bids, loads and scores a stored 4// model version, writes a new CHALLENGER (never a champion), and, as a separate call, promotes a version to 5// champion. The numerical fit itself lives in internal/platform/learn (pure, no DB); this package is only the 6// DB reads and writes, run under the BYPASSRLS migrator role. 7// 8// float64 lives ONLY in the offline solver and this offline scorer (an offline evaluation, never a request 9// path score); every coefficient WriteChallenger stores is re quantized to shopspring/decimal at the store's 10// numeric scale, so the SERVED path (internal/modules/ai) stays decimal exact. The served numbers are computed 11// by the ai package in decimal; this offline float scoring is used only to compare a challenger against the 12// champion on held out labels, which is a relative metric, not a surfaced number. 13package magstore 14 15import ( 16 "context" 17 "encoding/json" 18 "fmt" 19 "math" 20 21 "github.com/jackc/pgx/v5" 22 "github.com/jackc/pgx/v5/pgxpool" 23 "github.com/shopspring/decimal" 24 25 "github.com/missionworkforce/missionmatch/internal/platform/learn" 26) 27 28// SegmentBaseRateFeature is the one market structure feature the win probability trainer fits and the serving 29// path gathers: the win rate of a bid's customer agency segment. It is never a protected attribute or a proxy. 30const SegmentBaseRateFeature = "segment_base_rate" 31 32// interceptFeature is the reserved coefficient row carrying the model bias term (the log odds base rate), 33// matching the serving path's reserved name so a stored model round trips identically. 34const interceptFeature = "intercept" 35 36const ( 37 coefficientScale int32 = 10 // aedt_model_coefficients.coefficient is numeric(20,10) 38 calibrationScale int32 = 8 // aedt_calibration_bins figures are numeric(12,8) 39) 40 41// Model is a loaded model version: its intercept, its raw feature space coefficients keyed by feature name, and 42// its monotone calibration bins. Every figure is a float64 for the offline scorer; the stored source is decimal 43// and is read back exactly, then converted once here. 44type Model struct { 45 ModelKey string 46 Version int 47 OutputKind string 48 Intercept float64 49 Coeffs map[string]float64 50 Bins []Bin 51} 52 53// Bin is one calibration bin: a raw probability range [Lo, Hi] mapping to an EmpiricalRate. 54type Bin struct { 55 Lo, Hi, EmpiricalRate float64 56} 57 58// ReadLabeledBids reads every decided bid (win or loss) POOLED across tenants with its segment base rate 59// feature, the win rate of the bid's customer agency segment over the decided ledger, the SAME market base rate 60// the serving path gathers. Requirements are joined by the composite key (organization_id, id) so a cross 61// tenant join can never happen even under the BYPASSRLS migrator role. Ordered deterministically. 62func ReadLabeledBids(ctx context.Context, pool *pgxpool.Pool) ([]learn.LabeledRow, error) { 63 const q = ` 64WITH decided AS ( 65 SELECT b.id, 66 b.result, 67 COALESCE(NULLIF(r.customer_agency, ''), '(unspecified)') AS segment 68 FROM bid_outcomes b 69 LEFT JOIN requirements r 70 ON r.organization_id = b.organization_id AND r.id = b.requirement_id 71 WHERE b.result IN ('win', 'loss') 72), 73seg AS ( 74 SELECT segment, 75 (count(*) FILTER (WHERE result = 'win'))::numeric / NULLIF(count(*), 0) AS win_rate 76 FROM decided 77 GROUP BY segment 78) 79SELECT d.result, 80 COALESCE(s.win_rate, 0)::float8 AS segment_base_rate 81 FROM decided d 82 JOIN seg s ON s.segment = d.segment 83 ORDER BY d.id` 84 rows, err := pool.Query(ctx, q) 85 if err != nil { 86 return nil, err 87 } 88 defer rows.Close() 89 var out []learn.LabeledRow 90 for rows.Next() { 91 var result string 92 var rate float64 93 if err := rows.Scan(&result, &rate); err != nil { 94 return nil, err 95 } 96 out = append(out, learn.LabeledRow{ 97 Result: result, 98 Features: map[string]float64{SegmentBaseRateFeature: rate}, 99 }) 100 } 101 return out, rows.Err() 102} 103 104// ChampionVersion returns the champion version for a model_key, or an error when no champion is resolvable so an 105// evaluator never silently compares against nothing. 106func ChampionVersion(ctx context.Context, pool *pgxpool.Pool, modelKey string) (int, error) { 107 var v int 108 err := pool.QueryRow(ctx, 109 `SELECT version FROM aedt_learned_models WHERE model_key = $1 AND champion = true`, modelKey).Scan(&v) 110 if err != nil { 111 if err == pgx.ErrNoRows { 112 return 0, fmt.Errorf("magstore: no champion for %q", modelKey) 113 } 114 return 0, err 115 } 116 return v, nil 117} 118 119// LoadModel loads a specific model version: its output kind, its coefficients (the intercept plus every real 120// feature), and its calibration bins. It reads the decimal store exactly, then converts each figure to float64 121// once for the offline scorer. 122func LoadModel(ctx context.Context, pool *pgxpool.Pool, modelKey string, version int) (Model, error) { 123 m := Model{ModelKey: modelKey, Version: version, Coeffs: map[string]float64{}} 124 125 if err := pool.QueryRow(ctx, 126 `SELECT output_kind FROM aedt_learned_models WHERE model_key = $1 AND version = $2`, 127 modelKey, version).Scan(&m.OutputKind); err != nil { 128 if err == pgx.ErrNoRows { 129 return Model{}, fmt.Errorf("magstore: no model %q v%d", modelKey, version) 130 } 131 return Model{}, err 132 } 133 134 crows, err := pool.Query(ctx, 135 `SELECT feature_name, coefficient FROM aedt_model_coefficients WHERE model_key = $1 AND version = $2`, 136 modelKey, version) 137 if err != nil { 138 return Model{}, err 139 } 140 defer crows.Close() 141 for crows.Next() { 142 var name string 143 var coef decimal.Decimal 144 if err := crows.Scan(&name, &coef); err != nil { 145 return Model{}, err 146 } 147 f, _ := coef.Float64() 148 if name == interceptFeature { 149 m.Intercept = f 150 continue 151 } 152 m.Coeffs[name] = f 153 } 154 if err := crows.Err(); err != nil { 155 return Model{}, err 156 } 157 158 brows, err := pool.Query(ctx, 159 `SELECT bin_lo, bin_hi, empirical_rate FROM aedt_calibration_bins WHERE model_key = $1 AND version = $2 ORDER BY bin_index`, 160 modelKey, version) 161 if err != nil { 162 return Model{}, err 163 } 164 defer brows.Close() 165 for brows.Next() { 166 var lo, hi, rate decimal.Decimal 167 if err := brows.Scan(&lo, &hi, &rate); err != nil { 168 return Model{}, err 169 } 170 lf, _ := lo.Float64() 171 hf, _ := hi.Float64() 172 rf, _ := rate.Float64() 173 m.Bins = append(m.Bins, Bin{Lo: lf, Hi: hf, EmpiricalRate: rf}) 174 } 175 return m, brows.Err() 176} 177 178// Score computes the offline calibrated probability of a feature vector under a loaded model: raw = 179// sigmoid(intercept + SUM coef*value), then the monotone calibration remap (an empty bin set is the identity 180// map). It mirrors the serving path's math for the held out comparison; the served numbers themselves are 181// computed in decimal by the ai package. 182func (m Model) Score(features map[string]float64) float64 { 183 eta := m.Intercept 184 for name, v := range features { 185 eta += m.Coeffs[name] * v 186 } 187 raw := 1.0 / (1.0 + math.Exp(-eta)) 188 for _, b := range m.Bins { 189 if raw >= b.Lo && raw <= b.Hi { 190 raw = b.EmpiricalRate 191 break 192 } 193 } 194 if raw < 0 { 195 return 0 196 } 197 if raw > 1 { 198 return 1 199 } 200 return raw 201} 202 203// FeatureSpec builds the feature_spec jsonb value recording each feature, the Registry source intent it reads 204// from, and its TRAINING MEAN, mirroring the v0 seed shape so a challenger is spec compatible with the prior 205// only champion. The stored mean is what the serve path imputes for a feature with no history (a segment with 206// no decided bids), so a learned model returns the base rate for a no history segment instead of extrapolating 207// off a misleading zero. means may be nil (the prior only seed carries none) or aligned to featureNames. 208func FeatureSpec(featureNames []string, means []float64) map[string]any { 209 feats := make([]map[string]any, 0, len(featureNames)) 210 for i, name := range featureNames { 211 f := map[string]any{"name": name} 212 if name == SegmentBaseRateFeature { 213 f["source_intent"] = "bid_win_rate_by_segment" 214 f["column"] = "win_rate" 215 } 216 if i < len(means) { 217 f["mean"] = quant(means[i], coefficientScale).String() 218 } 219 feats = append(feats, f) 220 } 221 return map[string]any{"features": feats} 222} 223 224// Challenger is a fitted model ready to persist: its raw feature space intercept and coefficients (keyed by 225// feature name), its calibration bins, and the training summary for the notes column. 226type Challenger struct { 227 Family string // the model family, e.g. logistic or pairwise_logistic 228 OutputKind string // one of the legal MAG kinds: opportunity_score, display_order, confidence 229 FeatureNames []string 230 Intercept float64 231 Coeffs []float64 // aligned to FeatureNames 232 Means []float64 // training feature means, aligned to FeatureNames, stored for no history imputation 233 Bins []Bin 234 Notes string 235} 236 237// WriteChallenger writes a new challenger version in ONE migrator transaction: the next version number, its 238// coefficient rows (decimal), and its calibration bins. It sets champion=false ALWAYS; a champion flip is a 239// separate governance action (PromoteChampion), never a side effect of a fit. Versions are append only, an 240// audit trail of every fit. 241func WriteChallenger(ctx context.Context, pool *pgxpool.Pool, modelKey string, ch Challenger) (int, error) { 242 var version int 243 err := pgx.BeginFunc(ctx, pool, func(tx pgx.Tx) error { 244 if err := tx.QueryRow(ctx, 245 `SELECT COALESCE(MAX(version), 0) + 1 FROM aedt_learned_models WHERE model_key = $1`, modelKey).Scan(&version); err != nil { 246 return err 247 } 248 spec, err := json.Marshal(FeatureSpec(ch.FeatureNames, ch.Means)) 249 if err != nil { 250 return err 251 } 252 family := ch.Family 253 if family == "" { 254 family = "logistic" 255 } 256 outputKind := ch.OutputKind 257 if outputKind == "" { 258 outputKind = "opportunity_score" 259 } 260 if _, err := tx.Exec(ctx, 261 `INSERT INTO aedt_learned_models (model_key, version, family, output_kind, feature_spec, champion, notes) 262 VALUES ($1, $2, $3, $4, $5, false, $6)`, 263 modelKey, version, family, outputKind, spec, ch.Notes); err != nil { 264 return err 265 } 266 if _, err := tx.Exec(ctx, 267 `INSERT INTO aedt_model_coefficients (model_key, version, feature_name, coefficient) VALUES ($1, $2, 'intercept', $3)`, 268 modelKey, version, quant(ch.Intercept, coefficientScale)); err != nil { 269 return err 270 } 271 for j, name := range ch.FeatureNames { 272 if _, err := tx.Exec(ctx, 273 `INSERT INTO aedt_model_coefficients (model_key, version, feature_name, coefficient) VALUES ($1, $2, $3, $4)`, 274 modelKey, version, name, quant(ch.Coeffs[j], coefficientScale)); err != nil { 275 return err 276 } 277 } 278 for i, b := range ch.Bins { 279 if _, err := tx.Exec(ctx, 280 `INSERT INTO aedt_calibration_bins (model_key, version, bin_index, bin_lo, bin_hi, empirical_rate) 281 VALUES ($1, $2, $3, $4, $5, $6)`, 282 modelKey, version, i, quant(b.Lo, calibrationScale), quant(b.Hi, calibrationScale), quant(b.EmpiricalRate, calibrationScale)); err != nil { 283 return err 284 } 285 } 286 return nil 287 }) 288 return version, err 289} 290 291// PromoteChampion flips the champion to a challenger version in ONE migrator transaction: it clears the current 292// champion flag for the model_key, sets it on the target version, and appends one governance audit row, all 293// atomically. The partial unique index enforces at most one champion, so the flip is atomic and a second 294// champion is impossible. It errors if the target version does not exist. This is the ONLY place a champion 295// changes, and cmd/eval-model calls it only after the held out and fairness gates pass and a human passes the 296// promote flag. The audit row carries the model_key and the from and to versions, no PII and no org (a global 297// governance action). It returns the version the promotion replaced (0 when there was no prior champion). 298func PromoteChampion(ctx context.Context, pool *pgxpool.Pool, modelKey string, version int) (fromVersion int, err error) { 299 err = pgx.BeginFunc(ctx, pool, func(tx pgx.Tx) error { 300 var exists bool 301 if err := tx.QueryRow(ctx, 302 `SELECT EXISTS(SELECT 1 FROM aedt_learned_models WHERE model_key = $1 AND version = $2)`, 303 modelKey, version).Scan(&exists); err != nil { 304 return err 305 } 306 if !exists { 307 return fmt.Errorf("magstore: cannot promote missing model %q v%d", modelKey, version) 308 } 309 // The version being replaced (0 when there is no current champion), for the audit trail. 310 if err := tx.QueryRow(ctx, 311 `SELECT COALESCE(MAX(version), 0) FROM aedt_learned_models WHERE model_key = $1 AND champion = true`, 312 modelKey).Scan(&fromVersion); err != nil { 313 return err 314 } 315 if _, err := tx.Exec(ctx, 316 `UPDATE aedt_learned_models SET champion = false WHERE model_key = $1 AND champion = true`, modelKey); err != nil { 317 return err 318 } 319 if _, err := tx.Exec(ctx, 320 `UPDATE aedt_learned_models SET champion = true WHERE model_key = $1 AND version = $2`, modelKey, version); err != nil { 321 return err 322 } 323 newValues, merr := json.Marshal(map[string]any{"from_version": fromVersion, "to_version": version}) 324 if merr != nil { 325 return merr 326 } 327 if _, err := tx.Exec(ctx, 328 `INSERT INTO audit_logs (organization_id, actor_user_id, action, entity_type, entity_id, new_values) 329 VALUES (NULL, NULL, 'aedt.model.promoted', 'aedt_learned_model', $1, $2)`, 330 modelKey, newValues); err != nil { 331 return err 332 } 333 return nil 334 }) 335 return fromVersion, err 336} 337 338// quant re quantizes an offline float64 to a shopspring/decimal at a fixed scale, the boundary where a number 339// leaves the float solver and becomes the decimal the store holds and the serve path reads. 340func quant(v float64, scale int32) decimal.Decimal { 341 return decimal.NewFromFloat(v).Round(scale) 342} 343