This repository has been archived by the owner on Jun 12, 2022. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 1
/
search.go
84 lines (70 loc) · 1.56 KB
/
search.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
package search
import (
"encoding/json"
"net/http"
"github.com/huichen/wukong/types"
"github.com/rs/zerolog/log"
"github.com/Xuanwo/xuanzang/contexts"
"github.com/Xuanwo/xuanzang/model"
)
// Response is the response for search.
type Response struct {
Tokens []string `json:"tokens"`
Docs []Document `json:"docs"`
Total int `json:"total"`
}
// Document is the document that scored.
type Document struct {
Title string `json:"title"`
URL string `json:"url"`
ContentText string `json:"content_text"`
}
// Search is the method to do search.
func Search(w http.ResponseWriter, r *http.Request) {
w.Header().Add("Content-Type", "application/json")
text := ""
if v := r.URL.Query().Get("text"); v != "" {
text = v
}
if text == "" {
log.Info().Msgf("Text is required.")
w.WriteHeader(400)
return
}
result := contexts.Searcher.Search(types.SearchRequest{
Text: text,
RankOptions: &types.RankOptions{
MaxOutputs: 20,
},
})
resp := &Response{
Tokens: result.Tokens,
Total: result.NumDocs,
}
resp.Docs = make([]Document, len(result.Docs))
for k, v := range result.Docs {
doc, err := model.GetDocByID(v.DocId)
if err != nil {
log.Error().Msgf("Get doc by id failed for %v.", err)
w.WriteHeader(500)
return
}
if doc == nil {
continue
}
d := Document{
Title: doc.Title,
URL: doc.URL,
}
resp.Docs[k] = d
}
content, err := json.Marshal(resp)
if err != nil {
log.Error().Msgf("JSON marshal failed for %v.", err)
w.WriteHeader(500)
return
}
w.WriteHeader(200)
w.Write(content)
return
}