Janady13/missionmatch / cmd/embed-backfill/main.go

1// Command embed-backfill populates the embedding columns that the audit found dead: skills.embedding 2// (0 of 61) and candidates.resume_embedding (sparse). It composes a contents profile per row, embeds it 3// through the stateless MEGAMIND inference plane (384 dim, matching the vector(384) columns), L2 4// normalizes, and writes it back. After a run the HNSW indexes and the SearchSkillsByEmbedding / 5// SearchCandidatesHybrid queries become live. 6// 7// It targets ONE database (a tenant database under the database per tenant model, or the current shared 8// database). It writes only embeddings, never identifiers off box: text goes to the plane, vectors come 9// back. It writes a single append only AEDT audit summary at the end (counts and model only, no PII). 10// 11// Usage: 12// 13// embed-backfill -dsn $MM_MIGRATOR_DATABASE_URL -megamind $MM_MEGAMIND_URL -what all [-org <uuid>] [-only-missing] 14package main 15 16import ( 17 "context" 18 "encoding/json" 19 "flag" 20 "fmt" 21 "log" 22 "math" 23 "os" 24 "strconv" 25 "strings" 26 "time" 27 28 "github.com/jackc/pgx/v5/pgxpool" 29 30 "github.com/missionworkforce/missionmatch/internal/platform/embeddings" 31 "github.com/missionworkforce/missionmatch/internal/platform/inference" 32) 33 34func main() { 35 dsn := flag.String("dsn", os.Getenv("MM_MIGRATOR_DATABASE_URL"), "target database DSN (defaults to MM_MIGRATOR_DATABASE_URL)") 36 mega := flag.String("megamind", os.Getenv("MM_MEGAMIND_URL"), "MEGAMIND gateway base URL") 37 model := flag.String("model", envDefault("MM_MEGAMIND_EMBED_MODEL", "qllama/bge-small-en-v1.5"), "384 dim embed model") 38 what := flag.String("what", "all", "what to embed: skills | candidates | all") 39 org := flag.String("org", "", "optional organization_id filter (for the shared database)") 40 onlyMissing := flag.Bool("only-missing", false, "embed only rows whose embedding is currently NULL") 41 batch := flag.Int("batch", 32, "embedding batch size") 42 flag.Parse() 43 44 if *dsn == "" || *mega == "" { 45 log.Fatal("embed-backfill: -dsn and -megamind are required") 46 } 47 ctx := context.Background() 48 pool, err := pgxpool.New(ctx, *dsn) 49 if err != nil { 50 log.Fatalf("db: %v", err) 51 } 52 defer pool.Close() 53 54 client, err := inference.New(inference.Config{BaseURL: *mega, EmbedModel: *model, Dimension: 384, Timeout: 60 * time.Second, MaxRetries: 2}) 55 if err != nil { 56 log.Fatalf("inference: %v", err) 57 } 58 if err := client.Health(ctx); err != nil { 59 log.Fatalf("MEGAMIND not reachable at %s: %v", *mega, err) 60 } 61 62 var nSkills, nCands int 63 if *what == "skills" || *what == "all" { 64 if nSkills, err = backfillSkills(ctx, pool, client, *onlyMissing, *batch); err != nil { 65 log.Fatalf("skills: %v", err) 66 } 67 log.Printf("embed-backfill: embedded %d skills", nSkills) 68 } 69 if *what == "candidates" || *what == "all" { 70 if nCands, err = backfillCandidates(ctx, pool, client, *org, *onlyMissing, *batch); err != nil { 71 log.Fatalf("candidates: %v", err) 72 } 73 log.Printf("embed-backfill: embedded %d candidates", nCands) 74 } 75 writeAudit(ctx, pool, *org, *model, nSkills, nCands) 76 log.Printf("embed-backfill: DONE (skills=%d candidates=%d model=%s)", nSkills, nCands, *model) 77} 78 79func backfillSkills(ctx context.Context, pool *pgxpool.Pool, client *inference.Client, onlyMissing bool, batch int) (int, error) { 80 q := "SELECT id, name, COALESCE(category,'') FROM skills" 81 if onlyMissing { 82 q += " WHERE embedding IS NULL" 83 } 84 q += " ORDER BY id" 85 rows, err := pool.Query(ctx, q) 86 if err != nil { 87 return 0, err 88 } 89 type item struct { 90 id string 91 text string 92 } 93 var items []item 94 for rows.Next() { 95 var id, name, cat string 96 if err := rows.Scan(&id, &name, &cat); err != nil { 97 rows.Close() 98 return 0, err 99 } 100 items = append(items, item{id: id, text: embeddings.SkillEmbedText(name, cat)}) 101 } 102 rows.Close() 103 return embedAndWrite(ctx, pool, client, batch, len(items), 104 func(i int) (string, string) { return items[i].id, items[i].text }, 105 "UPDATE skills SET embedding = $1::vector WHERE id = $2") 106} 107 108func backfillCandidates(ctx context.Context, pool *pgxpool.Pool, client *inference.Client, org string, onlyMissing bool, batch int) (int, error) { 109 // Compose a contents profile per candidate from headline, skills, certs, customer agencies, and 110 // clearance, so the embedding represents what the candidate can do, not just a name. 111 q := ` 112SELECT c.id, 113 trim(coalesce(c.headline,'') 114 || ' | skills: ' || coalesce(s.skills,'') 115 || ' | certs: ' || coalesce(ce.certs,'') 116 || ' | agencies: '|| coalesce(cu.agencies,'') 117 || ' | clearance: '|| c.clearance_level 118 || ' | years: ' || c.years_experience) 119FROM candidates c 120LEFT JOIN (SELECT cs.candidate_id, string_agg(sk.name, ', ') skills 121 FROM candidate_skills cs JOIN skills sk ON sk.id = cs.skill_id GROUP BY cs.candidate_id) s ON s.candidate_id = c.id 122LEFT JOIN (SELECT candidate_id, string_agg(name, ', ') certs 123 FROM candidate_certifications GROUP BY candidate_id) ce ON ce.candidate_id = c.id 124LEFT JOIN (SELECT candidate_id, string_agg(agency, ', ') agencies 125 FROM candidate_customer_experience GROUP BY candidate_id) cu ON cu.candidate_id = c.id 126WHERE c.deleted_at IS NULL AND c.anonymized_at IS NULL` 127 args := []any{} 128 if org != "" { 129 args = append(args, org) 130 q += " AND c.organization_id = $1" 131 } 132 if onlyMissing { 133 q += " AND c.resume_embedding IS NULL" 134 } 135 q += " ORDER BY c.id" 136 rows, err := pool.Query(ctx, q, args...) 137 if err != nil { 138 return 0, err 139 } 140 type item struct{ id, text string } 141 var items []item 142 for rows.Next() { 143 var id, text string 144 if err := rows.Scan(&id, &text); err != nil { 145 rows.Close() 146 return 0, err 147 } 148 items = append(items, item{id: id, text: text}) 149 } 150 rows.Close() 151 return embedAndWrite(ctx, pool, client, batch, len(items), 152 func(i int) (string, string) { return items[i].id, items[i].text }, 153 "UPDATE candidates SET resume_embedding = $1::vector, updated_at = now() WHERE id = $2") 154} 155 156// embedAndWrite batches the embed calls and writes each normalized vector back via the given UPDATE. 157func embedAndWrite(ctx context.Context, pool *pgxpool.Pool, client *inference.Client, batch, n int, at func(int) (id, text string), update string) (int, error) { 158 if batch < 1 { 159 batch = 32 160 } 161 written := 0 162 for start := 0; start < n; start += batch { 163 end := start + batch 164 if end > n { 165 end = n 166 } 167 ids := make([]string, end-start) 168 texts := make([]string, end-start) 169 for i := start; i < end; i++ { 170 ids[i-start], texts[i-start] = at(i) 171 } 172 res, err := client.Embed(ctx, texts) 173 if err != nil { 174 return written, fmt.Errorf("embed batch [%d,%d): %w", start, end, err) 175 } 176 for j, vec := range res.Embeddings { 177 normalize(vec) 178 if _, err := pool.Exec(ctx, update, vectorLiteral(vec), ids[j]); err != nil { 179 return written, fmt.Errorf("write %s: %w", ids[j], err) 180 } 181 written++ 182 } 183 log.Printf("embed-backfill: %d/%d", written, n) 184 } 185 return written, nil 186} 187 188func writeAudit(ctx context.Context, pool *pgxpool.Pool, org, model string, nSkills, nCands int) { 189 payload, _ := json.Marshal(map[string]any{"skills": nSkills, "candidates": nCands, "model": model}) 190 var orgArg any 191 if org != "" { 192 orgArg = org 193 } 194 _, _ = pool.Exec(ctx, 195 "INSERT INTO audit_logs (organization_id, action, entity_type, new_values) VALUES ($1, 'aedt.embedding.backfill', 'embedding', $2)", 196 orgArg, payload) 197} 198 199// vectorLiteral formats a vector as a pgvector text literal for the $::vector cast. 200func vectorLiteral(v []float32) string { 201 var b strings.Builder 202 b.WriteByte('[') 203 for i, x := range v { 204 if i > 0 { 205 b.WriteByte(',') 206 } 207 b.WriteString(strconv.FormatFloat(float64(x), 'f', -1, 32)) 208 } 209 b.WriteByte(']') 210 return b.String() 211} 212 213// normalize L2 normalizes in place so cosine distance over the column is meaningful, matching the 214// embeddings package contract. 215func normalize(v []float32) { 216 var sum float64 217 for _, x := range v { 218 sum += float64(x) * float64(x) 219 } 220 if sum == 0 { 221 return 222 } 223 inv := float32(1.0 / math.Sqrt(sum)) 224 for i := range v { 225 v[i] *= inv 226 } 227} 228 229func envDefault(k, def string) string { 230 if v := os.Getenv(k); v != "" { 231 return v 232 } 233 return def 234} 235