casbin
/
casibase
1
0
Fork 1
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Improve filterTextFiles()

This commit is contained in:
Yang Luo 2023-08-14 02:47:44 +08:00
parent 081b88a2be
commit 66a5021eb7
1 changed files with 5 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
object/vector_embedding.go
View File

@ -28,14 +28,10 @@ import (
"golang.org/x/time/rate" "golang.org/x/time/rate"
) )
func isTxt(filename string) bool { func filterTextFiles(files []*storage.Object) []*storage.Object {
return strings.HasSuffix(filename, ".txt")
}
func filterTxtFiles(files []*storage.Object) []*storage.Object {
var res []*storage.Object var res []*storage.Object
for _, file := range files { for _, file := range files {
if isTxt(file.Key) { if strings.HasSuffix(file.Key, ".txt") || strings.HasSuffix(file.Key, ".md") {
res = append(res, file) res = append(res, file)
} }
} }
@ -43,13 +39,13 @@ func filterTxtFiles(files []*storage.Object) []*storage.Object {
return res return res
} }
func getTxtFiles(provider string, prefix string) ([]*storage.Object, error) { func getTextFiles(provider string, prefix string) ([]*storage.Object, error) {
files, err := storage.ListObjects(provider, prefix) files, err := storage.ListObjects(provider, prefix)
if err != nil { if err != nil {
return nil, err return nil, err
} }
return filterTxtFiles(files), nil return filterTextFiles(files), nil
} }
func getObjectReadCloser(object *storage.Object) (io.ReadCloser, error) { func getObjectReadCloser(object *storage.Object) (io.ReadCloser, error) {
@ -92,7 +88,7 @@ func addEmbeddedVector(authToken string, text string, storeName string, fileName
func setTxtObjectVector(authToken string, provider string, key string, storeName string) (bool, error) { func setTxtObjectVector(authToken string, provider string, key string, storeName string) (bool, error) {
lb := rate.NewLimiter(rate.Every(time.Minute), 3) lb := rate.NewLimiter(rate.Every(time.Minute), 3)
txtObjects, err := getTxtFiles(provider, key) txtObjects, err := getTextFiles(provider, key)
if err != nil { if err != nil {
return false, err return false, err
} }