282 lines
7.6 KiB
Go
282 lines
7.6 KiB
Go
|
package chatpipeline
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"regexp"
|
|||
|
|
"strings"
|
|||
|
|
"sync"
|
|||
|
|
"unicode"
|
|||
|
|
"unicode/utf8"
|
|||
|
|
|
|||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// runQueryExpansion performs query expansion when initial recall is low.
|
|||
|
|
// It generates query variants and runs concurrent retrieval across search targets.
|
|||
|
|
func (p *PluginSearch) runQueryExpansion(ctx context.Context, chatManage *types.ChatManage) []*types.SearchResult {
|
|||
|
|
pipelineInfo(ctx, "Search", "recall_low", map[string]interface{}{
|
|||
|
|
"current": len(chatManage.SearchResult),
|
|||
|
|
"threshold": chatManage.EmbeddingTopK,
|
|||
|
|
})
|
|||
|
|
expansions := p.expandQueries(ctx, chatManage)
|
|||
|
|
if len(expansions) == 0 {
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
pipelineInfo(ctx, "Search", "expansion_start", map[string]interface{}{
|
|||
|
|
"variants": len(expansions),
|
|||
|
|
})
|
|||
|
|
expTopK := max(chatManage.EmbeddingTopK*2, chatManage.RerankTopK*2)
|
|||
|
|
expKwTh := chatManage.KeywordThreshold * 0.8
|
|||
|
|
|
|||
|
|
// Concurrent expansion retrieval across queries and search targets
|
|||
|
|
expResults := make([]*types.SearchResult, 0, expTopK*len(expansions))
|
|||
|
|
var muExp sync.Mutex
|
|||
|
|
var wgExp sync.WaitGroup
|
|||
|
|
jobs := len(expansions) * len(chatManage.SearchTargets)
|
|||
|
|
capSem := 16
|
|||
|
|
if jobs < capSem {
|
|||
|
|
capSem = jobs
|
|||
|
|
}
|
|||
|
|
if capSem <= 0 {
|
|||
|
|
capSem = 1
|
|||
|
|
}
|
|||
|
|
sem := make(chan struct{}, capSem)
|
|||
|
|
pipelineInfo(ctx, "Search", "expansion_concurrency", map[string]interface{}{
|
|||
|
|
"jobs": jobs,
|
|||
|
|
"cap": capSem,
|
|||
|
|
})
|
|||
|
|
for _, q := range expansions {
|
|||
|
|
for _, target := range chatManage.SearchTargets {
|
|||
|
|
if target == nil || target.KnowledgeBaseID == "" {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
wgExp.Add(1)
|
|||
|
|
go func(q string, t *types.SearchTarget) {
|
|||
|
|
defer wgExp.Done()
|
|||
|
|
sem <- struct{}{}
|
|||
|
|
defer func() { <-sem }()
|
|||
|
|
vectorThreshold, keywordThreshold := t.RecallThresholds(
|
|||
|
|
chatManage.VectorThreshold,
|
|||
|
|
expKwTh,
|
|||
|
|
)
|
|||
|
|
paramsExp := types.SearchParams{
|
|||
|
|
QueryText: q,
|
|||
|
|
VectorThreshold: vectorThreshold,
|
|||
|
|
KeywordThreshold: keywordThreshold,
|
|||
|
|
MatchCount: expTopK,
|
|||
|
|
TagIDs: t.TagIDs,
|
|||
|
|
ScopeTagIDs: t.ScopeTagIDs,
|
|||
|
|
DisableVectorMatch: false,
|
|||
|
|
DisableKeywordsMatch: false,
|
|||
|
|
SkipContextEnrichment: true, // Pipeline handles context assembly in merge stage
|
|||
|
|
}
|
|||
|
|
// Apply knowledge ID filter if this is a partial KB search
|
|||
|
|
if t.Type == types.SearchTargetTypeKnowledge {
|
|||
|
|
paramsExp.KnowledgeIDs = t.KnowledgeIDs
|
|||
|
|
}
|
|||
|
|
res, err := p.knowledgeBaseService.HybridSearch(ctx, t.KnowledgeBaseID, paramsExp)
|
|||
|
|
if err != nil {
|
|||
|
|
pipelineWarn(ctx, "Search", "expansion_error", map[string]interface{}{
|
|||
|
|
"kb_id": t.KnowledgeBaseID,
|
|||
|
|
"error": err.Error(),
|
|||
|
|
})
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
if len(res) > 0 {
|
|||
|
|
for _, r := range res {
|
|||
|
|
r.KnowledgeBaseID = t.KnowledgeBaseID
|
|||
|
|
}
|
|||
|
|
pipelineInfo(ctx, "Search", "expansion_hits", map[string]interface{}{
|
|||
|
|
"kb_id": t.KnowledgeBaseID,
|
|||
|
|
"query": q,
|
|||
|
|
"hits": len(res),
|
|||
|
|
})
|
|||
|
|
muExp.Lock()
|
|||
|
|
expResults = append(expResults, res...)
|
|||
|
|
muExp.Unlock()
|
|||
|
|
}
|
|||
|
|
}(q, target)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
wgExp.Wait()
|
|||
|
|
|
|||
|
|
if len(expResults) > 0 {
|
|||
|
|
pipelineInfo(ctx, "Search", "expansion_done", map[string]interface{}{
|
|||
|
|
"added": len(expResults),
|
|||
|
|
})
|
|||
|
|
}
|
|||
|
|
return expResults
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// expandQueries generates query variants locally without LLM to improve keyword recall.
|
|||
|
|
// Uses simple techniques: word reordering, stopword removal, key phrase extraction.
|
|||
|
|
func (p *PluginSearch) expandQueries(ctx context.Context, chatManage *types.ChatManage) []string {
|
|||
|
|
query := strings.TrimSpace(chatManage.RewriteQuery)
|
|||
|
|
if query == "" {
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
expansions := make([]string, 0, 5)
|
|||
|
|
seen := make(map[string]struct{})
|
|||
|
|
seen[strings.ToLower(query)] = struct{}{}
|
|||
|
|
if q := strings.ToLower(chatManage.Query); q != "" {
|
|||
|
|
seen[q] = struct{}{}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
addIfNew := func(s string) {
|
|||
|
|
s = strings.TrimSpace(s)
|
|||
|
|
if s == "" || len(s) < 3 {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
key := strings.ToLower(s)
|
|||
|
|
if _, ok := seen[key]; ok {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
seen[key] = struct{}{}
|
|||
|
|
expansions = append(expansions, s)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 1. Remove common stopwords and create keyword-only variant
|
|||
|
|
keywords := extractKeywords(query)
|
|||
|
|
if len(keywords) >= 2 {
|
|||
|
|
addIfNew(strings.Join(keywords, " "))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 2. Extract quoted phrases or key segments
|
|||
|
|
phrases := extractPhrases(query)
|
|||
|
|
for _, phrase := range phrases {
|
|||
|
|
addIfNew(phrase)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 3. Split by common delimiters and use longest segment
|
|||
|
|
segments := splitByDelimiters(query)
|
|||
|
|
for _, seg := range segments {
|
|||
|
|
if len(seg) > 5 {
|
|||
|
|
addIfNew(seg)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 4. Remove question words (什么/如何/怎么/为什么/哪个 etc.)
|
|||
|
|
cleaned := removeQuestionWords(query)
|
|||
|
|
if cleaned != query {
|
|||
|
|
addIfNew(cleaned)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Limit to 5 expansions
|
|||
|
|
if len(expansions) > 5 {
|
|||
|
|
expansions = expansions[:5]
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
pipelineInfo(ctx, "Search", "local_expansion_result", map[string]interface{}{
|
|||
|
|
"variants": len(expansions),
|
|||
|
|
})
|
|||
|
|
return expansions
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Common Chinese and English stopwords
|
|||
|
|
var stopwords = map[string]struct{}{
|
|||
|
|
"的": {}, "是": {}, "在": {}, "了": {}, "和": {}, "与": {}, "或": {},
|
|||
|
|
"a": {}, "an": {}, "the": {}, "is": {}, "are": {}, "was": {}, "were": {},
|
|||
|
|
"be": {}, "been": {}, "being": {}, "have": {}, "has": {}, "had": {},
|
|||
|
|
"do": {}, "does": {}, "did": {}, "will": {}, "would": {}, "could": {},
|
|||
|
|
"should": {}, "may": {}, "might": {}, "must": {}, "can": {},
|
|||
|
|
"to": {}, "of": {}, "in": {}, "for": {}, "on": {}, "with": {}, "at": {},
|
|||
|
|
"by": {}, "from": {}, "as": {}, "into": {}, "through": {}, "about": {},
|
|||
|
|
"what": {}, "how": {}, "why": {}, "when": {}, "where": {}, "which": {},
|
|||
|
|
"who": {}, "whom": {}, "whose": {},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Question words in Chinese
|
|||
|
|
var questionWords = regexp.MustCompile(`^(什么是|什么|如何|怎么|怎样|为什么|为何|哪个|哪些|谁|何时|何地|请问|请告诉我|帮我|我想知道|我想了解)`)
|
|||
|
|
|
|||
|
|
func extractKeywords(text string) []string {
|
|||
|
|
words := tokenize(text)
|
|||
|
|
keywords := make([]string, 0, len(words))
|
|||
|
|
for _, w := range words {
|
|||
|
|
lower := strings.ToLower(w)
|
|||
|
|
if _, isStop := stopwords[lower]; !isStop && utf8.RuneCountInString(w) > 1 {
|
|||
|
|
keywords = append(keywords, w)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return keywords
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func extractPhrases(text string) []string {
|
|||
|
|
// Extract quoted content
|
|||
|
|
var phrases []string
|
|||
|
|
re := regexp.MustCompile(`["'"'「」『』]([^"'"'「」『』]+)["'"'「」『』]`)
|
|||
|
|
matches := re.FindAllStringSubmatch(text, -1)
|
|||
|
|
for _, m := range matches {
|
|||
|
|
if len(m) > 1 && len(m[1]) > 2 {
|
|||
|
|
phrases = append(phrases, m[1])
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return phrases
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func splitByDelimiters(text string) []string {
|
|||
|
|
// Split by common delimiters
|
|||
|
|
re := regexp.MustCompile(`[,，;；、。！？!?\s]+`)
|
|||
|
|
parts := re.Split(text, -1)
|
|||
|
|
var result []string
|
|||
|
|
for _, p := range parts {
|
|||
|
|
p = strings.TrimSpace(p)
|
|||
|
|
if p != "" {
|
|||
|
|
result = append(result, p)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return result
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func removeQuestionWords(text string) string {
|
|||
|
|
return strings.TrimSpace(questionWords.ReplaceAllString(text, ""))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func tokenize(text string) []string {
|
|||
|
|
var tokens []string
|
|||
|
|
var current strings.Builder
|
|||
|
|
currentIsHan := false
|
|||
|
|
|
|||
|
|
flush := func() {
|
|||
|
|
if current.Len() == 0 {
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if currentIsHan {
|
|||
|
|
// Use the existing search-mode dictionary for continuous Chinese text.
|
|||
|
|
for _, word := range types.Jieba.CutForSearch(current.String(), true) {
|
|||
|
|
word = strings.TrimSpace(word)
|
|||
|
|
if word != "" {
|
|||
|
|
tokens = append(tokens, word)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
} else {
|
|||
|
|
tokens = append(tokens, current.String())
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
current.Reset()
|
|||
|
|
currentIsHan = false
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
for _, r := range text {
|
|||
|
|
if unicode.Is(unicode.Han, r) {
|
|||
|
|
if current.Len() > 0 && !currentIsHan {
|
|||
|
|
flush()
|
|||
|
|
}
|
|||
|
|
currentIsHan = true
|
|||
|
|
current.WriteRune(r)
|
|||
|
|
} else if unicode.IsLetter(r) || unicode.IsDigit(r) {
|
|||
|
|
if current.Len() > 0 && currentIsHan {
|
|||
|
|
flush()
|
|||
|
|
}
|
|||
|
|
currentIsHan = false
|
|||
|
|
current.WriteRune(r)
|
|||
|
|
} else {
|
|||
|
|
flush()
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
flush()
|
|||
|
|
return tokens
|
|||
|
|
}
|