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

Source identity

Field	Recorded value
Source type	GitHub source
Original record	Janady13/missionmatch / cmd/embed-backfill/main.go
Date	Not specified in source
Repository	Janady13/missionmatch
Revision	2178e7298d0e5d9af83afa12402fd587f0ba6ed8
Branch	feature/aioe-upgrade
Path or archive member	cmd/embed-backfill/main.go
Original SHA-256	99edf6f2523704f13a5ed558b10f8fbb783571810874a4c86dabf8fe8e70077c
Git blob SHA	105ee8b84e2e403edd78769873fd4841db0ada5a
Reading-copy SHA-256	99edf6f2523704f13a5ed558b10f8fbb783571810874a4c86dabf8fe8e70077c
Original source ↗
Recovered source transcript

Download complete source text ↓245 lines · 8.4 KB
// Command embed-backfill populates the embedding columns that the audit found dead: skills.embedding
// (0 of 61) and candidates.resume_embedding (sparse). It composes a contents profile per row, embeds it
// through the stateless MEGAMIND inference plane (384 dim, matching the vector(384) columns), L2
// normalizes, and writes it back. After a run the HNSW indexes and the SearchSkillsByEmbedding /
// SearchCandidatesHybrid queries become live.
//
// It targets ONE database (a tenant database under the database per tenant model, or the current shared
// database). It writes only embeddings, never identifiers off box: text goes to the plane, vectors come
// back. It writes a single append only AEDT audit summary at the end (counts and model only, no PII).
//
// Usage:
//
//	embed-backfill -dsn $MM_MIGRATOR_DATABASE_URL -megamind $MM_MEGAMIND_URL -what all [-org <uuid>] [-only-missing]
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"math"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/missionworkforce/missionmatch/internal/platform/embeddings"
	"github.com/missionworkforce/missionmatch/internal/platform/inference"
	"github.com/missionworkforce/missionmatch/internal/platform/secrets"
)

func main() {
	dsn := flag.String("dsn", os.Getenv("MM_MIGRATOR_DATABASE_URL"), "target database DSN (defaults to MM_MIGRATOR_DATABASE_URL)")
	mega := flag.String("megamind", os.Getenv("MM_MEGAMIND_URL"), "MEGAMIND gateway base URL")
	model := flag.String("model", envDefault("MM_MEGAMIND_EMBED_MODEL", "qllama/bge-small-en-v1.5"), "384 dim embed model")
	what := flag.String("what", "all", "what to embed: skills | candidates | all")
	org := flag.String("org", "", "optional organization_id filter (for the shared database)")
	onlyMissing := flag.Bool("only-missing", false, "embed only rows whose embedding is currently NULL")
	tokenFlag := flag.String("token", "", "MEGAMIND gateway bearer token (defaults to the secret store MM_MEGAMIND_TOKEN)")
	batch := flag.Int("batch", 32, "embedding batch size")
	flag.Parse()

	if *dsn == "" || *mega == "" {
		log.Fatal("embed-backfill: -dsn and -megamind are required")
	}
	ctx := context.Background()
	pool, err := pgxpool.New(ctx, *dsn)
	if err != nil {
		log.Fatalf("db: %v", err)
	}
	defer pool.Close()

	// The gateway now requires a bearer token. Prefer the flag, else the secret store (MM_SECRETS_FILE
	// then env), so this cadence tool authenticates exactly like the API.
	token := *tokenFlag
	if token == "" {
		if st, serr := secrets.Open(os.Getenv("MM_SECRETS_FILE")); serr == nil {
			token, _ = st.Get("MM_MEGAMIND_TOKEN")
		}
	}
	client, err := inference.New(inference.Config{BaseURL: *mega, EmbedModel: *model, Dimension: 384, Timeout: 60 * time.Second, MaxRetries: 2, Token: token})
	if err != nil {
		log.Fatalf("inference: %v", err)
	}
	if err := client.Health(ctx); err != nil {
		log.Fatalf("MEGAMIND not reachable at %s: %v", *mega, err)
	}

	var nSkills, nCands int
	if *what == "skills" || *what == "all" {
		if nSkills, err = backfillSkills(ctx, pool, client, *onlyMissing, *batch); err != nil {
			log.Fatalf("skills: %v", err)
		}
		log.Printf("embed-backfill: embedded %d skills", nSkills)
	}
	if *what == "candidates" || *what == "all" {
		if nCands, err = backfillCandidates(ctx, pool, client, *org, *onlyMissing, *batch); err != nil {
			log.Fatalf("candidates: %v", err)
		}
		log.Printf("embed-backfill: embedded %d candidates", nCands)
	}
	writeAudit(ctx, pool, *org, *model, nSkills, nCands)
	log.Printf("embed-backfill: DONE (skills=%d candidates=%d model=%s)", nSkills, nCands, *model)
}

