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:
committed by
GitHub
GitHub
co-authored by
GitHub
GitHub
parent
140d86a4b9
commit
117f3c2cfd
@@ -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
|
||||
}
|
||||
@@ -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"
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user