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