feat: cherry-pick scrub-secrets, session-diversify, bench from #21

.gitignore

@@ -26,3 +26,4 @@ Thumbs.db
 # Solr data (managed by Docker volume)
 solr-data/
 /solr-mem-indexer
+/solr-mem-bench

Makefile

@@ -1,4 +1,4 @@
-.PHONY: build build-indexer install install-indexer test tidy run dev up down logs reset docker-build docker-up config systemd-install systemd-uninstall indexer-install indexer-uninstall launchd-install launchd-install-server launchd-install-indexer launchd-uninstall skill-install skill-uninstall
+.PHONY: build build-indexer build-bench install install-indexer test tidy run dev up down logs reset docker-build docker-up config systemd-install systemd-uninstall indexer-install indexer-uninstall launchd-install launchd-install-server launchd-install-indexer launchd-uninstall skill-install skill-uninstall bench
 
 # Go targets
 build:
@@ -7,7 +7,19 @@ build:
 build-indexer:
 	go build -o bin/solr-mem-indexer ./cmd/solr-mem-indexer
 
-build-all: build build-indexer
+build-all: build build-indexer build-bench
+
+build-bench:
+	go build -o bin/solr-mem-bench ./cmd/solr-mem-bench
+
+# Retrieval benchmark. Seeds the memories collection with namespaced bench-*
+# docs (safe to run against a live collection — only touches bench-* IDs) and
+# runs the shipped query set. Override BENCH_URL to point elsewhere.
+BENCH_URL ?= http://pax89.local:8983/solr/memories
+bench: build-bench
+	./bin/solr-mem-bench -solr-url $(BENCH_URL) -seed \
+	  -corpus cmd/solr-mem-bench/testdata/corpus.jsonl \
+	  -queries cmd/solr-mem-bench/testdata/queries.jsonl
 
 install:
 	go install ./cmd/solr-mem-server

cmd/solr-mem-bench/main.go

