Skip to content

Commit

Permalink
add search feature
Browse files Browse the repository at this point in the history
  • Loading branch information
quantonganh committed Apr 22, 2023
1 parent 7f63432 commit a1bb111
Show file tree
Hide file tree
Showing 4 changed files with 184 additions and 11 deletions.
1 change: 1 addition & 0 deletions go.mod
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Expand Up @@ -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=
Expand Down
70 changes: 59 additions & 11 deletions main.go
Expand Up @@ -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() {
Expand Down Expand Up @@ -398,37 +450,33 @@ 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
}
return nil
})

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).
Expand Down
122 changes: 122 additions & 0 deletions 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
}

0 comments on commit a1bb111

Please sign in to comment.