diff --git a/go.mod b/go.mod index b4f30c8..ac232ae 100644 --- a/go.mod +++ b/go.mod @@ -4,6 +4,7 @@ go 1.20 require ( github.com/gdamore/tcell/v2 v2.6.0 + github.com/kljensen/snowball v0.8.0 github.com/mitchellh/go-homedir v1.1.0 github.com/rivo/tview v0.0.0-20230320095235-84f9c0ff9de8 github.com/tidwall/buntdb v1.2.10 diff --git a/go.sum b/go.sum index 3fa5461..ebb0153 100644 --- a/go.sum +++ b/go.sum @@ -2,6 +2,8 @@ github.com/gdamore/encoding v1.0.0 h1:+7OoQ1Bc6eTm5niUzBa0Ctsh6JbMW6Ra+YNuAtDBdk github.com/gdamore/encoding v1.0.0/go.mod h1:alR0ol34c49FCSBLjhosxzcPHQbf2trDkoo5dl+VrEg= github.com/gdamore/tcell/v2 v2.6.0 h1:OKbluoP9VYmJwZwq/iLb4BxwKcwGthaa1YNBJIyCySg= github.com/gdamore/tcell/v2 v2.6.0/go.mod h1:be9omFATkdr0D9qewWW3d+MEvl5dha+Etb5y65J2H8Y= +github.com/kljensen/snowball v0.8.0 h1:WU4cExxK6sNW33AiGdbn4e8RvloHrhkAssu2mVJ11kg= +github.com/kljensen/snowball v0.8.0/go.mod h1:OGo5gFWjaeXqCu4iIrMl5OYip9XUJHGOU5eSkPjVg2A= github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY= github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mattn/go-runewidth v0.0.14 h1:+xnbZSEeDbOIg5/mE6JF0w6n9duR1l3/WmbinWVwUuU= diff --git a/main.go b/main.go index 8d61893..631ed70 100644 --- a/main.go +++ b/main.go @@ -149,8 +149,60 @@ func main() { AddItem(nil, 1, 1, false) } + searchInputField := tview.NewInputField() + searchInputField.SetTitle("Search") + searchInputField. + SetFieldWidth(50). + SetAcceptanceFunc(tview.InputFieldMaxLength(50)) + searchInputField.SetBorder(true) + searchInputField.SetDoneFunc(func(key tcell.Key) { + switch key { + case tcell.KeyEnter: + titles := make([]string, 0, len(m)) + db.View(func(tx *buntdb.Tx) error { + err := tx.Descend("time", func(key, value string) bool { + titles = append(titles, key) + return true + }) + return err + }) + + text := searchInputField.GetText() + if text != "" { + idx := make(index) + idx.add(titles) + r := idx.search(text) + list.Clear() + for _, i := range r { + list.AddItem(titles[i], "", rune(0), func() { + if c, ok := m[titles[i]]; ok { + textView.SetText(toConversation(c.Messages)) + } + }) + } + } else { + list.Clear() + for i := range titles { + list.AddItem(titles[i], "", rune(0), func() { + if c, ok := m[titles[i]]; ok { + textView.SetText(toConversation(c.Messages)) + } + }) + } + } + if list.GetItemCount() > 0 { + app.SetFocus(list) + } + } + }) + var hiddenItemCount int list.SetInputCapture(func(event *tcell.EventKey) *tcell.EventKey { + switch event.Key() { + case tcell.KeyESC: + app.SetFocus(searchInputField) + } + _, _, _, height := list.GetInnerRect() switch event.Rune() { @@ -398,24 +450,20 @@ func main() { return event }) - button := tview.NewButton("+ New chat") - button.SetFocusFunc(func() { - isNewChat = true - textView.Clear() - app.SetFocus(textArea) - }) - button.SetBorder(true) - app.SetInputCapture(func(event *tcell.EventKey) *tcell.EventKey { switch event.Key() { case tcell.KeyF1: - app.SetFocus(button) + isNewChat = true + textView.Clear() + app.SetFocus(textArea) case tcell.KeyF2: app.SetFocus(list) case tcell.KeyF3: app.SetFocus(textView) case tcell.KeyF4: app.SetFocus(textArea) + case tcell.KeyCtrlS: + app.SetFocus(searchInputField) default: return event } @@ -423,12 +471,12 @@ func main() { }) help := tview.NewTextView().SetRegions(true).SetDynamicColors(true) - help.SetText("F1: new chat, F2: history, F3: conversation, F4: question, enter: submit, j/k: down/up, ctrl-f/b: page down/up, e: edit, d: delete, ctrl-c: quit").SetTextAlign(tview.AlignCenter) + help.SetText("F1: new chat, F2: history, F3: conversation, F4: question, enter: submit, ctrl-s: search, j/k: down/up, e: edit, d: delete, ctrl-f/b: page down/up, ctrl-c: quit").SetTextAlign(tview.AlignCenter) mainFlex := tview.NewFlex().SetDirection(tview.FlexRow). AddItem(tview.NewFlex().SetDirection(tview.FlexColumn). AddItem(tview.NewFlex().SetDirection(tview.FlexRow). - AddItem(button, 3, 1, false). + AddItem(searchInputField, 3, 1, false). AddItem(list, 0, 1, false), 0, 1, false). AddItem(tview.NewFlex().SetDirection(tview.FlexRow). AddItem(textView, 0, 1, false). diff --git a/search.go b/search.go new file mode 100644 index 0000000..6a1c92a --- /dev/null +++ b/search.go @@ -0,0 +1,122 @@ +// https://artem.krylysov.com/blog/2020/07/28/lets-build-a-full-text-search-engine/ +package main + +import ( + "strings" + "unicode" + + snowballeng "github.com/kljensen/snowball/english" +) + +type index map[string][]int + +func (idx index) add(titles []string) { + for id, title := range titles { + for _, token := range analyze(title) { + if contains(idx[token], id) { + continue + } + idx[token] = append(idx[token], id) + } + } +} + +func analyze(text string) []string { + tokens := tokenize(text) + tokens = toLower(tokens) + tokens = removeCommonWords(tokens) + tokens = stem(tokens) + return tokens +} + +func tokenize(text string) []string { + return strings.FieldsFunc(text, func(r rune) bool { + return !unicode.IsLetter(r) && !unicode.IsNumber(r) + }) +} + +func toLower(tokens []string) []string { + r := make([]string, len(tokens)) + for i, token := range tokens { + r[i] = strings.ToLower(token) + } + return r +} + +var stopWords = map[string]struct{}{ + "a": {}, + "and": {}, + "be": {}, + "have": {}, + "i": {}, + "in": {}, + "of": {}, + "that": {}, + "the": {}, + "to": {}, +} + +func removeCommonWords(tokens []string) []string { + r := make([]string, 0, len(tokens)) + for _, token := range tokens { + if _, ok := stopWords[token]; !ok { + r = append(r, token) + } + } + return r +} + +func stem(tokens []string) []string { + r := make([]string, len(tokens)) + for i, token := range tokens { + r[i] = snowballeng.Stem(token, false) + } + return r +} + +func contains(slice []int, val int) bool { + for _, s := range slice { + if s == val { + return true + } + } + return false +} + +func intersection(a, b []int) []int { + maxLen := len(a) + if len(b) > maxLen { + maxLen = len(b) + } + + r := make([]int, 0, maxLen) + var i, j int + for i < len(a) && j < len(b) { + if a[i] < b[j] { + i++ + } else if a[i] > b[j] { + j++ + } else { + r = append(r, a[i]) + i++ + j++ + } + } + return r +} + +func (idx index) search(text string) []int { + var r []int + for _, token := range analyze(text) { + if ids, ok := idx[token]; ok { + if r == nil { + r = ids + } else { + r = intersection(r, ids) + } + } else { + return nil + } + } + return r +}