@@ -0,0 +1,230 @@
+// solr-mem-bench is a small harness for measuring memory retrieval quality.
+//
+// Usage:
+//
+//	solr-mem-bench -solr-url http://localhost:8983/solr -collection memories_bench \
+//	  -corpus testdata/corpus.jsonl -queries testdata/queries.jsonl -seed
+//
+// The -seed flag clears the target collection and writes the corpus before
+// running queries. Without -seed, the harness assumes the collection is
+// already populated with the gold IDs referenced by queries.
+//
+// Gold format (JSONL, one per line):
+//
+//	corpus:  {"id":"mem-001","title":"...","content":"...","tags":["..."]}
+//	queries: {"id":"q-001","text":"database N+1 fix","gold":["mem-007"]}
+package main
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"log"
+	"os"
+	"time"
+
+	"github.com/arreyder/solr-mem/internal/solr"
+)
+
+type corpusItem struct {
+	ID      string   `json:"id"`
+	Title   string   `json:"title"`
+	Content string   `json:"content"`
+	Tags    []string `json:"tags,omitempty"`
+}
+
+type queryItem struct {
+	ID   string   `json:"id"`
+	Text string   `json:"text"`
+	Gold []string `json:"gold"`
+}
+
+func main() {
+	var (
+		solrURL    = flag.String("solr-url", "http://localhost:8983/solr/memories", "Base URL of the memories collection")
+		corpusPath = flag.String("corpus", "cmd/solr-mem-bench/testdata/corpus.jsonl", "Path to corpus JSONL")
+		queryPath  = flag.String("queries", "cmd/solr-mem-bench/testdata/queries.jsonl", "Path to queries JSONL")
+		seed       = flag.Bool("seed", false, "Clear the collection and write the corpus before running queries")
+		topK       = flag.Int("topk", 10, "Number of hits to retrieve per query")
+		variant    = flag.String("variant", "bm25", "Retrieval variant name (for labeling output)")
+	)
+	flag.Parse()
+
+	corpus, err := loadCorpus(*corpusPath)
+	if err != nil {
+		log.Fatalf("load corpus: %v", err)
+	}
+	queries, err := loadQueries(*queryPath)
+	if err != nil {
+		log.Fatalf("load queries: %v", err)
+	}
+	log.Printf("loaded %d corpus items, %d queries", len(corpus), len(queries))
+
+	client := solr.NewClient(*solrURL)
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+
+	if *seed {
+		log.Printf("seeding collection at %s (clearing first)", *solrURL)
+		if err := seedCorpus(ctx, client, corpus); err != nil {
+			log.Fatalf("seed: %v", err)
+		}
+		// Give Solr a moment to commit.
+		time.Sleep(1 * time.Second)
+	}
+
+	agg := NewAggregate(1, 3, 5, 10)
+	perQuery := make([]queryResult, 0, len(queries))
+
+	start := time.Now()
+	for _, q := range queries {
+		hits, err := runQuery(ctx, client, q.Text, *topK)
+		if err != nil {
+			log.Printf("query %s failed: %v", q.ID, err)
+			continue
+		}
+		agg.Add(q.Gold, hits)
+		perQuery = append(perQuery, queryResult{ID: q.ID, Text: q.Text, Hits: hits, Gold: q.Gold})
+	}
+	elapsed := time.Since(start)
+
+	agg.Finalize()
+	printReport(*variant, agg, perQuery, elapsed)
+}
+
+type queryResult struct {
+	ID   string
+	Text string
+	Hits []string
+	Gold []string
+}
+
+func loadCorpus(path string) ([]corpusItem, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+	var out []corpusItem
+	s := bufio.NewScanner(f)
+	s.Buffer(make([]byte, 0, 64*1024), 1024*1024)
+	for s.Scan() {
+		line := s.Bytes()
+		if len(line) == 0 {
+			continue
+		}
+		var item corpusItem
+		if err := json.Unmarshal(line, &item); err != nil {
+			return nil, fmt.Errorf("parse corpus line: %w", err)
+		}
+		out = append(out, item)
+	}
+	return out, s.Err()
+}
+
+func loadQueries(path string) ([]queryItem, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+	var out []queryItem
+	s := bufio.NewScanner(f)
+	s.Buffer(make([]byte, 0, 64*1024), 1024*1024)
+	for s.Scan() {
+		line := s.Bytes()
+		if len(line) == 0 {
+			continue
+		}
+		var item queryItem
+		if err := json.Unmarshal(line, &item); err != nil {
+			return nil, fmt.Errorf("parse query line: %w", err)
+		}
+		out = append(out, item)
+	}
+	return out, s.Err()
+}
+
+// seedCorpus writes the bench corpus to Solr. It scopes its cleanup to
+// IDs starting with "bench-" so it is safe to run against a live memories
+// collection — it will only ever touch its own namespaced docs.
+func seedCorpus(ctx context.Context, client *solr.Client, items []corpusItem) error {
+	for _, it := range items {
+		if !isBenchID(it.ID) {
+			return fmt.Errorf("corpus id %q must start with 'bench-' for safe seeding", it.ID)
+		}
+	}
+	if err := client.DeleteByQuery(ctx, "id:bench-*"); err != nil {
+		return fmt.Errorf("clear bench docs: %w", err)
+	}
+	docs := make([]solr.Document, 0, len(items))
+	now := time.Now().UTC()
+	for _, it := range items {
+		docs = append(docs, solr.Document{
+			ID:        it.ID,
+			Title:     it.Title,
+			Content:   it.Content,
+			Tags:      it.Tags,
+			CreatedAt: now,
+			UpdatedAt: now,
+			Lifetime:  "permanent",
+			Format:    "prose",
+		})
+	}
+	return client.Add(ctx, docs...)
+}
+
+func isBenchID(s string) bool {
+	return len(s) > 6 && s[:6] == "bench-"
+}
+
+func runQuery(ctx context.Context, client *solr.Client, text string, topK int) ([]string, error) {
+	params := solr.QueryParams{
+		Query:  text,
+		Rows:   topK,
+		Fields: []string{"id"},
+	}
+	resp, err := client.Query(ctx, params)
+	if err != nil {
+		return nil, err
+	}
+	out := make([]string, 0, len(resp.Docs))
+	for _, d := range resp.Docs {
+		if id, ok := d["id"].(string); ok {
+			out = append(out, id)
+		}
+	}
+	return out, nil
+}
+
+func printReport(variant string, agg *Aggregate, results []queryResult, elapsed time.Duration) {
+	fmt.Printf("# solr-mem retrieval benchmark\n\n")
+	fmt.Printf("Variant: `%s` | Queries: %d | Elapsed: %s\n\n", variant, agg.Queries, elapsed.Round(time.Millisecond))
+	fmt.Printf("| Metric | Value |\n|---|---|\n")
+	for _, k := range []int{1, 3, 5, 10} {
+		if v, ok := agg.RecallAt[k]; ok {
+			fmt.Printf("| R@%d | %.3f |\n", k, v)
+		}
+	}
+	fmt.Printf("| MRR | %.3f |\n\n", agg.MRR)
+
+	fmt.Printf("## Per-query breakdown\n\n")
+	fmt.Printf("| Query | Gold | R@5 | MRR | Top hits |\n|---|---|---|---|---|\n")
+	for _, r := range results {
+		top := r.Hits
+		if len(top) > 5 {
+			top = top[:5]
+		}
+		fmt.Printf("| %s (%s) | %v | %.2f | %.2f | %v |\n",
+			r.ID, truncate(r.Text, 40), r.Gold, RecallAtK(r.Gold, r.Hits, 5), MRR(r.Gold, r.Hits), top)
+	}
+}
+
+func truncate(s string, n int) string {
+	if len(s) <= n {
+		return s
+	}
+	return s[:n-1] + "…"
+}

