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.
This commit is contained in:
2026-08-13 12:19:26 +01:00
committed by GitHub
co-authored by GitHub
parent ebc3f948c1
commit 669e184cf6
10 changed files with 394 additions and 160 deletions
+4
View File
@@ -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"
+6 -5
View File
@@ -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.
+19 -8
View File
@@ -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
}
}
}
+4 -2
View File
@@ -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:]))
}
}
+72
View File
@@ -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
}
+9
View File
@@ -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"
+100
View File
@@ -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
}
+143
View File
@@ -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)
}
}
+29 -135
View File
@@ -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)
}
}
+8 -10
View File
@@ -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"`
}
// 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