func backfillSkills(ctx context.Context, pool *pgxpool.Pool, client *inference.Client, onlyMissing bool, batch int) (int, error) {
	q := "SELECT id, name, COALESCE(category,'') FROM skills"
	if onlyMissing {
		q += " WHERE embedding IS NULL"
	}
	q += " ORDER BY id"
	rows, err := pool.Query(ctx, q)
	if err != nil {
		return 0, err
	}
	type item struct {
		id   string
		text string
	}
	var items []item
	for rows.Next() {
		var id, name, cat string
		if err := rows.Scan(&id, &name, &cat); err != nil {
			rows.Close()
			return 0, err
		}
		items = append(items, item{id: id, text: embeddings.SkillEmbedText(name, cat)})
	}
	rows.Close()
	return embedAndWrite(ctx, pool, client, batch, len(items),
		func(i int) (string, string) { return items[i].id, items[i].text },
		"UPDATE skills SET embedding = $1::vector WHERE id = $2")
}

func backfillCandidates(ctx context.Context, pool *pgxpool.Pool, client *inference.Client, org string, onlyMissing bool, batch int) (int, error) {
	// Compose a contents profile per candidate from headline, skills, certs, customer agencies, and
	// clearance, so the embedding represents what the candidate can do, not just a name.
	q := `
SELECT c.id,
  trim(coalesce(c.headline,'')
    || ' | skills: '  || coalesce(s.skills,'')
    || ' | certs: '   || coalesce(ce.certs,'')
    || ' | agencies: '|| coalesce(cu.agencies,'')
    || ' | clearance: '|| c.clearance_level
    || ' | years: '   || c.years_experience)
FROM candidates c
LEFT JOIN (SELECT cs.candidate_id, string_agg(sk.name, ', ') skills
           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
LEFT JOIN (SELECT candidate_id, string_agg(name, ', ') certs
           FROM candidate_certifications GROUP BY candidate_id) ce ON ce.candidate_id = c.id
LEFT JOIN (SELECT candidate_id, string_agg(agency, ', ') agencies
           FROM candidate_customer_experience GROUP BY candidate_id) cu ON cu.candidate_id = c.id
WHERE c.deleted_at IS NULL AND c.anonymized_at IS NULL`
	args := []any{}
	if org != "" {
		args = append(args, org)
		q += " AND c.organization_id = $1"
	}
	if onlyMissing {
		q += " AND c.resume_embedding IS NULL"
	}
	q += " ORDER BY c.id"
	rows, err := pool.Query(ctx, q, args...)
	if err != nil {
		return 0, err
	}
	type item struct{ id, text string }
	var items []item
	for rows.Next() {
		var id, text string
		if err := rows.Scan(&id, &text); err != nil {
			rows.Close()
			return 0, err
		}
		items = append(items, item{id: id, text: text})
	}
	rows.Close()
	return embedAndWrite(ctx, pool, client, batch, len(items),
		func(i int) (string, string) { return items[i].id, items[i].text },
		"UPDATE candidates SET resume_embedding = $1::vector, updated_at = now() WHERE id = $2")
}

// embedAndWrite batches the embed calls and writes each normalized vector back via the given UPDATE.
func embedAndWrite(ctx context.Context, pool *pgxpool.Pool, client *inference.Client, batch, n int, at func(int) (id, text string), update string) (int, error) {
	if batch < 1 {
		batch = 32
	}
	written := 0
	for start := 0; start < n; start += batch {
		end := start + batch
		if end > n {
			end = n
		}
		ids := make([]string, end-start)
		texts := make([]string, end-start)
		for i := start; i < end; i++ {
			ids[i-start], texts[i-start] = at(i)
		}
		res, err := client.Embed(ctx, texts)
		if err != nil {
			return written, fmt.Errorf("embed batch [%d,%d): %w", start, end, err)
		}
		for j, vec := range res.Embeddings {
			normalize(vec)
			if _, err := pool.Exec(ctx, update, vectorLiteral(vec), ids[j]); err != nil {
				return written, fmt.Errorf("write %s: %w", ids[j], err)
			}
			written++
		}
		log.Printf("embed-backfill: %d/%d", written, n)
	}
	return written, nil
}

func writeAudit(ctx context.Context, pool *pgxpool.Pool, org, model string, nSkills, nCands int) {
	payload, _ := json.Marshal(map[string]any{"skills": nSkills, "candidates": nCands, "model": model})
	var orgArg any
	if org != "" {
		orgArg = org
	}
	_, _ = pool.Exec(ctx,
		"INSERT INTO audit_logs (organization_id, action, entity_type, new_values) VALUES ($1, 'aedt.embedding.backfill', 'embedding', $2)",
		orgArg, payload)
}

// vectorLiteral formats a vector as a pgvector text literal for the $::vector cast.
func vectorLiteral(v []float32) string {
	var b strings.Builder
	b.WriteByte('[')
	for i, x := range v {
		if i > 0 {
			b.WriteByte(',')
		}
		b.WriteString(strconv.FormatFloat(float64(x), 'f', -1, 32))
	}
	b.WriteByte(']')
	return b.String()
}

// normalize L2 normalizes in place so cosine distance over the column is meaningful, matching the
// embeddings package contract.
func normalize(v []float32) {
	var sum float64
	for _, x := range v {
		sum += float64(x) * float64(x)
	}
	if sum == 0 {
		return
	}
	inv := float32(1.0 / math.Sqrt(sum))
	for i := range v {
		v[i] *= inv
	}
}

func envDefault(k, def string) string {
	if v := os.Getenv(k); v != "" {
		return v
	}
	return def
}