Knowledge search ranks by how near a chunk's meaning is, not by its place in two lists
Fused by rank, a full-text match on a question's everyday words outvoted a far nearer meaning: on the PoC, "how much leave do fathers get" found the annual-leave page before the paternity one (cosine distance 0.44 against 0.25). Now a chunk scores its similarity to the question plus up to 0.3 for a full-text match, each search brings four times k candidates, the text search drops Turkish and English stop words, and an operator's search says when the embedding model did not answer rather than quietly searching text. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
1ec9b83076
commit
bc5ee9f3aa
3 files changed
+153
-69
No files matched your search
@@ -450,11 +450,15 @@ func (a *API) kbSearch(w http.ResponseWriter, r *http.Request) {
|
||||
k = 5
|
||||
}
|
||||
q := store.KBSearch{OrgID: c.OrgID, Collections: []store.KBCollection{c}, Text: in.Query, K: k}
|
||||
// Without the model, the full-text half still answers.
|
||||
// Said, not swallowed: a search that quietly lost its meaning half
|
||||
// looks like a collection that does not know the answer.
|
||||
if a.kb.Embed != nil {
|
||||
if v, err := a.kb.Embed(r.Context(), c.EmbedModel, in.Query); err == nil {
|
||||
q.Embeddings = map[string][]float32{c.EmbedModel: v}
|
||||
v, err := a.kb.Embed(r.Context(), c.EmbedModel, in.Query)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, c.EmbedModel+" did not embed the question: "+err.Error())
|
||||
return
|
||||
}
|
||||
q.Embeddings = map[string][]float32{c.EmbedModel: v}
|
||||
}
|
||||
hits, err := a.kb.Store.SearchChunks(r.Context(), q)
|
||||
if err != nil {
|
||||
|
||||
+113
-65
@@ -89,8 +89,9 @@ type KBHit struct {
|
||||
Document string
|
||||
Ord int
|
||||
Body string
|
||||
// Score is the fused rank (reciprocal rank fusion over the vector and
|
||||
// text searches): only the order means anything.
|
||||
// Score is the cosine similarity to the question, plus up to
|
||||
// textWeight for a full-text match: higher is better, and only the
|
||||
// order is meant to be compared.
|
||||
Score float64
|
||||
}
|
||||
|
||||
@@ -626,13 +627,22 @@ type KBSearch struct {
|
||||
K int
|
||||
}
|
||||
|
||||
// rrfK is reciprocal rank fusion's constant: 60, as in the paper and nearly
|
||||
// everywhere it is used.
|
||||
const rrfK = 60
|
||||
// textWeight is how much a full-text match adds to a chunk's similarity to
|
||||
// the question, at most. The embedding model reads what a question means;
|
||||
// full text is there for what it cannot — a code, a name, a number. Fused by
|
||||
// rank, as first built, the two were equals and the gap between a close
|
||||
// meaning and a distant one was thrown away: "how much leave do fathers
|
||||
// get" found the annual-leave page over the paternity one on the PoC,
|
||||
// though the paternity page was far nearer (0.25 against 0.44).
|
||||
const textWeight = 0.3
|
||||
|
||||
// candidates is how many each search brings to the fusion for k results:
|
||||
// with only k from each, the meaning search's third-best never met the
|
||||
// text search's first, ties were broken by id, and the order was luck.
|
||||
func candidates(k int) int { return max(4*k, 20) }
|
||||
|
||||
// SearchChunks finds the best chunks for a query across collections: a
|
||||
// nearest-neighbour search in each collection and a full-text search across
|
||||
// them, fused by rank. Every collection must be in OrgID; the caller decides
|
||||
// nearest-neighbour search and a full-text search in each, put on one scale. Every collection must be in OrgID; the caller decides
|
||||
// which ones the requester may read, and the table refuses rows outside them
|
||||
// regardless.
|
||||
func (s *DB) SearchChunks(ctx context.Context, q KBSearch) ([]KBHit, error) {
|
||||
@@ -643,86 +653,93 @@ func (s *DB) SearchChunks(ctx context.Context, q KBSearch) ([]KBHit, error) {
|
||||
return nil, nil
|
||||
}
|
||||
ids := make([]string, 0, len(q.Collections))
|
||||
quoted := make([]string, 0, len(q.Collections))
|
||||
for _, c := range q.Collections {
|
||||
if c.OrgID != q.OrgID || !kbIDPattern.MatchString(c.ID) {
|
||||
return nil, fmt.Errorf("collection %s is not in organization %s", c.ID, q.OrgID)
|
||||
}
|
||||
ids = append(ids, c.ID)
|
||||
quoted = append(quoted, sqlLiteral(c.ID))
|
||||
}
|
||||
type key struct {
|
||||
doc string
|
||||
ord int
|
||||
}
|
||||
hits := map[key]*KBHit{}
|
||||
add := func(h KBHit, rank int) {
|
||||
k := key{h.DocID, h.Ord}
|
||||
if have, ok := hits[k]; ok {
|
||||
have.Score += 1.0 / float64(rrfK+rank)
|
||||
return
|
||||
}
|
||||
h.Score = 1.0 / float64(rrfK+rank)
|
||||
hits[k] = &h
|
||||
type found struct {
|
||||
hit KBHit
|
||||
// sim is the cosine similarity to the question, when the question
|
||||
// was embedded for this chunk's collection; text is its full-text
|
||||
// rank, 0 when the text search did not find it.
|
||||
sim, text float64
|
||||
hasSim bool
|
||||
}
|
||||
scan := func(rows *sql.Rows) ([]KBHit, error) {
|
||||
hits := map[key]*found{}
|
||||
add := func(h KBHit, sim sql.NullFloat64, text float64) {
|
||||
k := key{h.DocID, h.Ord}
|
||||
f, ok := hits[k]
|
||||
if !ok {
|
||||
f = &found{hit: h}
|
||||
hits[k] = f
|
||||
}
|
||||
if sim.Valid {
|
||||
f.sim, f.hasSim = sim.Float64, true
|
||||
}
|
||||
f.text = max(f.text, text)
|
||||
}
|
||||
query := func(tx conn, text bool, q string, args ...any) error {
|
||||
rows, err := tx.QueryContext(ctx, q, args...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []KBHit
|
||||
for rows.Next() {
|
||||
var h KBHit
|
||||
if err := rows.Scan(&h.Collection, &h.DocID, &h.Document, &h.Ord, &h.Body); err != nil {
|
||||
return nil, err
|
||||
var sim sql.NullFloat64
|
||||
var rank float64
|
||||
if err := rows.Scan(&h.Collection, &h.DocID, &h.Document, &h.Ord, &h.Body, &sim, &rank); err != nil {
|
||||
return err
|
||||
}
|
||||
out = append(out, h)
|
||||
if !text {
|
||||
rank = 0
|
||||
}
|
||||
add(h, sim, rank)
|
||||
}
|
||||
return out, rows.Err()
|
||||
return rows.Err()
|
||||
}
|
||||
tsq := orQuery(q.Text)
|
||||
err := s.inTx(ctx, func(tx conn) error {
|
||||
if err := kbScope(ctx, tx, q.OrgID, ids); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, c := range q.Collections {
|
||||
v, ok := q.Embeddings[c.EmbedModel]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if len(v) != c.Dims {
|
||||
return fmt.Errorf("the query has %d dimensions; collection %s takes %d", len(v), c.Name, c.Dims)
|
||||
}
|
||||
// The literal collection and cast are what let the planner use
|
||||
// that collection's partial index.
|
||||
cast := `::vector(` + strconv.Itoa(c.Dims) + `)`
|
||||
rows, err := tx.QueryContext(ctx, `SELECT k.collection, k.doc_id, d.name, k.ord, k.body
|
||||
FROM kb_chunks k JOIN kb_documents d ON d.id = k.doc_id
|
||||
WHERE k.org_id = ? AND k.collection = `+sqlLiteral(c.ID)+`
|
||||
ORDER BY k.embedding`+cast+` <=> ?`+cast+` LIMIT ?`,
|
||||
q.OrgID, vectorText(v), q.K)
|
||||
if err != nil {
|
||||
return err
|
||||
lit := sqlLiteral(c.ID)
|
||||
sim, args := "NULL::float8", []any{}
|
||||
v, embedded := q.Embeddings[c.EmbedModel]
|
||||
if embedded {
|
||||
if len(v) != c.Dims {
|
||||
return fmt.Errorf("the query has %d dimensions; collection %s takes %d", len(v), c.Name, c.Dims)
|
||||
}
|
||||
cast := `::vector(` + strconv.Itoa(c.Dims) + `)`
|
||||
sim = `1 - (k.embedding` + cast + ` <=> ?` + cast + `)`
|
||||
args = append(args, vectorText(v))
|
||||
if err := query(tx, false, `SELECT k.collection, k.doc_id, d.name, k.ord, k.body, `+sim+`, 0::float8
|
||||
FROM kb_chunks k JOIN kb_documents d ON d.id = k.doc_id
|
||||
WHERE k.org_id = ? AND k.collection = `+lit+`
|
||||
ORDER BY k.embedding`+cast+` <=> ?`+cast+` LIMIT ?`,
|
||||
vectorText(v), q.OrgID, vectorText(v), candidates(q.K)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
got, err := scan(rows)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for i, h := range got {
|
||||
add(h, i+1)
|
||||
}
|
||||
}
|
||||
if tsq := orQuery(q.Text); tsq != "" {
|
||||
rows, err := tx.QueryContext(ctx, `SELECT k.collection, k.doc_id, d.name, k.ord, k.body
|
||||
FROM kb_chunks k JOIN kb_documents d ON d.id = k.doc_id, to_tsquery('simple', ?) q
|
||||
WHERE k.org_id = ? AND k.collection IN (`+strings.Join(quoted, ", ")+`) AND k.tsv @@ q
|
||||
ORDER BY ts_rank(k.tsv, q) DESC LIMIT ?`,
|
||||
tsq, q.OrgID, q.K)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
got, err := scan(rows)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for i, h := range got {
|
||||
add(h, i+1)
|
||||
if tsq != "" {
|
||||
// The text search's finds are measured against the question's
|
||||
// meaning too, so the two are compared on one scale.
|
||||
if err := query(tx, true, `SELECT k.collection, k.doc_id, d.name, k.ord, k.body, `+sim+`, ts_rank(k.tsv, tq)::float8
|
||||
FROM kb_chunks k JOIN kb_documents d ON d.id = k.doc_id, to_tsquery('simple', ?) tq
|
||||
WHERE k.org_id = ? AND k.collection = `+lit+` AND k.tsv @@ tq
|
||||
ORDER BY ts_rank(k.tsv, tq) DESC LIMIT ?`,
|
||||
append(args, tsq, q.OrgID, candidates(q.K))...); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
@@ -730,9 +747,21 @@ func (s *DB) SearchChunks(ctx context.Context, q KBSearch) ([]KBHit, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// One scale: similarity, plus the text rank as a share of the best one,
|
||||
// weighted down.
|
||||
best := 0.0
|
||||
for _, f := range hits {
|
||||
best = max(best, f.text)
|
||||
}
|
||||
for _, f := range hits {
|
||||
f.hit.Score = f.sim
|
||||
if best > 0 {
|
||||
f.hit.Score += textWeight * f.text / best
|
||||
}
|
||||
}
|
||||
out := make([]KBHit, 0, len(hits))
|
||||
for _, h := range hits {
|
||||
out = append(out, *h)
|
||||
for _, f := range hits {
|
||||
out = append(out, f.hit)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
if out[i].Score != out[j].Score {
|
||||
@@ -749,6 +778,25 @@ func (s *DB) SearchChunks(ctx context.Context, q KBSearch) ([]KBHit, error) {
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// stopWords are words nearly every passage has, in Turkish and English:
|
||||
// matching on them only ranks whatever is longest first.
|
||||
var stopWords = func() map[string]bool {
|
||||
m := map[string]bool{}
|
||||
for _, w := range strings.Fields(`
|
||||
acaba ama ancak artık aslında az bana bazı belki ben beni benim beri bile bir birçok biri birkaç
|
||||
biz bize bizi bizim bu buna bunda bundan bunlar bunları bunların bunu bunun burada çok çünkü da
|
||||
daha dahi de defa diye en gibi göre hangi hangisi hem hep hepsi her hiç için ile ise kaç kadar
|
||||
ki kim kime kimi mı mi mu mü nasıl ne neden nerde nerede nereye niçin niye o olan olarak oldu
|
||||
olduğu olsa olur ona ondan onlar onları onların onu onun orada öyle şey şu şuna şunda şundan şunu
|
||||
tüm ve veya ya yani var yok mısın misin musun müsün midir mıdır nedir hangileri lütfen
|
||||
a an and are as at be by can do does for from has have how i in is it its me my of on or our
|
||||
should that the their there this to was what when where which who why will with you your
|
||||
`) {
|
||||
m[w] = true
|
||||
}
|
||||
return m
|
||||
}()
|
||||
|
||||
// orQuery turns a question into a tsquery matching any of its words. All of
|
||||
// them (plainto_tsquery) would match almost nothing a person types as a
|
||||
// question; the rank then prefers chunks with more of them.
|
||||
@@ -761,7 +809,7 @@ func orQuery(text string) string {
|
||||
for _, w := range strings.FieldsFunc(text, func(r rune) bool {
|
||||
return !unicode.IsLetter(r) && !unicode.IsDigit(r) && !unicode.IsMark(r)
|
||||
}) {
|
||||
if len([]rune(w)) < 2 || seen[w] {
|
||||
if len([]rune(w)) < 2 || seen[w] || stopWords[strings.ToLowerSpecial(unicode.TurkishCase, w)] {
|
||||
continue
|
||||
}
|
||||
seen[w] = true
|
||||
|
||||
@@ -352,10 +352,42 @@ func TestPostgresAPersonsCollectionsGoWithTheAccount(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestPostgresMeaningOutweighsSharedWords(t *testing.T) {
|
||||
// The PoC's case: the question shares its everyday words with one
|
||||
// passage and its meaning with another.
|
||||
ctx := context.Background()
|
||||
s := openKB(t)
|
||||
kbTwoOrgs(t, s)
|
||||
c := mustCollection(t, s, ScopeProject, "proj1", "", "docs")
|
||||
put := func(name, body string, v []float32) {
|
||||
t.Helper()
|
||||
if _, _, err := s.PutDocument(ctx, c, KBDocument{Name: name, Digest: name}, []KBChunk{{Body: body, Embedding: v}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
put("yillik.md", "yıllık izin kadar izin alır izin", unit(0))
|
||||
put("babalik.md", "babalık hakkı on gündür", unit(1))
|
||||
for _, f := range []string{"a", "b", "c"} {
|
||||
put(f+".md", "başka bir konu", unit(2))
|
||||
}
|
||||
hits, err := s.SearchChunks(ctx, KBSearch{OrgID: "org1", Collections: []KBCollection{c},
|
||||
Embeddings: map[string][]float32{"emb": unit(1)}, Text: "babalar ne kadar izin alır", K: 1})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hits) != 1 || hits[0].Document != "babalik.md" {
|
||||
t.Errorf("best hit %+v", hits)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOrQueryMatchesAnyWord(t *testing.T) {
|
||||
if got := orQuery("Yıllık izin kaç gün? İzin & 'x' -- a"); got != "Yıllık | izin | kaç | gün | İzin" {
|
||||
if got := orQuery("Yıllık izin kaç gün? İzin & 'x' -- a"); got != "Yıllık | izin | gün | İzin" {
|
||||
t.Errorf("orQuery gave %q", got)
|
||||
}
|
||||
// Words every passage has are no evidence of anything.
|
||||
if got := orQuery("Babalar NE kadar izin alır, ve İÇİN"); got != "Babalar | izin | alır" {
|
||||
t.Errorf("stop words kept: %q", got)
|
||||
}
|
||||
if got := orQuery("?!"); got != "" {
|
||||
t.Errorf("no words gave %q", got)
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user