profit/golangLAKEHOUSE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

corpusingest: extract reusable text→vector ingest pipeline

Browse Source 
Generalizes the staffing_500k driver's embed-and-push loop into
internal/corpusingest. Per docs/SPEC.md §3.4 component 1 (corpus
builders): adding a new staffing/code/playbook corpus is now one
Source impl + one main.go calling Run, not 200 lines of pipeline
copy-paste.

API:
  type Source interface { Next() (Row, error) }
  func Run(ctx, Config, Source) (Stats, error)

Library owns:
  - Index lifecycle (create, optional drop-existing, idempotent
    reuse on 409)
  - Parallel embed dispatcher (configurable workers + batch size)
  - Vectord push batching
  - Progress logging + Stats reporting
  - Partial-failure semantics (log + continue per-batch errors;
    operator decides on re-run via Stats.Embedded vs Scanned delta)

Per-corpus driver owns: source parsing + column→Row mapping +
post-ingest validation queries.

Refactor scripts/staffing_500k/main.go to use it. Driver is now
~190 lines (was 339), with the embed/add plumbing replaced by one
Run call. -drop flag added so callers can opt out of the destructive
DELETE-first behavior (default still true to keep the 500K test
clean-recall semantics).

Unit tests (internal/corpusingest/ingest_test.go, 8/8 PASS):
  - Pipeline shape: 50 rows / 16 batch → 4 embed + 4 add calls,
    every ID added exactly once, vectors at correct dimension
  - DropExisting fires DELETE
  - 409 on create → reuse existing index
  - Limit stops early
  - Empty Text rows skipped (counted as scanned, not added)
  - Required IndexName + Dimension validation
  - Context cancel stops mid-pipeline

Real bug caught and fixed by the test suite: if embedd ever returns
fewer vectors than texts in the request (degraded backend), the
addBatch loop would panic with index-out-of-range. Worker now
length-checks the response and logs+skips on mismatch.

