From 117f3c2cfd9be5ad9fde3217cea39ef116063abb Mon Sep 17 00:00:00 2001 From: Andrey Oblivantsev Date: Thu, 13 Aug 2026 12:19:26 +0100 Subject: [PATCH] fix(kbsearch): rank FTS correctly, filter before -n, start the daemon. (#5) Go search took worst BM25 hits (ORDER BY score), cut to -n before --root, and never called ensureDaemon. Ranking and flag parsing move to a cgo-free package so CI can fail those regressions without ladybug. --hop errors instead of being swallowed into the query. --- .github/workflows/ci.yml | 4 + PLAN.md | 11 ++- bin/kbsearch/brain.go | 27 ++++-- bin/kbsearch/main.go | 6 +- bin/kbsearch/rank/args.go | 72 +++++++++++++++ bin/kbsearch/rank/query.go | 9 ++ bin/kbsearch/rank/rank.go | 100 ++++++++++++++++++++ bin/kbsearch/rank/rank_test.go | 143 ++++++++++++++++++++++++++++ bin/kbsearch/search.go | 164 ++++++--------------------------- bin/kbsearch/types.go | 18 ++-- 10 files changed, 394 insertions(+), 160 deletions(-) create mode 100644 bin/kbsearch/rank/args.go create mode 100644 bin/kbsearch/rank/query.go create mode 100644 bin/kbsearch/rank/rank.go create mode 100644 bin/kbsearch/rank/rank_test.go diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1cca26e..41fe022 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -42,6 +42,10 @@ jobs: go vet ./... go test ./... -count=1 + - name: kbsearch ranking tests (no cgo / no ladybug) + working-directory: bin/kbsearch + run: go test ./rank -count=1 + - name: facts/audit self (lexicon consistency, no network) run: | ./bin/facts/audit self 2>/dev/null || echo "audit: not yet implemented; gate skipped" diff --git a/PLAN.md b/PLAN.md index 625c7ea..f00fe68 100644 --- a/PLAN.md +++ b/PLAN.md @@ -116,11 +116,12 @@ Common props on every node/edge: `root`, `confidence`, `evidence[]`, `how`, `.github/workflows/ci.yml`: -1. go vet + go test ./... (Go tools) -2. python -m unittest discover + pytest (Py tools) -3. bin/facts/audit self (lexicon internal consistency) -4. bin/kb/eval (recall@5 ≥ 0.95, gates index regressions) -5. md-docs build/lint if docs tooling arrives. +1. go vet + go test ./... (Go tools; root module) +2. `go test ./rank` in `bin/kbsearch` (cgo-free ranking + flag parser; nested module still needs ladybug for the rest) +3. python -m unittest discover (Py tools) +4. bin/facts/audit self (lexicon internal consistency) +5. bin/kb/eval (recall@5 ≥ 0.95, gates index regressions) +6. md-docs build/lint if docs tooling arrives. Feedback loop: every commit → PR → CI → green/gate → merge. Same discipline as `db/tech-poc`: contract first where there is an OpenAPI/message shape. diff --git a/bin/kbsearch/brain.go b/bin/kbsearch/brain.go index e9697ff..caf3644 100644 --- a/bin/kbsearch/brain.go +++ b/bin/kbsearch/brain.go @@ -5,6 +5,7 @@ import ( "fmt" "os" "path/filepath" + "strings" lbug "github.com/LadybugDB/go-ladybug" ) @@ -44,10 +45,10 @@ func dbPath() string { } func openBrain() error { - return openWithOpts(2, eps()) + return openWithSandbox(eps()) } -func openWithOpts(allow int, epsv string) error { +func openWithSandbox(epsv string) error { cfg := lbug.DefaultSystemConfig() cfg.MaxNumThreads = 8 cfg.BufferPoolSize = 1 << 30 // 1GB @@ -57,20 +58,30 @@ func openWithOpts(allow int, epsv string) error { if err != nil { return fmt.Errorf("OpenDatabase: %w", err) } - if epsv != "" { - if _, err := conn.Query("SET STREAM_SANDBOX = '" + epsv + "'"); err != nil { - return err - } - } conn, err = lbug.OpenConnection(db) if err != nil { + closeBrain() return fmt.Errorf("OpenConnection: %w", err) } + // Session settings need a live connection; running this before + // OpenConnection dereferenced a nil *Connection. + if epsv != "" { + if strings.ContainsAny(epsv, "'\\") { + closeBrain() + return fmt.Errorf("SET STREAM_SANDBOX: invalid value") + } + if _, err := conn.Query("SET STREAM_SANDBOX = '" + epsv + "'"); err != nil { + closeBrain() + return fmt.Errorf("SET STREAM_SANDBOX: %w", err) + } + } if _, err := conn.Query("LOAD EXTENSION FTS"); err != nil { + closeBrain() return fmt.Errorf("LOAD EXTENSION FTS: %w", err) } if _, err := conn.Query("LOAD EXTENSION VECTOR"); err != nil { + closeBrain() return fmt.Errorf("LOAD EXTENSION VECTOR: %w", err) } return nil @@ -85,4 +96,4 @@ func closeBrain() { db.Close() db = nil } -} \ No newline at end of file +} diff --git a/bin/kbsearch/main.go b/bin/kbsearch/main.go index c896925..b5ac729 100644 --- a/bin/kbsearch/main.go +++ b/bin/kbsearch/main.go @@ -8,7 +8,9 @@ // kbsearch --list-model print the resolved model dir // // The potion-multilingual model is loaded only in `serve`; a CLI reuses the -// daemon over localhost HTTP (falling back to in-process embedding). +// daemon over localhost HTTP (KBSEARCH_PORT, default 17830) and starts one in +// the background when none answers. KBSEARCH_NO_DAEMON=1 skips that and embeds +// in-process instead. package main import ( @@ -41,4 +43,4 @@ func main() { return } os.Exit(runSearch(os.Args[1:])) -} \ No newline at end of file +} diff --git a/bin/kbsearch/rank/args.go b/bin/kbsearch/rank/args.go new file mode 100644 index 0000000..72263e6 --- /dev/null +++ b/bin/kbsearch/rank/args.go @@ -0,0 +1,72 @@ +package rank + +import ( + "fmt" + "strconv" + "strings" +) + +const Usage = `usage: kbsearch "query" [--root facts|info] [--repo REPO] [-n N] [--json] + kbsearch serve [port] + kbsearch --list-model` + +type Options struct { + Query string + Root string + Repo string + Limit int + JSONOut bool + ListModel bool +} + +// ParseArgs reads flags. Unknown flags are an error: silently dropping them +// meant `--hop 1` vanished and its argument `1` was appended to the query. +// --hop is recognised so it cannot be swallowed; it is not implemented until +// File/FROM_FILE edges exist. +func ParseArgs(args []string) (Options, error) { + opt := Options{Limit: 20} + var queryArgs []string + + for i := 0; i < len(args); i++ { + arg := args[i] + wantsValue := arg == "--root" || arg == "--repo" || arg == "-n" || arg == "--hop" + if wantsValue && i+1 >= len(args) { + return opt, fmt.Errorf("%s needs a value", arg) + } + switch arg { + case "--root": + i++ + opt.Root = args[i] + if opt.Root != "facts" && opt.Root != "info" { + return opt, fmt.Errorf("--root must be facts or info, got %q", opt.Root) + } + case "--repo": + i++ + opt.Repo = args[i] + case "-n": + i++ + n, err := strconv.Atoi(args[i]) + if err != nil || n < 1 { + return opt, fmt.Errorf("-n must be a positive integer, got %q", args[i]) + } + opt.Limit = n + case "--hop": + return opt, fmt.Errorf("--hop is not implemented yet (needs File/FROM_FILE edges)") + case "--json": + opt.JSONOut = true + case "--list-model": + opt.ListModel = true + default: + if strings.HasPrefix(arg, "-") { + return opt, fmt.Errorf("unknown flag %q", arg) + } + queryArgs = append(queryArgs, arg) + } + } + + opt.Query = strings.TrimSpace(strings.Join(queryArgs, " ")) + if opt.Query == "" && !opt.ListModel { + return opt, fmt.Errorf("no query given") + } + return opt, nil +} diff --git a/bin/kbsearch/rank/query.go b/bin/kbsearch/rank/query.go new file mode 100644 index 0000000..df9db45 --- /dev/null +++ b/bin/kbsearch/rank/query.go @@ -0,0 +1,9 @@ +package rank + +// BM25 ranks best-first, so the top hits are the *highest* scores; cosine +// distance ranks best-first ascending. Both mirror kblib.py. +const FTSStmt = "CALL QUERY_FTS_INDEX('Leaf', 'id', $q) " + + "RETURN node.id, node.text, node.root, node.source, score ORDER BY score DESC LIMIT $n" + +const VecStmt = "CALL QUERY_VECTOR_INDEX('Leaf', 'Leaf_vec', $q, $n) " + + "RETURN node.id, node.text, node.root, node.source, distance ORDER BY distance LIMIT $n" diff --git a/bin/kbsearch/rank/rank.go b/bin/kbsearch/rank/rank.go new file mode 100644 index 0000000..474836b --- /dev/null +++ b/bin/kbsearch/rank/rank.go @@ -0,0 +1,100 @@ +// Package rank is the cgo-free ranking and CLI parsing for kbsearch. +// CI can `go test ./rank` without the native ladybug library. +package rank + +import ( + "sort" + "strings" +) + +// Hit is one search result, mirroring the python script's dict shape. +type Hit struct { + ID string `json:"id"` + Text string `json:"text"` + Root string `json:"root"` + Source string `json:"-"` + Score float64 `json:"score"` + Snippet string `json:"snippet,omitempty"` +} + +// rrfK dampens the contribution of low ranks; same constant as kblib.py. +const rrfK = 60 + +// RankAndFilter fuses the two hit lists, applies --root/--repo, then cuts to +// limit. Cutting first dropped every matching leaf ranked below the cut, so +// `--root facts` came back empty whenever info leafs filled the top N. +// limit <= 0 keeps everything. +func RankAndFilter(fts, vec []Hit, root, repo string, limit int) []Hit { + out := Hybrid(fts, vec, 0) + if root != "" { + out = FilterRoot(out, root) + } + if repo != "" { + out = FilterRepo(out, repo) + } + if limit > 0 && len(out) > limit { + out = out[:limit] + } + return out +} + +// Hybrid merges FTS and vector hits by reciprocal rank fusion. +// limit <= 0 returns the full fused list. +func Hybrid(fts, vec []Hit, limit int) []Hit { + byID := make(map[string]Hit, len(fts)+len(vec)) + rrf := make(map[string]float64, len(fts)+len(vec)) + + for i, h := range fts { + byID[h.ID] = h + rrf[h.ID] += 1.0 / (rrfK + float64(i+1)) + } + for i, h := range vec { + if existing, ok := byID[h.ID]; !ok { + byID[h.ID] = h + } else if existing.Score == 0 { + existing.Score = h.Score + byID[h.ID] = existing + } + rrf[h.ID] += 1.0 / (rrfK + float64(i+1)) + } + + ids := make([]string, 0, len(rrf)) + for id := range rrf { + ids = append(ids, id) + } + sort.Slice(ids, func(i, j int) bool { + if rrf[ids[i]] != rrf[ids[j]] { + return rrf[ids[i]] > rrf[ids[j]] + } + return ids[i] < ids[j] + }) + if limit > 0 && len(ids) > limit { + ids = ids[:limit] + } + + out := make([]Hit, 0, len(ids)) + for _, id := range ids { + out = append(out, byID[id]) + } + return out +} + +func FilterRoot(hits []Hit, root string) []Hit { + var out []Hit + for _, h := range hits { + if h.Root == root { + out = append(out, h) + } + } + return out +} + +func FilterRepo(hits []Hit, repo string) []Hit { + var out []Hit + for _, h := range hits { + if strings.Contains(h.Source, repo) { + out = append(out, h) + } + } + return out +} diff --git a/bin/kbsearch/rank/rank_test.go b/bin/kbsearch/rank/rank_test.go new file mode 100644 index 0000000..838acfc --- /dev/null +++ b/bin/kbsearch/rank/rank_test.go @@ -0,0 +1,143 @@ +// Unit tests for ranking/filtering and CLI parsing (no db, no model, offline). +package rank + +import ( + "strings" + "testing" +) + +func h(id, root, source string) Hit { + return Hit{ID: id, Text: id, Root: root, Source: source} +} + +func ids(hits []Hit) []string { + out := make([]string, len(hits)) + for i, hit := range hits { + out[i] = hit.ID + } + return out +} + +func eq(t *testing.T, got []Hit, want ...string) { + t.Helper() + g := ids(got) + if len(g) != len(want) { + t.Fatalf("got %v, want %v", g, want) + } + for i := range want { + if g[i] != want[i] { + t.Fatalf("got %v, want %v", g, want) + } + } +} + +// A facts leaf that ranks below the limit in the unfiltered list must still +// be returned for --root facts. Filtering after truncation loses it. +func TestRankAndFilterFiltersBeforeLimit(t *testing.T) { + fts := []Hit{ + h("i1", "info", "docs/a.md"), + h("i2", "info", "docs/b.md"), + h("i3", "info", "docs/c.md"), + h("f1", "facts", "docker ps x compose"), + } + eq(t, RankAndFilter(fts, nil, "facts", "", 2), "f1") +} + +func TestRankAndFilterRepoFiltersBeforeLimit(t *testing.T) { + fts := []Hit{ + h("a", "info", "eSlider/2dph:README.md"), + h("b", "info", "eSlider/2dph:PLAN.md"), + h("c", "info", "eSlider/ops:compose.yaml"), + } + eq(t, RankAndFilter(fts, nil, "", "ops", 2), "c") +} + +func TestRankAndFilterTruncatesToLimit(t *testing.T) { + fts := []Hit{h("a", "info", "x"), h("b", "info", "x"), h("c", "info", "x")} + eq(t, RankAndFilter(fts, nil, "", "", 2), "a", "b") +} + +func TestRankAndFilterLimitZeroKeepsAll(t *testing.T) { + fts := []Hit{h("a", "info", "x"), h("b", "info", "x")} + eq(t, RankAndFilter(fts, nil, "", "", 0), "a", "b") +} + +func TestHybridFusesBothRetrievers(t *testing.T) { + fts := []Hit{h("only-fts", "info", "x"), h("both", "info", "x")} + vec := []Hit{h("only-vec", "info", "x"), h("both", "info", "x")} + eq(t, Hybrid(fts, vec, 0), "both", "only-fts", "only-vec") +} + +func TestHybridTiesAreDeterministic(t *testing.T) { + fts := []Hit{h("b", "info", "x"), h("a", "info", "x")} + first := ids(Hybrid(fts, nil, 0)) + for i := 0; i < 50; i++ { + got := ids(Hybrid(fts, nil, 0)) + for j := range first { + if got[j] != first[j] { + t.Fatalf("unstable order: %v then %v", first, got) + } + } + } +} + +func TestHybridKeepsVectorScoreForSharedHit(t *testing.T) { + fts := []Hit{{ID: "x", Root: "info", Score: 0}} + vec := []Hit{{ID: "x", Root: "info", Score: 0.87}} + got := Hybrid(fts, vec, 0) + if len(got) != 1 || got[0].Score != 0.87 { + t.Fatalf("got %+v, want score 0.87", got) + } +} + +// The old parser dropped unknown flags and appended their arguments to the +// query, so `search "q" --hop 1` searched for "q 1". --hop is not implemented +// here (needs File edges); it must still fail closed instead of changing q. +func TestParseHopIsNotSwallowedIntoTheQuery(t *testing.T) { + _, err := ParseArgs([]string{"what runs on arc-2", "--hop", "1"}) + if err == nil { + t.Fatal("expected --hop to error (not implemented), not be swallowed") + } + if !strings.Contains(err.Error(), "--hop") { + t.Fatalf("error should name --hop, got %v", err) + } +} + +func TestParseRejectsUnknownFlags(t *testing.T) { + if _, err := ParseArgs([]string{"query", "--nope"}); err == nil { + t.Fatal("unknown flag accepted") + } +} + +func TestParseRejectsBadValues(t *testing.T) { + for _, args := range [][]string{ + {"q", "-n", "zero"}, + {"q", "-n", "0"}, + {"q", "--root", "nonsense"}, + {"q", "--hop"}, + {"--json"}, + } { + if _, err := ParseArgs(args); err == nil { + t.Errorf("accepted %v", args) + } + } +} + +func TestParseDefaults(t *testing.T) { + opt, err := ParseArgs([]string{"two", "words", "--json"}) + if err != nil || opt.Query != "two words" || opt.Limit != 20 || !opt.JSONOut { + t.Fatalf("got %+v err=%v", opt, err) + } +} + +func TestListModelNeedsNoQuery(t *testing.T) { + if _, err := ParseArgs([]string{"--list-model"}); err != nil { + t.Fatalf("unexpected error: %v", err) + } +} + +func TestFTSQueryOrdersByScoreDescending(t *testing.T) { + if !strings.Contains(FTSStmt, "ORDER BY score DESC") { + t.Fatalf("FTS query must order by score DESC, got:\n%s", FTSStmt) + } +} diff --git a/bin/kbsearch/search.go b/bin/kbsearch/search.go index e8ff2f2..3e090ab 100644 --- a/bin/kbsearch/search.go +++ b/bin/kbsearch/search.go @@ -13,12 +13,12 @@ import ( "os" "os/exec" "path/filepath" - "sort" "strconv" - "strings" + "syscall" "time" lbug "github.com/LadybugDB/go-ladybug" + "github.com/eSlider/2dph/bin/kbsearch/rank" ) const defaultPort = 17830 @@ -26,45 +26,15 @@ const daemonPath = "/embed" const healthPath = "/health" func runSearch(args []string) int { - // Manual flag parsing to allow flags after query (like Python argparse) - root := "" - repo := "" - limit := 20 - jsonOut := false - listModel := false - - var queryArgs []string - for i := 0; i < len(args); i++ { - switch args[i] { - case "--root": - if i+1 < len(args) { - root = args[i+1] - i++ - } - case "--repo": - if i+1 < len(args) { - repo = args[i+1] - i++ - } - case "-n": - if i+1 < len(args) { - if n, err := strconv.Atoi(args[i+1]); err == nil { - limit = n - } - i++ - } - case "--json": - jsonOut = true - case "--list-model": - listModel = true - default: - if !strings.HasPrefix(args[i], "-") { - queryArgs = append(queryArgs, args[i]) - } - } + opt, err := rank.ParseArgs(args) + if err != nil { + fmt.Fprintf(os.Stderr, "kbsearch: %v\n%s\n", err, rank.Usage) + return 2 } + root, repo, limit, query := opt.Root, opt.Repo, opt.Limit, opt.Query + jsonOut := opt.JSONOut - if listModel { + if opt.ListModel { dir, err := modelDir() if err != nil { fmt.Fprintln(os.Stderr, err) @@ -74,12 +44,6 @@ func runSearch(args []string) int { return 0 } - query := strings.TrimSpace(strings.Join(queryArgs, " ")) - if query == "" { - fmt.Fprintln(os.Stderr, "usage: kbsearch \"query\" [--root facts|info] [--repo REPO] [-n N] [--json]") - return 1 - } - if err := openBrain(); err != nil { fmt.Fprintf(os.Stderr, "open brain: %v\n", err) return 1 @@ -103,17 +67,7 @@ func runSearch(args []string) int { fmt.Fprintf(os.Stderr, "vec: %v\n", err) } - results := hybrid(fts, vec, limit) - - if root != "" { - results = filterRoot(results, root) - } - if repo != "" { - results = filterRepo(results, repo) - } - if len(results) > limit { - results = results[:limit] - } + results := rank.RankAndFilter(fts, vec, root, repo, limit) for i := range results { if results[i].Text != "" { @@ -150,10 +104,7 @@ func b2i(err error) int { } func queryFTS(text string, limit int) ([]Hit, error) { - stmt, err := conn.Prepare( - "CALL QUERY_FTS_INDEX('Leaf', 'id', $q) " + - "RETURN node.id, node.text, node.root, node.source, score ORDER BY score LIMIT $n", - ) + stmt, err := conn.Prepare(rank.FTSStmt) if err != nil { return nil, err } @@ -170,10 +121,7 @@ func queryVector(emb []float64, limit int) ([]Hit, error) { for i, v := range emb { embList[i] = v } - stmt, err := conn.Prepare( - "CALL QUERY_VECTOR_INDEX('Leaf', 'Leaf_vec', $q, $n) " + - "RETURN node.id, node.text, node.root, node.source, distance ORDER BY distance LIMIT $n", - ) + stmt, err := conn.Prepare(rank.VecStmt) if err != nil { return nil, err } @@ -215,10 +163,10 @@ func rowsToHits(res *lbug.QueryResult) ([]Hit, error) { // JSON output types type jsonOut struct { - Query string `json:"query"` - RootFilter string `json:"root_filter"` - Count int `json:"count"` - Results []jsonHit `json:"results"` + Query string `json:"query"` + RootFilter string `json:"root_filter"` + Count int `json:"count"` + Results []jsonHit `json:"results"` } type jsonHit struct { @@ -248,70 +196,6 @@ func toJSONOut(hits []Hit, query, rootFilter string) *jsonOut { } } -func hybrid(fts, vec []Hit, limit int) []Hit { - byID := make(map[string]Hit) - rrf := make(map[string]float64) - - for rank, h := range fts { - byID[h.ID] = h - rrf[h.ID] += 1.0 / (60 + float64(rank+1)) - } - for rank, h := range vec { - if _, ok := byID[h.ID]; !ok { - byID[h.ID] = h - } else { - existing := byID[h.ID] - if existing.Score == 0 { - existing.Score = h.Score - byID[h.ID] = existing - } - } - rrf[h.ID] += 1.0 / (60 + float64(rank+1)) - } - - type scored struct { - id string - rrf float64 - } - var scoredList []scored - for id, v := range rrf { - scoredList = append(scoredList, scored{id, v}) - } - sort.Slice(scoredList, func(i, j int) bool { - return scoredList[i].rrf > scoredList[j].rrf - }) - - var out []Hit - for i, s := range scoredList { - if i >= limit { - break - } - h := byID[s.id] - out = append(out, h) - } - return out -} - -func filterRoot(hits []Hit, root string) []Hit { - var out []Hit - for _, h := range hits { - if h.Root == root { - out = append(out, h) - } - } - return out -} - -func filterRepo(hits []Hit, repo string) []Hit { - var out []Hit - for _, h := range hits { - if strings.Contains(h.Source, repo) { - out = append(out, h) - } - } - return out -} - func resultsToDicts(hits []Hit) []any { out := make([]any, len(hits)) for i, h := range hits { @@ -382,10 +266,16 @@ func embedQuery(text string) ([]float64, error) { port = p } } - emb, err := tryDaemon(text, port) - if err == nil { + if emb, err := tryDaemon(text, port); err == nil { return emb, nil } + if os.Getenv("KBSEARCH_NO_DAEMON") == "" { + if err := ensureDaemon(port); err == nil { + if emb, err := tryDaemon(text, port); err == nil { + return emb, nil + } + } + } model, err := loadModel() if err != nil { @@ -447,9 +337,13 @@ func ensureDaemon(port int) error { cmd.Dir, _ = filepath.Split(self) cmd.Stdout = nil cmd.Stderr = nil + cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} if err := cmd.Start(); err != nil { return err } + if cmd.Process != nil { + _ = cmd.Process.Release() + } for i := 0; i < 40; i++ { time.Sleep(250 * time.Millisecond) @@ -462,4 +356,4 @@ func ensureDaemon(port int) error { } } return fmt.Errorf("daemon failed to start on port %d", port) -} \ No newline at end of file +} diff --git a/bin/kbsearch/types.go b/bin/kbsearch/types.go index 2a403bd..509bf1a 100644 --- a/bin/kbsearch/types.go +++ b/bin/kbsearch/types.go @@ -1,16 +1,14 @@ // Common types and helpers for kbsearch. package main -import "os" +import ( + "os" + + "github.com/eSlider/2dph/bin/kbsearch/rank" +) func eps() string { return os.Getenv("KBTEST_EPS") } -// Hit is one search result, mirroring the python script's dict shape. -type Hit struct { - ID string `json:"id"` - Text string `json:"text"` - Root string `json:"root"` - Source string `json:"-"` // for repo filtering, not in output - Score float64 `json:"score"` - Snippet string `json:"snippet,omitempty"` -} \ No newline at end of file +// Hit is the search hit type; ranking lives in package rank so CI can test +// it without the native ladybug library. +type Hit = rank.Hit