cmd/solr-mem-bench/metrics.go

@@ -0,0 +1,78 @@
+package main
+
+// RecallAtK returns the fraction of gold IDs that appear in the top-K hits.
+// gold is the set of relevant IDs for a query; hits is the ranked result list.
+func RecallAtK(gold, hits []string, k int) float64 {
+	if len(gold) == 0 {
+		return 0
+	}
+	if k <= 0 || k > len(hits) {
+		k = len(hits)
+	}
+	goldSet := make(map[string]struct{}, len(gold))
+	for _, g := range gold {
+		goldSet[g] = struct{}{}
+	}
+	found := 0
+	for i := 0; i < k; i++ {
+		if _, ok := goldSet[hits[i]]; ok {
+			found++
+		}
+	}
+	return float64(found) / float64(len(gold))
+}
+
+// MRR returns the reciprocal rank of the first gold hit in the ranked list,
+// or 0 if no gold item was retrieved. Ranks are 1-indexed.
+func MRR(gold, hits []string) float64 {
+	if len(gold) == 0 || len(hits) == 0 {
+		return 0
+	}
+	goldSet := make(map[string]struct{}, len(gold))
+	for _, g := range gold {
+		goldSet[g] = struct{}{}
+	}
+	for i, h := range hits {
+		if _, ok := goldSet[h]; ok {
+			return 1.0 / float64(i+1)
+		}
+	}
+	return 0
+}
+
+// Aggregate holds per-run averaged metrics.
+type Aggregate struct {
+	Queries  int
+	RecallAt map[int]float64 // k -> mean recall
+	MRR      float64
+}
+
+// NewAggregate accumulates per-query metrics over a set of queries.
+func NewAggregate(ks ...int) *Aggregate {
+	m := make(map[int]float64, len(ks))
+	for _, k := range ks {
+		m[k] = 0
+	}
+	return &Aggregate{RecallAt: m}
+}
+
+// Add records one query's hits against its gold.
+func (a *Aggregate) Add(gold, hits []string) {
+	a.Queries++
+	for k := range a.RecallAt {
+		a.RecallAt[k] += RecallAtK(gold, hits, k)
+	}
+	a.MRR += MRR(gold, hits)
+}
+
+// Finalize divides sums by query count, producing mean metrics.
+func (a *Aggregate) Finalize() {
+	if a.Queries == 0 {
+		return
+	}
+	n := float64(a.Queries)
+	for k, v := range a.RecallAt {
+		a.RecallAt[k] = v / n
+	}
+	a.MRR /= n
+}