refactor: one Go module; brain search in bin/brain + internal/brain.
Tests / Test (push) Skipped
Tests / Release (semver) (push) Skipped
Tests / Test (push) Skipped
Tests / Release (semver) (push) Skipped
Collapse nested kbsearch/chats go.mod into the root module. Ranking stays cgo-free under internal/brain/rank so CI does not need ladybug. bin/kb/search is a deprecation wrapper that still sets CGO and builds the binary.
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
// Commands in this directory are shebang mains (search.go).
|
||||
// search.go is behind the system_ladybug build tag (cgo).
|
||||
package main
|
||||
Executable
+23
@@ -0,0 +1,23 @@
|
||||
//usr/bin/env go run -tags=system_ladybug "$0" "$@"; exit
|
||||
//go:build cgo && system_ladybug
|
||||
//
|
||||
// bin/brain/search.go - deduction search over the 2dph brain.
|
||||
//
|
||||
// ./bin/brain/search.go "query" [--root facts|info] [--repo P] [-n N] [--json]
|
||||
// ./bin/brain/search.go serve [port]
|
||||
// ./bin/brain/search.go --list-model
|
||||
//
|
||||
// Needs CGO + libladybug (CGO_CFLAGS/CGO_LDFLAGS). Prefer the wrapper
|
||||
// bin/kb/search which sets those and builds a binary for the embed daemon.
|
||||
// NOTE: never run `gofmt -w` on this file — it breaks the shebang.
|
||||
package main
|
||||
|
||||
import (
|
||||
"os"
|
||||
|
||||
"github.com/eSlider/2dph/internal/brain"
|
||||
)
|
||||
|
||||
func main() {
|
||||
os.Exit(brain.Main(os.Args[1:]))
|
||||
}
|
||||
@@ -1,3 +0,0 @@
|
||||
module github.com/eSlider/2dph/bin/chats
|
||||
|
||||
go 1.25.0
|
||||
+24
-21
@@ -1,34 +1,37 @@
|
||||
#!/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.
|
||||
# bin/kb/search — deprecated wrapper. Use bin/brain/search.go.
|
||||
# Sets CGO for ladybug, builds a binary (embed daemon needs a real executable),
|
||||
# then execs it. Prints one deprecation line.
|
||||
set -euo pipefail
|
||||
|
||||
KB="$(cd "$(dirname "$0")/../.." && pwd)"
|
||||
BIN="$KB/var/bin/kbsearch"
|
||||
SRC="$KB/bin/kbsearch"
|
||||
ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
|
||||
BIN="$ROOT/var/bin/brain-search"
|
||||
SRC="$ROOT/internal/brain"
|
||||
CMD="$ROOT/bin/brain"
|
||||
|
||||
mkdir -p "$KB/var/bin"
|
||||
mkdir -p "$ROOT/var/bin"
|
||||
|
||||
# Rebuild if binary missing or any .go source newer
|
||||
need_build=0
|
||||
if [ ! -x "$BIN" ]; then
|
||||
need_build=1
|
||||
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)
|
||||
while IFS= read -r -d '' f; do
|
||||
if [ "$f" -nt "$BIN" ]; then
|
||||
need_build=1
|
||||
break
|
||||
fi
|
||||
done < <(find "$SRC" "$CMD" -name '*.go' -print0 2>/dev/null)
|
||||
fi
|
||||
|
||||
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
|
||||
echo "Building brain/search..." >&2
|
||||
(
|
||||
cd "$ROOT" &&
|
||||
CGO_CFLAGS="-I$ROOT/lib-ladybug" \
|
||||
CGO_LDFLAGS="-L$ROOT/lib-ladybug -Wl,-rpath,$ROOT/lib-ladybug" \
|
||||
go build -tags system_ladybug -o "$BIN" ./bin/brain
|
||||
) || exit 1
|
||||
fi
|
||||
|
||||
exec "$BIN" "$@"
|
||||
echo "bin/kb/search is deprecated; use bin/brain/search.go" >&2
|
||||
exec "$BIN" "$@"
|
||||
|
||||
@@ -1,99 +0,0 @@
|
||||
// Brain connection management using go-ladybug.
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
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 openWithSandbox(eps())
|
||||
}
|
||||
|
||||
func openWithSandbox(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)
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
func closeBrain() {
|
||||
if conn != nil {
|
||||
conn.Close()
|
||||
conn = nil
|
||||
}
|
||||
if db != nil {
|
||||
db.Close()
|
||||
db = nil
|
||||
}
|
||||
}
|
||||
@@ -1,23 +0,0 @@
|
||||
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
|
||||
)
|
||||
@@ -1,44 +0,0 @@
|
||||
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=
|
||||
@@ -1,46 +0,0 @@
|
||||
// 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 (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 (
|
||||
"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:]))
|
||||
}
|
||||
@@ -1,236 +0,0 @@
|
||||
// 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
|
||||
}
|
||||
@@ -1,60 +0,0 @@
|
||||
// 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)")
|
||||
}
|
||||
@@ -1,72 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -1,9 +0,0 @@
|
||||
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"
|
||||
@@ -1,100 +0,0 @@
|
||||
// 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
|
||||
}
|
||||
@@ -1,143 +0,0 @@
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
@@ -1,359 +0,0 @@
|
||||
// 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"
|
||||
"strconv"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
lbug "github.com/LadybugDB/go-ladybug"
|
||||
"github.com/eSlider/2dph/bin/kbsearch/rank"
|
||||
)
|
||||
|
||||
const defaultPort = 17830
|
||||
const daemonPath = "/embed"
|
||||
const healthPath = "/health"
|
||||
|
||||
func runSearch(args []string) int {
|
||||
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 opt.ListModel {
|
||||
dir, err := modelDir()
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
return 1
|
||||
}
|
||||
fmt.Println(dir)
|
||||
return 0
|
||||
}
|
||||
|
||||
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 := rank.RankAndFilter(fts, vec, root, repo, 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(rank.FTSStmt)
|
||||
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(rank.VecStmt)
|
||||
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 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
|
||||
}
|
||||
}
|
||||
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 {
|
||||
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
|
||||
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)
|
||||
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)
|
||||
}
|
||||
@@ -1,14 +0,0 @@
|
||||
// Common types and helpers for kbsearch.
|
||||
package main
|
||||
|
||||
import (
|
||||
"os"
|
||||
|
||||
"github.com/eSlider/2dph/bin/kbsearch/rank"
|
||||
)
|
||||
|
||||
func eps() string { return os.Getenv("KBTEST_EPS") }
|
||||
|
||||
// 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
|
||||
@@ -1,107 +0,0 @@
|
||||
// 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
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
"""D14 layout: bin/{subject}/{method}.go, libs in internal/, one go.mod."""
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
class BinLayoutTest(unittest.TestCase):
|
||||
def test_brain_search_shebang_exists(self) -> None:
|
||||
p = ROOT / "bin" / "brain" / "search.go"
|
||||
self.assertTrue(p.is_file(), "missing bin/brain/search.go")
|
||||
first = p.read_text().splitlines()[0]
|
||||
self.assertTrue(
|
||||
first.startswith("//usr/bin/env go run"),
|
||||
f"shebang first line, got {first!r}",
|
||||
)
|
||||
|
||||
def test_no_nested_go_mod_under_bin(self) -> None:
|
||||
nested = list((ROOT / "bin").rglob("go.mod"))
|
||||
self.assertEqual(nested, [], f"nested go.mod files: {nested}")
|
||||
|
||||
def test_rank_lives_in_internal_brain(self) -> None:
|
||||
self.assertTrue(
|
||||
(ROOT / "internal" / "brain" / "rank" / "rank.go").is_file(),
|
||||
"ranking must live in internal/brain/rank (cgo-free)",
|
||||
)
|
||||
self.assertFalse(
|
||||
(ROOT / "bin" / "kbsearch").exists(),
|
||||
"bin/kbsearch nested module must be gone",
|
||||
)
|
||||
|
||||
def test_no_main_go_under_bin_brain(self) -> None:
|
||||
main = ROOT / "bin" / "brain" / "main.go"
|
||||
self.assertFalse(main.exists(), "bin/brain/main.go is not a method")
|
||||
Reference in New Issue
Block a user