mesh-catalog/modules/mesh-knowledge/cmd/mesh-knowledge/index.go

package main

// The search index (novox/hq to-be 55 §1, ADR 0597 §12): a full-text index in the module's own memory, written with the
// pure Go library bleve. At start the module reads every head from the ledger and the bodies of every current version into
// it, then keeps it current as it writes; `ask` and `holds` answer from it, and refuse with "index not built" while it is
// built. Nothing but the module writes it and nothing outside it reads it.

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/blevesearch/bleve/v2"
	"github.com/blevesearch/bleve/v2/analysis/lang/en"
	"github.com/blevesearch/bleve/v2/mapping"
	"github.com/blevesearch/bleve/v2/search/query"
)

// ErrIndexNotBuilt is the answer of a question asked while the index is built.
var ErrIndexNotBuilt = errors.New("index not built: the search index is being built from the ledger; ask again in a moment (get, list and propose answer at once)")

// searchIndex is the index and what it holds.
type searchIndex struct {
	mu      sync.Mutex
	idx     bleve.Index
	built   bool
	builtAt time.Time
	docs    int
	// docsOf are the documents each record has in the index, so a new version replaces them.
	docsOf map[Key][]string
}

func newSearchIndex() (*searchIndex, error) {
	m := bleve.NewIndexMapping()
	text := mapping.NewTextFieldMapping()
	text.Analyzer = en.AnalyzerName
	doc := bleve.NewDocumentMapping()
	for _, f := range []string{"title", "topic", "text", "clause"} {
		doc.AddFieldMappingsAt(f, text)
	}
	m.DefaultMapping = doc
	idx, err := bleve.NewMemOnly(m)
	if err != nil {
		return nil, err
	}
	return &searchIndex{idx: idx, docsOf: map[Key][]string{}}, nil
}

type indexDoc struct {
	Title  string `json:"title,omitempty"`
	Topic  string `json:"topic,omitempty"`
	Text   string `json:"text,omitempty"`
	Clause string `json:"clause,omitempty"`
}

// put replaces the documents of a record by those of its served version: its title, topic and sections as one document,
// and each clause in force or proposed as one of its own. A reserved, withdrawn or retracted record has none.
func (s *searchIndex) put(r *Record) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	key := r.Key()
	b := s.idx.NewBatch()
	for _, id := range s.docsOf[key] {
		b.Delete(id)
	}
	delete(s.docsOf, key)
	if r.Status != StatusReserved && r.Status != StatusWithdrawn && r.Status != StatusRetracted {
		var body strings.Builder
		for _, sec := range r.Sections {
			body.WriteString(sec.Text + "\n")
		}
		id := fmt.Sprintf("%s/%d", key.Kind, key.Number)
		if err := b.Index(id, indexDoc{Title: r.Title, Topic: r.Topic, Text: body.String()}); err != nil {
			return err
		}
		ids := []string{id}
		for _, c := range r.Clauses {
			if c.Status == ClauseRetracted {
				continue
			}
			cid := fmt.Sprintf("%s#%d", id, c.N)
			if err := b.Index(cid, indexDoc{Clause: c.Text}); err != nil {
				return err
			}
			ids = append(ids, cid)
		}
		s.docsOf[key] = ids
	}
	s.docs = 0
	for _, ids := range s.docsOf {
		s.docs += len(ids)
	}
	return s.idx.Batch(b)
}

// search ranks the documents against the question's words: a title weighs twice, as the in-memory ranking did.
func (s *searchIndex) search(question string, limit int) ([]Hit, error) {
	s.mu.Lock()
	built := s.built
	s.mu.Unlock()
	if !built {
		return nil, ErrIndexNotBuilt
	}
	var words []string
	for _, w := range wordsOf(question) {
		if !askingWords[w] && len(w) >= 3 {
			words = append(words, w)
		}
	}
	if len(words) == 0 {
		return nil, nil
	}
	var should []query.Query
	for _, w := range words {
		for field, boost := range map[string]float64{"title": 2, "topic": 1, "text": 0.05, "clause": 1} {
			q := bleve.NewMatchQuery(w)
			q.SetField(field)
			q.SetBoost(boost)
			should = append(should, q)
		}
	}
	req := bleve.NewSearchRequestOptions(bleve.NewDisjunctionQuery(should...), 5000, 0, false)
	res, err := s.idx.Search(req)
	if err != nil {
		return nil, err
	}
	var hits []Hit
	for _, h := range res.Hits {
		var kind string
		var number, clause int
		rest := h.ID
		if i := strings.IndexByte(rest, '#'); i >= 0 {
			if _, err := fmt.Sscanf(rest[i+1:], "%d", &clause); err != nil {
				continue
			}
			rest = rest[:i]
		}
		if i := strings.IndexByte(rest, '/'); i >= 0 {
			kind = rest[:i]
			if _, err := fmt.Sscanf(rest[i+1:], "%d", &number); err != nil {
				continue
			}
		}
		hits = append(hits, Hit{Key: Key{kind, number}, Clause: clause, Rank: h.Score})
	}
	sort.SliceStable(hits, func(i, j int) bool {
		if hits[i].Rank != hits[j].Rank {
			return hits[i].Rank > hits[j].Rank
		}
		if hits[i].Key != hits[j].Key {
			return hits[i].Key.Number > hits[j].Key.Number
		}
		return hits[i].Clause < hits[j].Clause
	})
	if limit > 0 && len(hits) > limit {
		hits = hits[:limit]
	}
	return hits, nil
}

// IndexState says how the index stands: built or not, when, and how many documents it holds.
type IndexState struct {
	Built   bool      `json:"built"`
	BuiltAt time.Time `json:"built-at,omitempty"`
	Docs    int       `json:"documents"`
}

func (s *searchIndex) state() IndexState {
	s.mu.Lock()
	defer s.mu.Unlock()
	return IndexState{Built: s.built, BuiltAt: s.builtAt, Docs: s.docs}
}

// BuildIndex reads every record from the ledger into the index and marks it built; asked again it builds it anew.
func (s *ledgerStore) BuildIndex(ctx context.Context) error {
	idx, err := newSearchIndex()
	if err != nil {
		return err
	}
	all, err := s.All(ctx)
	if err != nil {
		return err
	}
	for _, e := range all {
		if err := idx.put(e.Current); err != nil {
			return err
		}
	}
	idx.mu.Lock()
	idx.built, idx.builtAt = true, time.Now().UTC()
	idx.mu.Unlock()
	s.mu.Lock()
	s.index = idx
	s.mu.Unlock()
	return nil
}

// IndexState is how the module's index stands.
func (s *ledgerStore) IndexState() IndexState {
	s.mu.Lock()
	idx := s.index
	s.mu.Unlock()
	if idx == nil {
		return IndexState{}
	}
	return idx.state()
}

// indexed keeps the index current with a version the store wrote.
func (s *ledgerStore) indexed(r *Record, current bool) {
	if s.index == nil || !current {
		return
	}
	if err := s.index.put(r); err != nil && s.logf != nil {
		s.logf("[mesh-knowledge] the search index could not take %s: %v", r.Key(), err)
	}
}