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. 37// The verdict is held PER ORGANIZATION, which matters once tenants have their own databases. Each tenant 38// database carries its own vectors and is backfilled independently, so "verified" is a property of a 39// store, not of the process. A single shared verdict would let the first tenant that happened to be 40// backfilled correctly open the gate for a tenant whose store is stale, half re embedded, or on the old 41// model, which is exactly the silent mis ranking this gate exists to prevent. In shared mode every 42// request resolves to the same empty key, so there is one verdict and the behaviour is unchanged. 43type VectorGate struct { 44 mu sync.Mutex 45 states map[string]*gateState 46 verify func(ctx context.Context) error 47} 48 49// gateState is one store's verdict. ok is atomic so the hot path after verification takes no lock. 50type gateState struct { 51 ok atomic.Bool 52 mismatch bool // a permanent negative verdict was reached for this store 53 reason string // the last failure detail, for the log and the 503 message 54} 55 56// NewVectorGate builds a gate over a verification closure. A nil closure means no verifier is configured 57// (the dev or test posture); such a gate is always open, so existing call sites that do not wire the 58// check keep serving. 59func NewVectorGate(verify func(ctx context.Context) error) *VectorGate { 60 return &VectorGate{verify: verify, states: map[string]*gateState{}} 61} 62 63// stateFor returns the verdict holder for the caller's organization, creating it on first sight. An 64// unresolved organization uses the empty key, which is the shared mode case and also the boot self check. 65func (g *VectorGate) stateFor(ctx context.Context) *gateState { 66 key, _ := tenant.OrgFromContext(ctx) 67 g.mu.Lock() 68 defer g.mu.Unlock() 69 if g.states == nil { 70 g.states = map[string]*gateState{} 71 } 72 s, ok := g.states[key] 73 if !ok { 74 s = &gateState{} 75 g.states[key] = s 76 } 77 return s 78} 79 80// Check returns nil when the brain may serve THIS caller's store, or a typed error when it may not: 81// ErrVectorModelMismatch (a model mismatch or wrong width or unbackfilled store, a permanent closed gate 82// for that store) or an inference.ErrPlaneUnavailable wrap (the check could not reach the plane yet, 83// retried on the next call). 84func (g *VectorGate) Check(ctx context.Context) error { 85 s := g.stateFor(ctx) 86 if s.ok.Load() { 87 return nil 88 } 89 g.mu.Lock() 90 defer g.mu.Unlock() 91 if s.ok.Load() { 92 return nil 93 } 94 if s.mismatch { 95 return fmt.Errorf("%w: %s", ErrVectorModelMismatch, s.reason) 96 } 97 if g.verify == nil { 98 s.ok.Store(true) 99 return nil 100 } 101 err := g.verify(ctx) 102 if err == nil { 103 s.ok.Store(true) 104 return nil 105 } 106 if errors.Is(err, ErrVectorModelMismatch) { 107 s.mismatch = true 108 s.reason = err.Error() 109 return err 110 } 111 // Transient (the plane was unreachable for the re embed): stay unverified and surface a plane outage, 112 // so the caller gets a typed 503 and the gate retries on the next query rather than darkening forever. 113 s.reason = err.Error() 114 return fmt.Errorf("%w: stored vector model not yet verified: %v", inference.ErrPlaneUnavailable, err) 115} 116 117// VerifyVectorModel is the boot self check. It asserts the vector column dimension for skills and 118// candidates and confirms the stored skill vectors were produced by the configured embed model, by re 119// embedding a sample of stored skill NAMES (the embedded text for a skill is its name, so it is 120// re embeddable; a candidate's resume text is CUI and is never re embedded here, its column dimension is 121// checked through the catalog instead). The skills table is global and shares the one backfill model with 122// candidates, so a match here proves the configured model is the store's model. 123func VerifyVectorModel(ctx context.Context, db *tenant.DB, client *inference.Client, dim int) error { 124 if db == nil { 125 return errors.New("ai: vector model check needs a database") 126 } 127 if client == nil { 128 return errors.New("ai: vector model check needs an inference client") 129 } 130 // Acquire rather than reading the DB's own pool directly. Under per tenant routing the routing DB has no pool of 131 // its own and the vectors being verified live in the CALLER'S tenant database, so the check has to run 132 // against the same store the retrieval will read. In shared mode this is a no op wrapper over the one 133 // pool. At boot there is no organization in context, so this fails as a TRANSIENT error, which is 134 // already the designed path: the gate stays unverified and self heals on the first real query, where an 135 // organization does exist. 136 leaf, release, err := db.Acquire(ctx) 137 if err != nil { 138 return fmt.Errorf("ai: vector model check needs a tenant database: %w", err) 139 } 140 defer release() 141 if leaf.Pool == nil { 142 return errors.New("ai: vector model check needs a database") 143 } 144 pool := leaf.Pool 145 // Dimension contract through the catalog (the app role sees no candidate rows without a tenant, and 146 // the column type carries the width), so a wrong width column is caught without reading any PII. 147 if err := assertVectorColumnDim(ctx, pool, "candidates", "resume_embedding", dim); err != nil { 148 return err 149 } 150 if err := assertVectorColumnDim(ctx, pool, "skills", "embedding", dim); err != nil { 151 return err 152 } 153 // Model match on the global skills table. Re embed the SAME text the backfill embedded (the skill 154 // name with its category), via the shared embeddings.SkillEmbedText, or a categorized skill would 155 // look like a model mismatch when only the text differs. 156 var texts []string 157 var stored [][]float32 158 rows, err := pool.Query(ctx, "SELECT name, COALESCE(category,''), embedding FROM skills WHERE embedding IS NOT NULL ORDER BY id LIMIT 5") 159 if err != nil { 160 return fmt.Errorf("ai: sample stored skill vectors: %w", err) 161 } 162 defer rows.Close() 163 for rows.Next() { 164 var name, category string 165 var v pgvector.Vector 166 if err := rows.Scan(&name, &category, &v); err != nil { 167 return err 168 } 169 texts = append(texts, embeddings.SkillEmbedText(name, category)) 170 stored = append(stored, v.Slice()) 171 } 172 if err := rows.Err(); err != nil { 173 return err 174 } 175 if len(texts) == 0 { 176 return fmt.Errorf("%w: no backfilled skill embeddings to verify against (run the bge backfill before serving the brain)", ErrVectorModelMismatch) 177 } 178 res, err := client.Embed(ctx, texts) 179 if err != nil { 180 // Could not reach the plane to re embed: a transient condition, not a proven mismatch. 181 return fmt.Errorf("ai: re embed for the vector model check: %w", err) 182 } 183 return vectorsMatchModel(texts, stored, res.Embeddings, dim) 184} 185 186// vectorsMatchModel is the pure comparison core (no DB, no plane), so the match and mismatch decision is 187// unit testable. Every stored vector must be dim wide and cosine close to the fresh embedding of its name. 188func vectorsMatchModel(names []string, stored, fresh [][]float32, dim int) error { 189 if len(fresh) != len(stored) { 190 return fmt.Errorf("ai: re embed returned %d vectors for %d stored", len(fresh), len(stored)) 191 } 192 for i := range stored { 193 if len(stored[i]) != dim { 194 return fmt.Errorf("%w: stored skill %q vector width %d, expected %d", ErrVectorModelMismatch, names[i], len(stored[i]), dim) 195 } 196 if len(fresh[i]) != dim { 197 return fmt.Errorf("ai: fresh embedding width %d, expected %d (plane or config error)", len(fresh[i]), dim) 198 } 199 c := cosineSim(stored[i], fresh[i]) 200 if c < vectorMatchTolerance { 201 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) 202 } 203 } 204 return nil 205} 206 207// assertVectorColumnDim reads the pgvector column dimension from the catalog (pgvector stores the 208// dimension directly in atttypmod, with no offset) and asserts it equals dim. 209func assertVectorColumnDim(ctx context.Context, pool *pgxpool.Pool, table, column string, dim int) error { 210 var typmod int 211 err := pool.QueryRow(ctx, 212 `SELECT a.atttypmod FROM pg_attribute a 213 JOIN pg_class c ON c.oid = a.attrelid 214 JOIN pg_namespace n ON n.oid = c.relnamespace 215 WHERE n.nspname = 'public' AND c.relname = $1 AND a.attname = $2 AND NOT a.attisdropped`, 216 table, column).Scan(&typmod) 217 if err != nil { 218 return fmt.Errorf("ai: read %s.%s vector dimension: %w", table, column, err) 219 } 220 if typmod != dim { 221 return fmt.Errorf("%w: column %s.%s is vector(%d), expected vector(%d)", ErrVectorModelMismatch, table, column, typmod, dim) 222 } 223 return nil 224} 225 226// cosineSim computes the true cosine (it divides by both norms), so it is correct whether or not the 227// stored vector was normalized at backfill. 228func cosineSim(a, b []float32) float64 { 229 var dot, na, nb float64 230 for i := range a { 231 dot += float64(a[i]) * float64(b[i]) 232 na += float64(a[i]) * float64(a[i]) 233 nb += float64(b[i]) * float64(b[i]) 234 } 235 if na == 0 || nb == 0 { 236 return 0 237 } 238 return dot / (math.Sqrt(na) * math.Sqrt(nb)) 239} 240