diff --git a/.gitignore b/.gitignore index 20bf5c2..defce95 100644 --- a/.gitignore +++ b/.gitignore @@ -8,4 +8,5 @@ __pycache__/ .DS_Store *.env .env -.secrets/ \ No newline at end of file +.secrets/ +lib-ladybug/ diff --git a/bin/kb/search b/bin/kb/search index e759371..3c24a47 100755 --- a/bin/kb/search +++ b/bin/kb/search @@ -1,76 +1,34 @@ -#!/usr/bin/env python3 -"""kb/search - deduction search over the 2dph brain. +#!/usr/bin/env bash +# bin/kb/search - Go deduction search over the brain (model served by daemon). +# Builds the kbsearch binary on first run / when source changes, then execs it. +set -euo pipefail - bin/kb/search "query" # hybrid facts+info, YAML out - bin/kb/search "query" --root facts # confirmed facts only - bin/kb/search "query" --hop 1 # follow graph edges after hitting - bin/kb/search "query" --json | yq '.' - bin/kb/search "query" -n 5 # more results +KB="$(cd "$(dirname "$0")/../.." && pwd)" +BIN="$KB/var/bin/kbsearch" +SRC="$KB/bin/kbsearch" -Deduction order: facts root first (confirmed answers with evidence links), -then info root (marked `(not confirmed)`). --root restricts to one root. ---hop N walks FROM_FILE edges (sibling leafs in the same source file). -""" -from __future__ import annotations +mkdir -p "$KB/var/bin" -import json -import sys -from pathlib import Path +# Rebuild if binary missing or any .go source newer +need_build=0 +if [ ! -x "$BIN" ]; then + need_build=1 +else + # Check if any .go in kbsearch is newer than binary + while IFS= read -r -d '' f; do + if [ "$f" -nt "$BIN" ]; then + need_build=1 + break + fi + done < <(find "$SRC" -name '*.go' -print0 2>/dev/null) +fi -ROOT = Path(__file__).resolve().parents[2] -sys.path.insert(0, str(ROOT / "bin" / "tools")) +if [ "$need_build" -eq 1 ]; then + echo "Building kbsearch..." >&2 + (cd "$SRC" && \ + CGO_CFLAGS="-I$KB/lib-ladybug" \ + CGO_LDFLAGS="-L$KB/lib-ladybug -Wl,-rpath,$KB/lib-ladybug" \ + go build -tags system_ladybug -o "$BIN" .) || exit 1 +fi -from kblib import connect, hybrid_search, init_schema, open_readonly, query_fts # noqa: E402 -from yamlout import to_yaml # noqa: E402 -import ladybug # noqa: E402 - - -def main(argv: list[str]) -> int: - import argparse - p = argparse.ArgumentParser(description="deduction search over the brain") - p.add_argument("query") - p.add_argument("--root", choices=("facts", "info", None), default=None) - p.add_argument("--repo", default=None, help="filter results to one repo (source prefix)") - p.add_argument("--hop", type=int, default=0) - p.add_argument("-n", "--limit", type=int, default=10) - p.add_argument("--json", action="store_true") - a = p.parse_args(argv) - - try: - db, conn = open_readonly() - except FileNotFoundError as e: - print(e, file=sys.stderr) - return 1 - - from model2vec import StaticModel - model = StaticModel.from_pretrained("minishlab/potion-multilingual-128M") - emb = model.encode([a.query])[0].astype(float).tolist() - - rhs: list[dict] = [] - try: - rhs = query_fts(conn, a.query, a.limit * 2) - except Exception: - rhs = [] - - results = hybrid_search(conn, emb, rhs, a.limit) - if a.root: - results = [h for h in results if h["root"] == a.root] - if a.repo: - repo = a.repo - results = [h for h in results if repo in (h.get("source") or "")] - - for hit in results: - hit.pop("rrf", None) - if hit.get("text"): - hit["snippet"] = hit["text"][:280] - - out = {"query": a.query, "root_filter": a.root or "facts+info", - "count": len(results), "results": results} - print(json.dumps(out, indent=2, ensure_ascii=False) if a.json else to_yaml(out)) - conn.close() - db.close() - return 0 - - -if __name__ == "__main__": - sys.exit(main(sys.argv[1:])) \ No newline at end of file +exec "$BIN" "$@" \ No newline at end of file diff --git a/bin/kbsearch/brain.go b/bin/kbsearch/brain.go new file mode 100644 index 0000000..e9697ff --- /dev/null +++ b/bin/kbsearch/brain.go @@ -0,0 +1,88 @@ +// Brain connection management using go-ladybug. +package main + +import ( + "fmt" + "os" + "path/filepath" + + lbug "github.com/LadybugDB/go-ladybug" +) + +var ( + db *lbug.Database + conn *lbug.Connection +) + +func repoRoot() string { + // Try KB_ROOT env, then walk up from binary + if v := os.Getenv("KB_ROOT"); v != "" { + return v + } + self, err := os.Executable() + if err == nil { + dir := filepath.Dir(self) + for i := 0; i < 5; i++ { + if _, err := os.Stat(filepath.Join(dir, "var")); err == nil { + return dir + } + if _, err := os.Stat(filepath.Join(dir, ".git")); err == nil { + return dir + } + parent := filepath.Dir(dir) + if parent == dir { + break + } + dir = parent + } + } + return "." +} + +func dbPath() string { + return filepath.Join(repoRoot(), "var", "kb.lbug") +} + +func openBrain() error { + return openWithOpts(2, eps()) +} + +func openWithOpts(allow int, epsv string) error { + cfg := lbug.DefaultSystemConfig() + cfg.MaxNumThreads = 8 + cfg.BufferPoolSize = 1 << 30 // 1GB + + var err error + db, err = lbug.OpenDatabase(dbPath(), cfg) + 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 { + return fmt.Errorf("OpenConnection: %w", err) + } + if _, err := conn.Query("LOAD EXTENSION FTS"); err != nil { + return fmt.Errorf("LOAD EXTENSION FTS: %w", err) + } + if _, err := conn.Query("LOAD EXTENSION VECTOR"); err != nil { + return fmt.Errorf("LOAD EXTENSION VECTOR: %w", err) + } + return nil +} + +func closeBrain() { + if conn != nil { + conn.Close() + conn = nil + } + if db != nil { + db.Close() + db = nil + } +} \ No newline at end of file diff --git a/bin/kbsearch/go.mod b/bin/kbsearch/go.mod new file mode 100644 index 0000000..5054e8e --- /dev/null +++ b/bin/kbsearch/go.mod @@ -0,0 +1,23 @@ +module github.com/eSlider/2dph/bin/kbsearch + +go 1.26.0 + +require ( + github.com/LadybugDB/go-ladybug v0.17.0 + github.com/chewxy/math32 v1.11.2 + github.com/daulet/tokenizers v1.27.0 +) + +require ( + github.com/apache/arrow-go/v18 v18.6.0 // indirect + github.com/goccy/go-json v0.10.6 // indirect + github.com/google/flatbuffers v25.12.19+incompatible // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/klauspost/compress v1.18.5 // indirect + github.com/klauspost/cpuid/v2 v2.3.0 // indirect + github.com/pierrec/lz4/v4 v4.1.26 // indirect + github.com/shopspring/decimal v1.4.0 // indirect + github.com/zeebo/xxh3 v1.1.0 // indirect + golang.org/x/exp v0.0.0-20260112195511-716be5621a96 // indirect + golang.org/x/sys v0.43.0 // indirect +) diff --git a/bin/kbsearch/go.sum b/bin/kbsearch/go.sum new file mode 100644 index 0000000..0064c1b --- /dev/null +++ b/bin/kbsearch/go.sum @@ -0,0 +1,44 @@ +github.com/LadybugDB/go-ladybug v0.17.0 h1:RXDbkBjrbRmLdEbhGl4CLOIEzSt09gbP0n9UbKDEfwI= +github.com/LadybugDB/go-ladybug v0.17.0/go.mod h1:GeIXmE8XyF5TFS94NAuTag7vgCC+no/HTBMRA6Rd5Cs= +github.com/andybalholm/brotli v1.2.1 h1:R+f5xP285VArJDRgowrfb9DqL18yVK0gKAW/F+eTWro= +github.com/andybalholm/brotli v1.2.1/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= +github.com/apache/arrow-go/v18 v18.6.0 h1:GX/Jyd3R7mCLiECAwY9FWbbaYblie2WXBSz4Sw8fNpM= +github.com/apache/arrow-go/v18 v18.6.0/go.mod h1:gm3MiPpY82fLYK5VKPB3WoJbsiLVDfT7flD5/vHReKw= +github.com/apache/thrift v0.22.0 h1:r7mTJdj51TMDe6RtcmNdQxgn9XcyfGDOzegMDRg47uc= +github.com/apache/thrift v0.22.0/go.mod h1:1e7J/O1Ae6ZQMTYdy9xa3w9k+XHWPfRvdPyJeynQ+/g= +github.com/chewxy/math32 v1.11.2 h1:IufN08Zwr1NKuWfY+4Tz55BcwKmyKKNdOP7KtumehnM= +github.com/chewxy/math32 v1.11.2/go.mod h1:dOB2rcuFrCn6UHrze36WSLVPKtzPMRAQvBvUwkSsLqs= +github.com/daulet/tokenizers v1.27.0 h1:MmFYAEDFz69s/nNQfHg59DWqHz3v94m99kEZ/JbL+s4= +github.com/daulet/tokenizers v1.27.0/go.mod h1:YjFY1o1HGMyWkQgbXJDghhvke/yFDp2vGdIO2hYs4MQ= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= +github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/goccy/go-json v0.10.6 h1:p8HrPJzOakx/mn/bQtjgNjdTcN+/S6FcG2CTtQOrHVU= +github.com/goccy/go-json v0.10.6/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= +github.com/google/flatbuffers v25.12.19+incompatible h1:haMV2JRRJCe1998HeW/p0X9UaMTK6SDo0ffLn2+DbLs= +github.com/google/flatbuffers v25.12.19+incompatible/go.mod h1:1AeVuKshWv4vARoZatz6mlQ0JxURH0Kv5+zNeJKJCa8= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/klauspost/compress v1.18.5 h1:/h1gH5Ce+VWNLSWqPzOVn6XBO+vJbCNGvjoaGBFW2IE= +github.com/klauspost/compress v1.18.5/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= +github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y= +github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= +github.com/pierrec/lz4/v4 v4.1.26 h1:GrpZw1gZttORinvzBdXPUXATeqlJjqUG/D87TKMnhjY= +github.com/pierrec/lz4/v4 v4.1.26/go.mod h1:EoQMVJgeeEOMsCqCzqFm2O0cJvljX2nGZjcRIPL34O4= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U= +github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/shopspring/decimal v1.4.0 h1:bxl37RwXBklmTi0C79JfXCEBD1cqqHt0bbgBAGFp81k= +github.com/shopspring/decimal v1.4.0/go.mod h1:gawqmDU56v4yIKSwfBSFip1HdCCXN8/+DMd9qYNcwME= +github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= +github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +github.com/zeebo/assert v1.3.0 h1:g7C04CbJuIDKNPFHmsk4hwZDO5O+kntRxzaUoNXj+IQ= +github.com/zeebo/assert v1.3.0/go.mod h1:Pq9JiuJQpG8JLJdtkwrJESF0Foym2/D9XMU5ciN/wJ0= +github.com/zeebo/xxh3 v1.1.0 h1:s7DLGDK45Dyfg7++yxI0khrfwq9661w9EN78eP/UZVs= +github.com/zeebo/xxh3 v1.1.0/go.mod h1:IisAie1LELR4xhVinxWS5+zf1lA4p0MW4T+w+W07F5s= +golang.org/x/exp v0.0.0-20260112195511-716be5621a96 h1:Z/6YuSHTLOHfNFdb8zVZomZr7cqNgTJvA8+Qz75D8gU= +golang.org/x/exp v0.0.0-20260112195511-716be5621a96/go.mod h1:nzimsREAkjBCIEFtHiYkrJyT+2uy9YZJB7H1k68CXZU= +golang.org/x/sys v0.43.0 h1:Rlag2XtaFTxp19wS8MXlJwTvoh8ArU6ezoyFsMyCTNI= +golang.org/x/sys v0.43.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= +gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/bin/kbsearch/main.go b/bin/kbsearch/main.go new file mode 100644 index 0000000..c896925 --- /dev/null +++ b/bin/kbsearch/main.go @@ -0,0 +1,44 @@ +// bin/kbsearch - the Go implementation of bin/kb/search (nested module so the +// root `go test ./...` and CI never compile it against native ladyships). +// +// Usage (built/run by ./bin/kb/search): +// +// kbsearch "query" [--root facts|info] [--repo P] [-n N] [--json] +// kbsearch serve [port] start the embedding daemon +// 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). +package main + +import ( + "fmt" + "log" + "os" + "strconv" +) + +func main() { + if len(os.Args) > 1 && os.Args[1] == "serve" { + port := 17830 + if len(os.Args) > 2 { + if p, err := strconv.Atoi(os.Args[2]); err == nil { + port = p + } + } + if err := serve(port); err != nil { + log.Fatalf("kbsearch serve: %v", err) + } + return + } + if len(os.Args) > 1 && os.Args[1] == "--list-model" { + dir, err := modelDir() + if err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + fmt.Println(dir) + return + } + os.Exit(runSearch(os.Args[1:])) +} \ No newline at end of file diff --git a/bin/kbsearch/model.go b/bin/kbsearch/model.go new file mode 100644 index 0000000..c8735fa --- /dev/null +++ b/bin/kbsearch/model.go @@ -0,0 +1,236 @@ +// StaticModel wraps the potion-multilingual-128m embedding model. +// +// Mirrors model2vec.StaticModel: tokenizer (daulet Unigram) + safetensors matrix. +// Embed(text) applies the same preprocessing: median_token_length pre-truncation, +// add_special_tokens=false, drop unk (id=1), truncate to 512, mean pool, L2 normalize +1e-32. +package main + +import ( + "encoding/json" + "fmt" + "io" + "math" + "os" + "path/filepath" + "sort" + + "github.com/chewxy/math32" + "github.com/daulet/tokenizers" +) + +type StaticModel struct { + tok *tokenizers.Tokenizer + mat []float32 // row-major: vocab_size x 128 + medianLen int + vocabSize int + dim int +} + +func loadModel() (*StaticModel, error) { + dir, err := modelDir() + if err != nil { + return nil, err + } + + tok, err := tokenizers.FromFile(filepath.Join(dir, "tokenizer.json")) + if err != nil { + return nil, fmt.Errorf("tokenizer: %w", err) + } + + mat, vocabSize, dim, err := loadMatrix(filepath.Join(dir, "model.safetensors")) + if err != nil { + return nil, fmt.Errorf("safetensors: %w", err) + } + + median := medianTokenLength(filepath.Join(dir, "tokenizer.json")) + + return &StaticModel{ + tok: tok, + mat: mat, + vocabSize: vocabSize, + dim: dim, + medianLen: median, + }, nil +} + +func (m *StaticModel) Close() error { + if m.tok != nil { + m.tok.Close() + m.tok = nil + } + return nil +} + +func (m *StaticModel) Embed(text string) ([]float64, error) { + const maxLen = 512 + + if m.medianLen > 0 { + maxChars := maxLen * m.medianLen + runes := []rune(text) + if len(runes) > maxChars { + text = string(runes[:maxChars]) + } + } + + ids, _, err := m.tok.EncodeErr(text, false) + if err != nil { + return nil, fmt.Errorf("encode: %w", err) + } + + filtered := make([]uint32, 0, len(ids)) + for _, id := range ids { + if id != 1 { + filtered = append(filtered, id) + } + if len(filtered) >= maxLen { + break + } + } + if len(filtered) == 0 { + return make([]float64, m.dim), nil + } + + acc := make([]float32, m.dim) + for _, id := range filtered { + if int(id) >= m.vocabSize { + continue + } + off := int(id) * m.dim + for d := 0; d < m.dim; d++ { + acc[d] += m.mat[off+d] + } + } + inv := 1.0 / float32(len(filtered)) + for d := 0; d < m.dim; d++ { + acc[d] *= inv + } + + var norm float32 + for d := 0; d < m.dim; d++ { + norm += acc[d] * acc[d] + } + norm = math32.Sqrt(norm) + 1e-32 + for d := 0; d < m.dim; d++ { + acc[d] /= norm + } + + out := make([]float64, m.dim) + for d := 0; d < m.dim; d++ { + out[d] = float64(acc[d]) + } + return out, nil +} + +func medianTokenLength(tokenizerPath string) int { + data, err := os.ReadFile(tokenizerPath) + if err != nil { + return 0 + } + var parsed struct { + Model struct { + Vocab [][]json.RawMessage `json:"vocab"` + } `json:"model"` + } + if err := json.Unmarshal(data, &parsed); err != nil { + return 0 + } + vocab := parsed.Model.Vocab + if len(vocab) == 0 { + return 0 + } + lengths := make([]int, 0, len(vocab)) + for _, pair := range vocab { + if len(pair) < 1 { + continue + } + var tok string + if err := json.Unmarshal(pair[0], &tok); err != nil { + continue + } + lengths = append(lengths, len([]rune(tok))) + } + if len(lengths) == 0 { + return 0 + } + sort.Ints(lengths) + return lengths[len(lengths)/2] +} + +func loadMatrix(path string) ([]float32, int, int, error) { + f, err := os.Open(path) + if err != nil { + return nil, 0, 0, err + } + defer f.Close() + + var hdrLen uint64 + if err := binaryRead(f, &hdrLen); err != nil { + return nil, 0, 0, err + } + hdrBytes := make([]byte, hdrLen) + if _, err := io.ReadFull(f, hdrBytes); err != nil { + return nil, 0, 0, err + } + + var hdr struct { + Embeddings struct { + Dtype string `json:"dtype"` + Shape []int `json:"shape"` + Offset []uint64 `json:"data_offsets"` + } `json:"embeddings"` + } + if err := json.Unmarshal(hdrBytes, &hdr); err != nil { + return nil, 0, 0, err + } + if hdr.Embeddings.Dtype != "F32" { + return nil, 0, 0, fmt.Errorf("unsupported dtype %s", hdr.Embeddings.Dtype) + } + if len(hdr.Embeddings.Shape) != 2 { + return nil, 0, 0, fmt.Errorf("expected 2D shape, got %v", hdr.Embeddings.Shape) + } + vocabSize := hdr.Embeddings.Shape[0] + dim := hdr.Embeddings.Shape[1] + if len(hdr.Embeddings.Offset) != 2 { + return nil, 0, 0, fmt.Errorf("bad offsets") + } + start := hdr.Embeddings.Offset[0] + end := hdr.Embeddings.Offset[1] + size := end - start + if size != uint64(vocabSize*dim*4) { + return nil, 0, 0, fmt.Errorf("size mismatch") + } + + if _, err := f.Seek(int64(8+hdrLen+start), io.SeekStart); err != nil { + return nil, 0, 0, err + } + + buf := make([]byte, size) + if _, err := io.ReadFull(f, buf); err != nil { + return nil, 0, 0, err + } + + mat := make([]float32, vocabSize*dim) + for i := 0; i < len(mat); i++ { + off := i * 4 + mat[i] = math.Float32frombits( + uint32(buf[off]) | + uint32(buf[off+1])<<8 | + uint32(buf[off+2])<<16 | + uint32(buf[off+3])<<24, + ) + } + return mat, vocabSize, dim, nil +} + +func binaryRead(r io.Reader, v any) error { + switch p := v.(type) { + case *uint64: + var b [8]byte + if _, err := io.ReadFull(r, b[:]); err != nil { + return err + } + *p = uint64(b[0]) | uint64(b[1])<<8 | uint64(b[2])<<16 | uint64(b[3])<<24 | + uint64(b[4])<<32 | uint64(b[5])<<40 | uint64(b[6])<<48 | uint64(b[7])<<56 + } + return nil +} \ No newline at end of file diff --git a/bin/kbsearch/modeldir.go b/bin/kbsearch/modeldir.go new file mode 100644 index 0000000..7913745 --- /dev/null +++ b/bin/kbsearch/modeldir.go @@ -0,0 +1,60 @@ +// modelDir returns the resolved potion-multilingual-128m model directory. +package main + +import ( + "fmt" + "os" + "path/filepath" + "strings" +) + +func modelDir() (string, error) { + // 1. Explicit env + if v := os.Getenv("KBSEARCH_MODEL"); v != "" { + return v, nil + } + // 2. Next to the binary (dev or installed) + self, err := os.Executable() + if err == nil { + if dir, err := filepath.EvalSymlinks(filepath.Dir(self)); err == nil { + cand := filepath.Join(dir, "potion-multilingual-128m") + if st, err := os.Stat(cand); err == nil && st.IsDir() { + return cand, nil + } + } + } + // 3. Repo root lib/ (where other scripts expect it) + if v := os.Getenv("KB_ROOT"); v != "" { + cand := filepath.Join(v, "lib", "potion-multilingual-128m") + if st, err := os.Stat(cand); err == nil && st.IsDir() { + return cand, nil + } + } + // 4. HF cache (new layout: models--*/snapshots/*) + if v := os.Getenv("HF_HOME"); v != "" { + base := filepath.Join(v, "hub") + if entries, err := os.ReadDir(base); err == nil { + for _, e := range entries { + if strings.HasPrefix(e.Name(), "models--") { + snapDir := filepath.Join(base, e.Name(), "snapshots") + if snaps, err := os.ReadDir(snapDir); err == nil { + for _, s := range snaps { + cand := filepath.Join(snapDir, s.Name()) + if st, _ := os.Stat(cand); st != nil && st.IsDir() { + return cand, nil + } + } + } + } + } + } + } + // 5. Legacy HF cache (symlinked model dir) + if v := os.Getenv("HF_HOME"); v != "" { + cand := filepath.Join(v, "potion-multilingual-128m") + if st, err := os.Stat(cand); err == nil && st.IsDir() { + return cand, nil + } + } + return "", fmt.Errorf("model not found (set KBSEARCH_MODEL or KB_ROOT, or download to HF cache)") +} \ No newline at end of file diff --git a/bin/kbsearch/search.go b/bin/kbsearch/search.go new file mode 100644 index 0000000..e8ff2f2 --- /dev/null +++ b/bin/kbsearch/search.go @@ -0,0 +1,465 @@ +// Hybrid FTS + vector search implementation, plus daemon client/server. +package main + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "log" + "net" + "net/http" + "os" + "os/exec" + "path/filepath" + "sort" + "strconv" + "strings" + "time" + + lbug "github.com/LadybugDB/go-ladybug" +) + +const defaultPort = 17830 +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]) + } + } + } + + if listModel { + dir, err := modelDir() + if err != nil { + fmt.Fprintln(os.Stderr, err) + return 1 + } + fmt.Println(dir) + 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 + } + defer closeBrain() + + emb, err := embedQuery(query) + if err != nil { + fmt.Fprintf(os.Stderr, "embed: %v\n", err) + return 1 + } + + fts, err := queryFTS(query, limit*3) + if err != nil { + fmt.Fprintf(os.Stderr, "fts: %v\n", err) + return 1 + } + + var vec []Hit + if vec, err = queryVector(emb, limit*3); err != nil { + 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] + } + + for i := range results { + if results[i].Text != "" { + runes := []rune(results[i].Text) + if len(runes) > 280 { + runes = runes[:280] + } + results[i].Snippet = string(runes) + } + } + + out := Dict{ + {"query", query}, + {"root_filter", root}, + {"count", len(results)}, + {"results", resultsToDicts(results)}, + } + + if jsonOut { + enc := json.NewEncoder(os.Stdout) + enc.SetIndent("", " ") + enc.SetEscapeHTML(false) + return b2i(enc.Encode(toJSONOut(results, query, root))) + } + fmt.Print(toYAML(out, 0)) + return 0 +} + +func b2i(err error) int { + if err != nil { + return 1 + } + return 0 +} + +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", + ) + if err != nil { + return nil, err + } + defer stmt.Close() + res, err := conn.Execute(stmt, map[string]any{"q": text, "n": limit}) + if err != nil { + return nil, err + } + return rowsToHits(res) +} + +func queryVector(emb []float64, limit int) ([]Hit, error) { + embList := make([]any, len(emb)) + 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", + ) + if err != nil { + return nil, err + } + defer stmt.Close() + res, err := conn.Execute(stmt, map[string]any{"q": embList, "n": limit}) + if err != nil { + return nil, err + } + hits, err := rowsToHits(res) + if err != nil { + return nil, err + } + for i := range hits { + hits[i].Score = 1.0 - hits[i].Score + } + return hits, nil +} + +func rowsToHits(res *lbug.QueryResult) ([]Hit, error) { + var hits []Hit + for res.HasNext() { + row, err := res.Next() + if err != nil { + return nil, err + } + vals, err := row.GetAsSlice() + if err != nil || len(vals) < 5 { + continue + } + id := fmt.Sprint(vals[0]) + text := fmt.Sprint(vals[1]) + root := fmt.Sprint(vals[2]) + source := fmt.Sprint(vals[3]) + score := float64(vals[4].(float64)) + hits = append(hits, Hit{ID: id, Text: text, Root: root, Source: source, Score: score}) + } + return hits, nil +} + +// JSON output types +type jsonOut struct { + Query string `json:"query"` + RootFilter string `json:"root_filter"` + Count int `json:"count"` + Results []jsonHit `json:"results"` +} + +type jsonHit struct { + ID string `json:"id"` + Text string `json:"text"` + Root string `json:"root"` + Score float64 `json:"score"` + Snippet string `json:"snippet,omitempty"` +} + +func toJSONOut(hits []Hit, query, rootFilter string) *jsonOut { + out := make([]jsonHit, len(hits)) + for i, h := range hits { + out[i] = jsonHit{ + ID: h.ID, + Text: h.Text, + Root: h.Root, + Score: h.Score, + Snippet: h.Snippet, + } + } + return &jsonOut{ + Query: query, + RootFilter: rootFilter, + Count: len(hits), + Results: out, + } +} + +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 { + d := Dict{ + {"id", h.ID}, + {"text", h.Text}, + {"root", h.Root}, + {"score", h.Score}, + } + if h.Snippet != "" { + d = append(d, KV{"snippet", h.Snippet}) + } + out[i] = d + } + return out +} + +// --- Daemon server --- +func serve(port int) error { + model, err := loadModel() + if err != nil { + return fmt.Errorf("load model: %w", err) + } + defer model.Close() + + mux := http.NewServeMux() + mux.HandleFunc(healthPath, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + mux.HandleFunc(daemonPath, func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + w.WriteHeader(http.StatusMethodNotAllowed) + return + } + var req struct { + Text string `json:"text"` + } + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + w.WriteHeader(http.StatusBadRequest) + return + } + vec, err := model.Embed(req.Text) + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + json.NewEncoder(w).Encode(map[string]string{"error": err.Error()}) + return + } + json.NewEncoder(w).Encode(map[string]any{"vector": vec}) + }) + + addr := fmt.Sprintf("127.0.0.1:%d", port) + log.Printf("kbsearch daemon listening on %s", addr) + return http.ListenAndServe(addr, mux) +} + +// --- Daemon client --- +var daemonClient = &http.Client{ + Timeout: 5 * time.Second, + Transport: &http.Transport{ + DialContext: (&net.Dialer{Timeout: 2 * time.Second}).DialContext, + }, +} + +func embedQuery(text string) ([]float64, error) { + port := defaultPort + if envPort := os.Getenv("KBSEARCH_PORT"); envPort != "" { + if p, err := strconv.Atoi(envPort); err == nil { + port = p + } + } + emb, err := tryDaemon(text, port) + if err == nil { + return emb, nil + } + + model, err := loadModel() + if err != nil { + return nil, fmt.Errorf("fallback load model: %w", err) + } + defer model.Close() + return model.Embed(text) +} + +func tryDaemon(text string, port int) ([]float64, error) { + url := fmt.Sprintf("http://127.0.0.1:%d%s", port, daemonPath) + payload := map[string]string{"text": text} + body, _ := json.Marshal(payload) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + req, _ := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + resp, err := daemonClient.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("daemon HTTP %d", resp.StatusCode) + } + var r struct { + Vector []float64 `json:"vector"` + Error string `json:"error"` + } + if err := json.NewDecoder(resp.Body).Decode(&r); err != nil { + return nil, err + } + if r.Error != "" { + return nil, errors.New(r.Error) + } + return r.Vector, nil +} + +func ensureDaemon(port int) error { + url := fmt.Sprintf("http://127.0.0.1:%d%s", port, healthPath) + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancel() + req, _ := http.NewRequestWithContext(ctx, "GET", url, nil) + if resp, err := daemonClient.Do(req); err == nil { + resp.Body.Close() + if resp.StatusCode == http.StatusOK { + return nil + } + } + + self, err := os.Executable() + if err != nil { + return err + } + cmd := exec.Command(self, "serve", strconv.Itoa(port)) + cmd.Dir, _ = filepath.Split(self) + cmd.Stdout = nil + cmd.Stderr = nil + if err := cmd.Start(); err != nil { + return err + } + + for i := 0; i < 40; i++ { + time.Sleep(250 * time.Millisecond) + req, _ := http.NewRequestWithContext(context.Background(), "GET", url, nil) + if resp, err := daemonClient.Do(req); err == nil { + resp.Body.Close() + if resp.StatusCode == http.StatusOK { + return nil + } + } + } + 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 new file mode 100644 index 0000000..2a403bd --- /dev/null +++ b/bin/kbsearch/types.go @@ -0,0 +1,16 @@ +// Common types and helpers for kbsearch. +package main + +import "os" + +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 diff --git a/bin/kbsearch/yaml.go b/bin/kbsearch/yaml.go new file mode 100644 index 0000000..74c4f18 --- /dev/null +++ b/bin/kbsearch/yaml.go @@ -0,0 +1,107 @@ +// YAML emitter ported from bin/kb/yamlout.py — preserves insertion order. +package main + +import ( + "fmt" + "strconv" + "strings" +) + +// KV is an ordered key-value pair for maps. +type KV struct { + K string + V any +} + +// Dict is an ordered map (slice of KV). +type Dict []KV + +func toYAML(node any, indent int) string { + pad := strings.Repeat(" ", indent) + switch n := node.(type) { + case Dict: + if len(n) == 0 { + return pad + "{}\n" + } + var b strings.Builder + for _, kv := range n { + switch nv := kv.V.(type) { + case Dict: + if len(nv) == 0 { + b.WriteString(pad + kv.K + ": {}\n") + } else { + b.WriteString(pad + kv.K + ":\n" + toYAML(nv, indent+1)) + } + case []any: + if len(nv) == 0 { + b.WriteString(pad + kv.K + ": []\n") + } else { + b.WriteString(pad + kv.K + ":\n" + toYAML(nv, indent+1)) + } + default: + b.WriteString(pad + kv.K + ": " + scalar(nv) + "\n") + } + } + return b.String() + + case []any: + if len(n) == 0 { + return pad + "[]\n" + } + var b strings.Builder + for _, item := range n { + if d, ok := item.(Dict); ok { + b.WriteString(pad + "-\n" + toYAML(d, indent+1)) + } else { + b.WriteString(pad + "- " + scalar(item) + "\n") + } + } + return b.String() + + default: + return pad + scalar(node) + "\n" + } +} + +func scalar(v any) string { + switch t := v.(type) { + case nil: + return "null" + case bool: + if t { + return "true" + } + return "false" + case int: + return strconv.Itoa(t) + case int64: + return strconv.FormatInt(t, 10) + case float64: + return fmtFloat(t) + case float32: + return fmtFloat(float64(t)) + case string: + return quoteIfNeeded(t) + default: + // fallback + return fmt.Sprintf("%v", v) + } +} + +func fmtFloat(f float64) string { + s := strconv.FormatFloat(f, 'g', -1, 64) + if !strings.ContainsAny(s, ".eE") { + s += ".0" + } + return s +} + +func quoteIfNeeded(s string) string { + if strings.Contains(s, "\n") { + return strconv.Quote(s) + } + if s == "" || strings.ContainsAny(s, ":#'\"[]{}&*!|>%@`") || s != strings.TrimSpace(s) { + return strconv.Quote(s) + } + return s +} \ No newline at end of file