12-smoke regression sweep all green (D1-D6, G1, G1P, G2,
storaged_cap, pathway, matrix). vet clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
root
2026-04-29 18:47:18 -05:00
co-authored by Claude Opus 4.7
parent c1d96b7b60
commit 166470f532
3 changed files with 916 additions and 251 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/staffing_500k/main.go
+130 -251
View File
@@ -1,13 +1,14 @@
// Staffing co-pilot scale test driver.
// Staffing co-pilot scale test driver — workers_500k corpus.
//
// Pipeline: workers_500k.csv → /v1/embed (batched, parallel) →
// /v1/vectors/index/workers_500k/add (batched). Then runs a handful
// of semantic queries against the populated index and prints the
// top hits — the human-readable check that "find workers like X"
// actually returns relevant workers.
// Pipeline: workers_500k.csv → /v1/embed → /v1/vectors/index/workers_500k/add.
// The pipeline itself lives in internal/corpusingest; this driver
// provides the CSV → Row mapping and the post-ingest semantic queries
// that are the human-readable check ("does forklift OSHA-30 actually
// retrieve forklift workers?").
//
// Designed to be re-run; index gets DELETEd at the start so leftover
// state from prior runs doesn't bias recall.
// Designed to be re-run safely; index gets DELETEd at the start
// when -drop is set so leftover state doesn't bias recall.
package main
import (
@@ -22,62 +23,123 @@ import (
"net/http"
"os"
"strings"
"sync"
"sync/atomic"
"time"
"git.agentview.dev/profit/golangLAKEHOUSE/internal/corpusingest"
)
const (
indexName = "workers_500k"
dim = 768
embedConcurrency = 8 // matches Ollama-on-A4000 sweet spot
embedBatchSize = 16 // texts per /v1/embed call
addBatchSize = 1000 // items per /v1/vectors/index/add call
maxColPhone = 4
maxColCity = 5
maxColState = 6
maxColRole = 2
maxColSkills = 8
maxColCerts = 9
maxColResume = 17
colWorkerID = 0
colName = 1
// Column indexes in workers_500k.csv. Stable contract; if the CSV
// schema changes these need updating.
colWorkerID = 0
colName = 1
colRole = 2
colCity = 5
colState = 6
colSkills = 8
colCerts = 9
colResume = 17
)
// workersCSV implements corpusingest.Source. CSV reader state +
// row → Row mapping live here; the embed/add pipeline is generic.
type workersCSV struct {
cr *csv.Reader
}
func (s *workersCSV) Next() (corpusingest.Row, error) {
for {
row, err := s.cr.Read()
if err != nil {
return corpusingest.Row{}, err
}
if len(row) <= colResume {
continue // skip malformed rows; matches prior behavior
}
id := strings.TrimSpace(row[colWorkerID])
return corpusingest.Row{
ID: "w-" + id,
Text: buildWorkerText(row),
Metadata: map[string]any{
"name": row[colName],
"role": row[colRole],
"city": row[colCity],
"state": row[colState],
},
}, nil
}
}
// buildWorkerText concatenates staffing-relevant columns into the
// embed-text. Order: role first (most semantically dense), then
// location, skills, certs, prose resume. Embedding models weight
// earlier tokens slightly more, so the front matter matters.
func buildWorkerText(row []string) string {
var b strings.Builder
b.WriteString(row[colRole])
b.WriteString(" in ")
b.WriteString(row[colCity])
b.WriteString(", ")
b.WriteString(row[colState])
b.WriteString(". Skills: ")
b.WriteString(row[colSkills])
b.WriteString(". Certifications: ")
b.WriteString(row[colCerts])
b.WriteString(". ")
b.WriteString(row[colResume])
return b.String()
}
func main() {
var (
gateway = flag.String("gateway", "http://127.0.0.1:3110", "gateway base URL")
csvPath = flag.String("csv", "/tmp/rs/workers_500k.csv", "path to workers CSV")
limit = flag.Int("limit", 0, "limit rows (0 = all)")
queries = flag.String("queries", "default", "default | <semicolon-separated query strings>")
skipPop = flag.Bool("skip-populate", false, "skip embed+add, only run queries")
gateway = flag.String("gateway", "http://127.0.0.1:3110", "gateway base URL")
csvPath = flag.String("csv", "/tmp/rs/workers_500k.csv", "path to workers CSV")
limit = flag.Int("limit", 0, "limit rows (0 = all)")
queries = flag.String("queries", "default", "default | <semicolon-separated query strings>")
skipPop = flag.Bool("skip-populate", false, "skip embed+add, only run queries")
drop = flag.Bool("drop", true, "DELETE index before populate (default true for clean recall)")
)
flag.Parse()
hc := &http.Client{Timeout: 5 * time.Minute}
ctx := context.Background()
if !*skipPop {
// Tear down any prior index so recall is on a fresh build.
fmt.Printf("[sc] DELETE %s/v1/vectors/index/%s (idempotent cleanup)\n", *gateway, indexName)
_ = httpDelete(hc, *gateway+"/v1/vectors/index/"+indexName)
// Create the index.
body := map[string]any{"name": indexName, "dimension": dim, "distance": "cosine"}
if code, msg := httpPostJSON(hc, *gateway+"/v1/vectors/index", body); code != 201 {
log.Fatalf("create index: %d %s", code, msg)
f, err := os.Open(*csvPath)
if err != nil {
log.Fatalf("open csv: %v", err)
}
fmt.Println("[sc] created index workers_500k dim=768 cosine")
t0 := time.Now()
if err := populate(hc, *gateway, *csvPath, *limit); err != nil {
log.Fatal(err)
defer f.Close()
cr := csv.NewReader(f)
cr.FieldsPerRecord = -1
if _, err := cr.Read(); err != nil { // skip header
log.Fatalf("read header: %v", err)
}
fmt.Printf("[sc] populate complete in %v\n", time.Since(t0))
stats, err := corpusingest.Run(ctx, corpusingest.Config{
GatewayURL: *gateway,
IndexName: indexName,
Dimension: dim,
Distance: "cosine",
EmbedBatch: 16, // matches Ollama-on-A4000 sweet spot
EmbedWorkers: 8, // matches Ollama-on-A4000 sweet spot
AddBatch: 1000, // empirically fine; vectord BatchAdd lock-amortized at f1c1883
Limit: *limit,
DropExisting: *drop,
HTTPClient: hc,
LogProgress: 10 * time.Second,
}, &workersCSV{cr: cr})
if err != nil {
log.Fatalf("ingest: %v", err)
}
fmt.Printf("[sc] populate done: scanned=%d embedded=%d added=%d wall=%v\n",
stats.Scanned, stats.Embedded, stats.Added, stats.Wall.Round(time.Millisecond))
}
// Validate semantic queries.
// Validate semantic queries against the populated index.
qs := defaultQueries()
if *queries != "default" {
qs = strings.Split(*queries, ";")
@@ -97,196 +159,35 @@ func defaultQueries() []string {
}
}
func populate(hc *http.Client, gateway, csvPath string, limit int) error {
f, err := os.Open(csvPath)
if err != nil {
return fmt.Errorf("open csv: %w", err)
}
defer f.Close()
cr := csv.NewReader(f)
cr.FieldsPerRecord = -1
if _, err := cr.Read(); err != nil { // header
return fmt.Errorf("read header: %w", err)
}
type job struct {
ids []string
texts []string
metas []json.RawMessage
}
jobs := make(chan job, embedConcurrency*2)
var wg sync.WaitGroup
var (
totalEmbedded int64
totalAdded int64
)
for i := 0; i < embedConcurrency; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for j := range jobs {
vecs, err := embedBatch(hc, gateway, j.texts)
if err != nil {
log.Printf("embed batch (%d items): %v", len(j.texts), err)
continue
}
atomic.AddInt64(&totalEmbedded, int64(len(vecs)))
if err := addBatch(hc, gateway, j.ids, vecs, j.metas); err != nil {
log.Printf("add batch (%d items): %v", len(j.ids), err)
continue
}
atomic.AddInt64(&totalAdded, int64(len(j.ids)))
}
}()
}
progressTicker := time.NewTicker(10 * time.Second)
go func() {
for range progressTicker.C {
fmt.Printf("[sc] progress: embedded=%d added=%d\n",
atomic.LoadInt64(&totalEmbedded), atomic.LoadInt64(&totalAdded))
}
}()
defer progressTicker.Stop()
curIDs := make([]string, 0, embedBatchSize)
curTexts := make([]string, 0, embedBatchSize)
curMetas := make([]json.RawMessage, 0, embedBatchSize)
rows := 0
for {
row, err := cr.Read()
if err == io.EOF {
break
}
if err != nil {
return fmt.Errorf("csv read row %d: %w", rows, err)
}
if len(row) <= maxColResume {
continue
}
id := strings.TrimSpace(row[colWorkerID])
text := buildSearchText(row)
meta, _ := json.Marshal(map[string]any{
"name": row[colName],
"role": row[maxColRole],
"city": row[maxColCity],
"state": row[maxColState],
})
curIDs = append(curIDs, "w-"+id)
curTexts = append(curTexts, text)
curMetas = append(curMetas, meta)
if len(curIDs) >= embedBatchSize {
jobs <- job{ids: curIDs, texts: curTexts, metas: curMetas}
curIDs = make([]string, 0, embedBatchSize)
curTexts = make([]string, 0, embedBatchSize)
curMetas = make([]json.RawMessage, 0, embedBatchSize)
}
rows++
if limit > 0 && rows >= limit {
break
}
}
if len(curIDs) > 0 {
jobs <- job{ids: curIDs, texts: curTexts, metas: curMetas}
}
close(jobs)
wg.Wait()
fmt.Printf("[sc] final: scanned=%d embedded=%d added=%d\n",
rows, atomic.LoadInt64(&totalEmbedded), atomic.LoadInt64(&totalAdded))
return nil
}
// buildSearchText concatenates the staffing-relevant columns into
// the text that gets embedded. Order: role first (most semantically
// dense), then skills + certs, city/state, finally the prose
// resume_text. Embedding models weight earlier tokens slightly more.
func buildSearchText(row []string) string {
var b strings.Builder
b.WriteString(row[maxColRole])
b.WriteString(" in ")
b.WriteString(row[maxColCity])
b.WriteString(", ")
b.WriteString(row[maxColState])
b.WriteString(". Skills: ")
b.WriteString(row[maxColSkills])
b.WriteString(". Certifications: ")
b.WriteString(row[maxColCerts])
b.WriteString(". ")
b.WriteString(row[maxColResume])
return b.String()
}
func embedBatch(hc *http.Client, gateway string, texts []string) ([][]float32, error) {
body := map[string]any{"texts": texts}
bs, _ := json.Marshal(body)
req, _ := http.NewRequest(http.MethodPost, gateway+"/v1/embed", bytes.NewReader(bs))
req.Header.Set("Content-Type", "application/json")
resp, err := hc.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
if resp.StatusCode != 200 {
preview, _ := io.ReadAll(io.LimitReader(resp.Body, 256))
return nil, fmt.Errorf("embed status %d: %s", resp.StatusCode, string(preview))
}
var er struct {
Vectors [][]float32 `json:"vectors"`
}
if err := json.NewDecoder(resp.Body).Decode(&er); err != nil {
return nil, err
}
return er.Vectors, nil
}
type addItem struct {
ID string `json:"id"`
Vector []float32 `json:"vector"`
Metadata json.RawMessage `json:"metadata"`
}
func addBatch(hc *http.Client, gateway string, ids []string, vecs [][]float32, metas []json.RawMessage) error {
items := make([]addItem, len(ids))
for i := range ids {
items[i] = addItem{ID: ids[i], Vector: vecs[i], Metadata: metas[i]}
}
bs, _ := json.Marshal(map[string]any{"items": items})
req, _ := http.NewRequest(http.MethodPost,
gateway+"/v1/vectors/index/"+indexName+"/add", bytes.NewReader(bs))
req.Header.Set("Content-Type", "application/json")
resp, err := hc.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != 200 {
preview, _ := io.ReadAll(io.LimitReader(resp.Body, 256))
return fmt.Errorf("add status %d: %s", resp.StatusCode, string(preview))
}
return nil
}
// runQuery embeds a query, searches the index, prints top hits.
// Stays in this driver (not corpusingest) — query validation is
// per-corpus concern, not part of the ingest pipeline.
func runQuery(hc *http.Client, gateway, q string) {
t0 := time.Now()
// 1. Embed the query.
vecs, err := embedBatch(hc, gateway, []string{q})
if err != nil || len(vecs) == 0 {
body, _ := json.Marshal(map[string]any{"texts": []string{q}})
req, _ := http.NewRequest(http.MethodPost, gateway+"/v1/embed", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
resp, err := hc.Do(req)
if err != nil {
fmt.Printf("[sc] query %q: embed err: %v\n", q, err)
return
}
defer resp.Body.Close()
var er struct {
Vectors [][]float32 `json:"vectors"`
}
if err := json.NewDecoder(resp.Body).Decode(&er); err != nil || len(er.Vectors) == 0 {
fmt.Printf("[sc] query %q: embed decode err: %v\n", q, err)
return
}
embedDur := time.Since(t0)
t1 := time.Now()
// 2. Search.
body := map[string]any{"vector": vecs[0], "k": 5}
bs, _ := json.Marshal(body)
req, _ := http.NewRequest(http.MethodPost,
gateway+"/v1/vectors/index/"+indexName+"/search", bytes.NewReader(bs))
body, _ = json.Marshal(map[string]any{"vector": er.Vectors[0], "k": 5})
req, _ = http.NewRequest(http.MethodPost,
gateway+"/v1/vectors/index/"+indexName+"/search", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
resp, err := hc.Do(req)
resp, err = hc.Do(req)
if err != nil {
fmt.Printf("[sc] query %q: search err: %v\n", q, err)
return
@@ -310,29 +211,7 @@ func runQuery(hc *http.Client, gateway, q string) {
}
}
func httpPostJSON(hc *http.Client, url string, body any) (int, string) {
bs, _ := json.Marshal(body)
req, _ := http.NewRequest(http.MethodPost, url, bytes.NewReader(bs))
req.Header.Set("Content-Type", "application/json")
resp, err := hc.Do(req)
if err != nil {
return 0, err.Error()
}
defer resp.Body.Close()
preview, _ := io.ReadAll(io.LimitReader(resp.Body, 256))
return resp.StatusCode, string(preview)
}
func httpDelete(hc *http.Client, url string) error {
req, _ := http.NewRequest(http.MethodDelete, url, nil)
resp, err := hc.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
io.Copy(io.Discard, resp.Body)
return nil
}
// keep context.Background reachable in case future paths use it
var _ = context.Background
// io.EOF imported transitively via corpusingest; keep the explicit
// reference so a hypothetical future "EOF means done" check in this
// driver's Source impl doesn't need a fresh import line.
var _ = io.EOF