okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

dawntui: listView, confirm and one fuzzy matcher for list screens to share8101869

main
5.0 KiB161 linesraw
1package tui
2
3import (
4	"albedo/cli/internal/daemon"
5	"slices"
6	"strings"
7	"unicode"
8
9	"github.com/sahilm/fuzzy"
10)
11
12// mergeModels lists seeds first, then the rest of listed, each id once. A
13// listed entry lends its catalog facts to the seed of the same id.
14func mergeModels(seeds []string, listed []daemon.Model) []daemon.Model {
15	byID := make(map[string]daemon.Model, len(listed))
16	for _, m := range listed {
17		byID[m.ID] = m
18	}
19	seen := map[string]bool{"": true}
20	var out []daemon.Model
21	add := func(m daemon.Model) {
22		if !seen[m.ID] {
23			seen[m.ID] = true
24			out = append(out, m)
25		}
26	}
27	for _, id := range seeds {
28		m, ok := byID[id]
29		if !ok {
30			m = daemon.Model{ID: id}
31		}
32		add(m)
33	}
34	for _, m := range listed {
35		add(m)
36	}
37	return out
38}
39
40var searchNoise = strings.NewReplacer(" ", "", "-", "", "_", "", ".", "", ":", "", "/", "")
41
42func searchKey(s string) string { return searchNoise.Replace(strings.ToLower(s)) }
43
44// matchRank orders search matches, compared element by element: an exact id
45// first, then words that start the id or one of its parts, then the fuzzy
46// score. A fuzzy bonus never outweighs a word typed from the start.
47type matchRank [3]int
48
49func (a matchRank) compare(b matchRank) int { return slices.Compare(a[:], b[:]) }
50
51// matchModel ranks a model against the search words. Each word must fuzzy
52// match the model id or the profile name; typing profile and id together
53// is an exact match too.
54func matchModel(words []string, profile, id string) (rank matchRank, hits []int, ok bool) {
55	rank, hits, ok = matchFields(words, id, profile)
56	if whole := strings.Join(words, ""); ok && whole != "" && whole == searchKey(profile+id) {
57		rank[0] = 1
58	}
59	return rank, hits, ok
60}
61
62// searchWords splits a query into the keys its words match by.
63func searchWords(query string) []string {
64	var words []string
65	for field := range strings.FieldsSeq(query) {
66		if word := searchKey(field); word != "" {
67			words = append(words, word)
68		}
69	}
70	return words
71}
72
73// matchFields ranks a row's fields against the search words, which every
74// list screen shares: each word must fuzzy match some field. The hits are
75// the characters of the first field to underline. An empty search matches
76// everything equally.
77func matchFields(words []string, fields ...string) (rank matchRank, hits []int, ok bool) {
78	if whole := strings.Join(words, ""); whole != "" && slices.ContainsFunc(fields, func(f string) bool { return whole == searchKey(f) }) {
79		rank[0] = 1
80	}
81	for _, word := range words {
82		var best matchRank
83		var bestHits []int
84		found := fuzzy.Find(word, fields)
85		for i, f := range found {
86			if r := (matchRank{0, startTier(fields[f.Index], word), f.Score}); i == 0 || r.compare(best) > 0 {
87				best, bestHits = r, nil
88				if f.Index == 0 {
89					bestHits = f.MatchedIndexes
90				}
91			}
92		}
93		if len(found) == 0 {
94			return matchRank{}, nil, false
95		}
96		rank[1] += best[1]
97		rank[2] += best[2]
98		hits = append(hits, bestHits...)
99	}
100	return rank, hits, true
101}
102
103// startTier is 2 when word starts s, 1 when it starts a part of s after a
104// separator, else 0. Separators fold as in searchKey, so "gpt5" starts "gpt-5".
105func startTier(s, word string) int {
106	if strings.HasPrefix(searchKey(s), word) {
107		return 2
108	}
109	for i := 1; i < len(s); i++ {
110		if strings.ContainsRune(searchSeparators, rune(s[i-1])) && strings.HasPrefix(searchKey(s[i:]), word) {
111			return 1
112		}
113	}
114	return 0
115}
116
117// refilter rebuilds the rows from the catalogs and the search text. The
118// cursor stays on its row unless reset.
119func (m *ModelPickerModel) refilter(reset bool) {
120	keep := ""
121	if r, ok := m.highlighted(); ok && !reset {
122		keep = r.key()
123	}
124	query := strings.TrimSpace(m.search.Value())
125	words := searchWords(query)
126
127	// Rows rank within their profile, and profiles by their best
128	// row. Ties, and an empty search, keep the catalog order.
129	type scored struct {
130		row  modelRow
131		rank matchRank
132	}
133	var groups [][]scored
134	var typed []modelRow
135	canType := query != "" && !strings.ContainsFunc(query, unicode.IsSpace) && validateModelID(query) == nil
136	for _, c := range m.catalogs {
137		var group []scored
138		for _, model := range c.models {
139			if rank, hits, ok := matchModel(words, c.name, model.ID); ok {
140				group = append(group, scored{modelRow{profile: c.name, model: model, hits: hits}, rank})
141			}
142		}
143		if len(group) > 0 {
144			slices.SortStableFunc(group, func(a, b scored) int { return b.rank.compare(a.rank) })
145			groups = append(groups, group)
146		}
147		// An id has no spaces, so a search of several words offers none.
148		if canType && !slices.ContainsFunc(c.models, func(model daemon.Model) bool { return model.ID == query }) {
149			typed = append(typed, modelRow{profile: c.name, model: daemon.Model{ID: query}, typed: true})
150		}
151	}
152	slices.SortStableFunc(groups, func(a, b []scored) int { return b[0].rank.compare(a[0].rank) })
153	m.rows = nil
154	for _, group := range groups {
155		for _, s := range group {
156			m.rows = append(m.rows, s.row)
157		}
158	}
159	m.rows = append(m.rows, typed...)
160	m.cursor = max(0, slices.IndexFunc(m.rows, func(r modelRow) bool { return r.key() == keep }))
161}