Janady13/missionmatch / internal/modules/ai/vectorguard.go

1package ai 2 3import ( 4 "context" 5 "errors" 6 "fmt" 7 "math" 8 "sync" 9 "sync/atomic" 10 11 "github.com/jackc/pgx/v5/pgxpool" 12 "github.com/pgvector/pgvector-go" 13 14 "github.com/missionworkforce/missionmatch/internal/platform/embeddings" 15 "github.com/missionworkforce/missionmatch/internal/platform/inference" 16 "github.com/missionworkforce/missionmatch/internal/platform/tenant" 17) 18 19// ErrVectorModelMismatch means the stored vectors were not produced by the configured embed model (or 20// the vector columns are the wrong width, or the store is not backfilled). Mixing embed models silently 21// mis ranks every retrieval with no error, so the brain MUST refuse to serve until the store is re 22// embedded with the configured model. The handler maps this to a typed 503 vector_model_mismatch. 23var ErrVectorModelMismatch = errors.New("ai: stored vectors were not produced by the configured embed model") 24 25// vectorMatchTolerance is the minimum cosine between a stored skill vector and a fresh embedding of the 26// same skill name by the configured model. The same text under the same model, normalized, lands at 27// cosine ~1.0; a DIFFERENT model (the all-minilm to bge corruption mode) lands far lower (the synonym 28// and unrelated calibration showed a different meaning, never mind a different model, sits at 0.4 to 29// 0.9), so 0.98 cleanly separates same model from a model swap while tolerating float drift. 30const vectorMatchTolerance = 0.98 31 32// VectorGate is the enforced re embed gate. It holds the boot time verdict on whether the stored vectors 33// were produced by the configured embed model, and /ai/query consults it before serving. A genuine model 34// mismatch closes the gate permanently (until a re embed and restart); a transient failure at boot (the 35// plane was briefly unreachable) leaves the gate to self heal on the first query after the plane recovers, 36// so a momentary outage does not dark the brain forever. The rest of the API never consults this gate. 37type VectorGate struct { 38 ok atomic.Bool // fast path once verified, no lock 39 mu sync.Mutex 40 mismatch bool // a permanent negative verdict was reached 41 reason string // the last failure detail, for the log and the 503 message 42 verify func(ctx context.Context) error 43} 44 45// NewVectorGate builds a gate over a verification closure. A nil closure means no verifier is configured 46// (the dev or test posture); such a gate is always open, so existing call sites that do not wire the 47// check keep serving. 48func NewVectorGate(verify func(ctx context.Context) error) *VectorGate { 49 return &VectorGate{verify: verify} 50} 51 52// Check returns nil when the brain may serve, or a typed error when it may not: ErrVectorModelMismatch 53// (a model mismatch or wrong width or unbackfilled store, a permanent closed gate) or an 54// inference.ErrPlaneUnavailable wrap (the check could not reach the plane yet, retried on the next call). 55func (g *VectorGate) Check(ctx context.Context) error { 56 if g.ok.Load() { 57 return nil 58 } 59 g.mu.Lock() 60 defer g.mu.Unlock() 61 if g.ok.Load() { 62 return nil 63 } 64 if g.mismatch { 65 return fmt.Errorf("%w: %s", ErrVectorModelMismatch, g.reason) 66 } 67 if g.verify == nil { 68 g.ok.Store(true) 69 return nil 70 } 71 err := g.verify(ctx) 72 if err == nil { 73 g.ok.Store(true) 74 return nil 75 } 76 if errors.Is(err, ErrVectorModelMismatch) { 77 g.mismatch = true 78 g.reason = err.Error() 79 return err 80 } 81 // Transient (the plane was unreachable for the re embed): stay unverified and surface a plane outage, 82 // so the caller gets a typed 503 and the gate retries on the next query rather than darkening forever. 83 g.reason = err.Error() 84 return fmt.Errorf("%w: stored vector model not yet verified: %v", inference.ErrPlaneUnavailable, err) 85} 86 87// VerifyVectorModel is the boot self check. It asserts the vector column dimension for skills and 88// candidates and confirms the stored skill vectors were produced by the configured embed model, by re 89// embedding a sample of stored skill NAMES (the embedded text for a skill is its name, so it is 90// re embeddable; a candidate's resume text is CUI and is never re embedded here, its column dimension is 91// checked through the catalog instead). The skills table is global and shares the one backfill model with 92// candidates, so a match here proves the configured model is the store's model. 93func VerifyVectorModel(ctx context.Context, db *tenant.DB, client *inference.Client, dim int) error { 94 if db == nil || db.Pool == nil { 95 return errors.New("ai: vector model check needs a database") 96 } 97 if client == nil { 98 return errors.New("ai: vector model check needs an inference client") 99 } 100 // Dimension contract through the catalog (the app role sees no candidate rows without a tenant, and 101 // the column type carries the width), so a wrong width column is caught without reading any PII. 102 if err := assertVectorColumnDim(ctx, db.Pool, "candidates", "resume_embedding", dim); err != nil { 103 return err 104 } 105 if err := assertVectorColumnDim(ctx, db.Pool, "skills", "embedding", dim); err != nil { 106 return err 107 } 108 // Model match on the global skills table. Re embed the SAME text the backfill embedded (the skill 109 // name with its category), via the shared embeddings.SkillEmbedText, or a categorized skill would 110 // look like a model mismatch when only the text differs. 111 var texts []string 112 var stored [][]float32 113 rows, err := db.Pool.Query(ctx, "SELECT name, COALESCE(category,''), embedding FROM skills WHERE embedding IS NOT NULL ORDER BY id LIMIT 5") 114 if err != nil { 115 return fmt.Errorf("ai: sample stored skill vectors: %w", err) 116 } 117 defer rows.Close() 118 for rows.Next() { 119 var name, category string 120 var v pgvector.Vector 121 if err := rows.Scan(&name, &category, &v); err != nil { 122 return err 123 } 124 texts = append(texts, embeddings.SkillEmbedText(name, category)) 125 stored = append(stored, v.Slice()) 126 } 127 if err := rows.Err(); err != nil { 128 return err 129 } 130 if len(texts) == 0 { 131 return fmt.Errorf("%w: no backfilled skill embeddings to verify against (run the bge backfill before serving the brain)", ErrVectorModelMismatch) 132 } 133 res, err := client.Embed(ctx, texts) 134 if err != nil { 135 // Could not reach the plane to re embed: a transient condition, not a proven mismatch. 136 return fmt.Errorf("ai: re embed for the vector model check: %w", err) 137 } 138 return vectorsMatchModel(texts, stored, res.Embeddings, dim) 139} 140 141// vectorsMatchModel is the pure comparison core (no DB, no plane), so the match and mismatch decision is 142// unit testable. Every stored vector must be dim wide and cosine close to the fresh embedding of its name. 143func vectorsMatchModel(names []string, stored, fresh [][]float32, dim int) error { 144 if len(fresh) != len(stored) { 145 return fmt.Errorf("ai: re embed returned %d vectors for %d stored", len(fresh), len(stored)) 146 } 147 for i := range stored { 148 if len(stored[i]) != dim { 149 return fmt.Errorf("%w: stored skill %q vector width %d, expected %d", ErrVectorModelMismatch, names[i], len(stored[i]), dim) 150 } 151 if len(fresh[i]) != dim { 152 return fmt.Errorf("ai: fresh embedding width %d, expected %d (plane or config error)", len(fresh[i]), dim) 153 } 154 c := cosineSim(stored[i], fresh[i]) 155 if c < vectorMatchTolerance { 156 return fmt.Errorf("%w: stored vector for skill %q matches the configured model at cosine %.4f, below tolerance %.2f (the store was embedded by a different model, re embed before serving)", ErrVectorModelMismatch, names[i], c, vectorMatchTolerance) 157 } 158 } 159 return nil 160} 161 162// assertVectorColumnDim reads the pgvector column dimension from the catalog (pgvector stores the 163// dimension directly in atttypmod, with no offset) and asserts it equals dim. 164func assertVectorColumnDim(ctx context.Context, pool *pgxpool.Pool, table, column string, dim int) error { 165 var typmod int 166 err := pool.QueryRow(ctx, 167 `SELECT a.atttypmod FROM pg_attribute a 168 JOIN pg_class c ON c.oid = a.attrelid 169 JOIN pg_namespace n ON n.oid = c.relnamespace 170 WHERE n.nspname = 'public' AND c.relname = $1 AND a.attname = $2 AND NOT a.attisdropped`, 171 table, column).Scan(&typmod) 172 if err != nil { 173 return fmt.Errorf("ai: read %s.%s vector dimension: %w", table, column, err) 174 } 175 if typmod != dim { 176 return fmt.Errorf("%w: column %s.%s is vector(%d), expected vector(%d)", ErrVectorModelMismatch, table, column, typmod, dim) 177 } 178 return nil 179} 180 181// cosineSim computes the true cosine (it divides by both norms), so it is correct whether or not the 182// stored vector was normalized at backfill. 183func cosineSim(a, b []float32) float64 { 184 var dot, na, nb float64 185 for i := range a { 186 dot += float64(a[i]) * float64(b[i]) 187 na += float64(a[i]) * float64(a[i]) 188 nb += float64(b[i]) * float64(b[i]) 189 } 190 if na == 0 || nb == 0 { 191 return 0 192 } 193 return dot / (math.Sqrt(na) * math.Sqrt(nb)) 194} 195